module Freckle.App.OpenTelemetry.Context
  ( HasHeaders (..)
  , CustomTraceContext (..)
  , extractContext
  , injectContext
  , processWithContext
  ) where

import Freckle.App.Prelude

import Control.Error.Util (hush)
import Control.Lens (Lens', lens, (&), (.~), (^.))
import Control.Monad.Catch (MonadMask)
import Data.Aeson (FromJSON, ToJSON)
import Data.ByteString (ByteString)
import Data.CaseInsensitive (CI)
import Data.CaseInsensitive qualified as CI
import Faktory.Job (Job, custom, jobOptions)
import Faktory.Job.Custom (fromCustom, toCustom)
import Faktory.JobOptions (JobOptions (..))
import Freckle.App.OpenTelemetry
  ( MonadTracer (..)
  , SpanArguments
  , inSpan
  )
import Freckle.App.OpenTelemetry.ThreadContext (withTraceContext)
import Network.HTTP.Client (Request, requestHeaders)
import Network.HTTP.Simple (setRequestHeaders)
import Network.HTTP.Types.Header (Header)
import OpenTelemetry.Context (Context)
import OpenTelemetry.Context.ThreadLocal (attachContext, getContext)
import OpenTelemetry.Propagator (Propagator, extract, inject)
import OpenTelemetry.Trace.Core
  ( getTracerProviderPropagators
  , getTracerTracerProvider
  )

class HasHeaders a where
  headersL :: Lens' a [Header]

instance HasHeaders [Header] where
  headersL :: Lens' [Header] [Header]
headersL = ([Header] -> f [Header]) -> [Header] -> f [Header]
forall a. a -> a
id

instance HasHeaders Request where
  headersL :: Lens' Request [Header]
headersL = (Request -> [Header])
-> (Request -> [Header] -> Request) -> Lens' Request [Header]
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Request -> [Header]
requestHeaders ((Request -> [Header] -> Request) -> Lens' Request [Header])
-> (Request -> [Header] -> Request) -> Lens' Request [Header]
forall a b. (a -> b) -> a -> b
$ \Request
req [Header]
hs -> [Header] -> Request -> Request
setRequestHeaders [Header]
hs Request
req

instance HasHeaders (Job a) where
  headersL :: Lens' (Job a) [Header]
headersL = (JobOptions -> f JobOptions) -> Job a -> f (Job a)
forall a (f :: * -> *).
Functor f =>
(JobOptions -> f JobOptions) -> Job a -> f (Job a)
optionsL ((JobOptions -> f JobOptions) -> Job a -> f (Job a))
-> (([Header] -> f [Header]) -> JobOptions -> f JobOptions)
-> ([Header] -> f [Header])
-> Job a
-> f (Job a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CustomTraceContext -> f CustomTraceContext)
-> JobOptions -> f JobOptions
Lens' JobOptions CustomTraceContext
customTraceContextL ((CustomTraceContext -> f CustomTraceContext)
 -> JobOptions -> f JobOptions)
-> (([Header] -> f [Header])
    -> CustomTraceContext -> f CustomTraceContext)
-> ([Header] -> f [Header])
-> JobOptions
-> f JobOptions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Header] -> f [Header])
-> CustomTraceContext -> f CustomTraceContext
forall a. HasHeaders a => Lens' a [Header]
Lens' CustomTraceContext [Header]
headersL

optionsL :: Lens' (Job a) JobOptions
optionsL :: forall a (f :: * -> *).
Functor f =>
(JobOptions -> f JobOptions) -> Job a -> f (Job a)
optionsL = (Job a -> JobOptions)
-> (Job a -> JobOptions -> Job a)
-> forall {f :: * -> *}.
   Functor f =>
   (JobOptions -> f JobOptions) -> Job a -> f (Job a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Job a -> JobOptions
forall arg. Job arg -> JobOptions
jobOptions ((Job a -> JobOptions -> Job a)
 -> forall {f :: * -> *}.
    Functor f =>
    (JobOptions -> f JobOptions) -> Job a -> f (Job a))
-> (Job a -> JobOptions -> Job a)
-> forall {f :: * -> *}.
   Functor f =>
   (JobOptions -> f JobOptions) -> Job a -> f (Job a)
forall a b. (a -> b) -> a -> b
$ \Job a
x JobOptions
y -> Job a
x {jobOptions = y}

customTraceContextL :: Lens' JobOptions CustomTraceContext
customTraceContextL :: Lens' JobOptions CustomTraceContext
customTraceContextL = (JobOptions -> CustomTraceContext)
-> (JobOptions -> CustomTraceContext -> JobOptions)
-> Lens' JobOptions CustomTraceContext
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens JobOptions -> CustomTraceContext
forall {a}. (Monoid a, FromJSON a) => JobOptions -> a
get JobOptions -> CustomTraceContext -> JobOptions
forall {a}. ToJSON a => JobOptions -> a -> JobOptions
set
 where
  get :: JobOptions -> a
get JobOptions
jo = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
forall a. Monoid a => a
mempty (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$ Either String a -> Maybe a
forall a b. Either a b -> Maybe b
hush (Either String a -> Maybe a)
-> (Custom -> Either String a) -> Custom -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Custom -> Either String a
forall a. FromJSON a => Custom -> Either String a
fromCustom (Custom -> Maybe a) -> Maybe Custom -> Maybe a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< JobOptions -> Maybe Custom
joCustom JobOptions
jo
  set :: JobOptions -> a -> JobOptions
set JobOptions
jo = (JobOptions
jo <>) (JobOptions -> JobOptions) -> (a -> JobOptions) -> a -> JobOptions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Custom -> JobOptions
forall a. ToJSON a => a -> JobOptions
custom (Custom -> JobOptions) -> (a -> Custom) -> a -> JobOptions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Custom
forall a. ToJSON a => a -> Custom
toCustom

-- | A type that can be stored as the @custom@ field of a Faktory 'Job'
newtype CustomTraceContext = CustomTraceContext
  { CustomTraceContext -> [(Text, Text)]
traceHeaders :: [(Text, Text)]
  }
  deriving stock ((forall x. CustomTraceContext -> Rep CustomTraceContext x)
-> (forall x. Rep CustomTraceContext x -> CustomTraceContext)
-> Generic CustomTraceContext
forall x. Rep CustomTraceContext x -> CustomTraceContext
forall x. CustomTraceContext -> Rep CustomTraceContext x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CustomTraceContext -> Rep CustomTraceContext x
from :: forall x. CustomTraceContext -> Rep CustomTraceContext x
$cto :: forall x. Rep CustomTraceContext x -> CustomTraceContext
to :: forall x. Rep CustomTraceContext x -> CustomTraceContext
Generic)
  deriving newtype (NonEmpty CustomTraceContext -> CustomTraceContext
CustomTraceContext -> CustomTraceContext -> CustomTraceContext
(CustomTraceContext -> CustomTraceContext -> CustomTraceContext)
-> (NonEmpty CustomTraceContext -> CustomTraceContext)
-> (forall b.
    Integral b =>
    b -> CustomTraceContext -> CustomTraceContext)
-> Semigroup CustomTraceContext
forall b.
Integral b =>
b -> CustomTraceContext -> CustomTraceContext
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
$c<> :: CustomTraceContext -> CustomTraceContext -> CustomTraceContext
<> :: CustomTraceContext -> CustomTraceContext -> CustomTraceContext
$csconcat :: NonEmpty CustomTraceContext -> CustomTraceContext
sconcat :: NonEmpty CustomTraceContext -> CustomTraceContext
$cstimes :: forall b.
Integral b =>
b -> CustomTraceContext -> CustomTraceContext
stimes :: forall b.
Integral b =>
b -> CustomTraceContext -> CustomTraceContext
Semigroup, Semigroup CustomTraceContext
CustomTraceContext
Semigroup CustomTraceContext =>
CustomTraceContext
-> (CustomTraceContext -> CustomTraceContext -> CustomTraceContext)
-> ([CustomTraceContext] -> CustomTraceContext)
-> Monoid CustomTraceContext
[CustomTraceContext] -> CustomTraceContext
CustomTraceContext -> CustomTraceContext -> CustomTraceContext
forall a.
Semigroup a =>
a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
$cmempty :: CustomTraceContext
mempty :: CustomTraceContext
$cmappend :: CustomTraceContext -> CustomTraceContext -> CustomTraceContext
mappend :: CustomTraceContext -> CustomTraceContext -> CustomTraceContext
$cmconcat :: [CustomTraceContext] -> CustomTraceContext
mconcat :: [CustomTraceContext] -> CustomTraceContext
Monoid)
  deriving anyclass (Value -> Parser [CustomTraceContext]
Value -> Parser CustomTraceContext
(Value -> Parser CustomTraceContext)
-> (Value -> Parser [CustomTraceContext])
-> FromJSON CustomTraceContext
forall a.
(Value -> Parser a) -> (Value -> Parser [a]) -> FromJSON a
$cparseJSON :: Value -> Parser CustomTraceContext
parseJSON :: Value -> Parser CustomTraceContext
$cparseJSONList :: Value -> Parser [CustomTraceContext]
parseJSONList :: Value -> Parser [CustomTraceContext]
FromJSON, [CustomTraceContext] -> Value
[CustomTraceContext] -> Encoding
CustomTraceContext -> Value
CustomTraceContext -> Encoding
(CustomTraceContext -> Value)
-> (CustomTraceContext -> Encoding)
-> ([CustomTraceContext] -> Value)
-> ([CustomTraceContext] -> Encoding)
-> ToJSON CustomTraceContext
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> ToJSON a
$ctoJSON :: CustomTraceContext -> Value
toJSON :: CustomTraceContext -> Value
$ctoEncoding :: CustomTraceContext -> Encoding
toEncoding :: CustomTraceContext -> Encoding
$ctoJSONList :: [CustomTraceContext] -> Value
toJSONList :: [CustomTraceContext] -> Value
$ctoEncodingList :: [CustomTraceContext] -> Encoding
toEncodingList :: [CustomTraceContext] -> Encoding
ToJSON)

instance HasHeaders CustomTraceContext where
  headersL :: Lens' CustomTraceContext [Header]
headersL = (CustomTraceContext -> [Header])
-> (CustomTraceContext -> [Header] -> CustomTraceContext)
-> Lens' CustomTraceContext [Header]
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (((Text, Text) -> Header) -> [(Text, Text)] -> [Header]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Text) -> Header
encode ([(Text, Text)] -> [Header])
-> (CustomTraceContext -> [(Text, Text)])
-> CustomTraceContext
-> [Header]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CustomTraceContext -> [(Text, Text)]
traceHeaders) ((CustomTraceContext -> [Header] -> CustomTraceContext)
 -> Lens' CustomTraceContext [Header])
-> (CustomTraceContext -> [Header] -> CustomTraceContext)
-> Lens' CustomTraceContext [Header]
forall a b. (a -> b) -> a -> b
$ \CustomTraceContext
x [Header]
y -> CustomTraceContext
x {traceHeaders = map decode y}

encode :: (Text, Text) -> (CI ByteString, ByteString)
encode :: (Text, Text) -> Header
encode = (Text -> CI ByteString)
-> (Text -> ByteString) -> (Text, Text) -> Header
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (ByteString -> CI ByteString
forall s. FoldCase s => s -> CI s
CI.mk (ByteString -> CI ByteString)
-> (Text -> ByteString) -> Text -> CI ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
encodeUtf8) Text -> ByteString
encodeUtf8

decode :: (CI ByteString, ByteString) -> (Text, Text)
decode :: Header -> (Text, Text)
decode = (CI ByteString -> Text)
-> (ByteString -> Text) -> Header -> (Text, Text)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (ByteString -> Text
decodeUtf8 (ByteString -> Text)
-> (CI ByteString -> ByteString) -> CI ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> ByteString
forall s. CI s -> s
CI.original) ByteString -> Text
decodeUtf8

-- | Update our trace context from that extracted from the given item's headers
extractContext
  :: (MonadIO m, MonadTracer m, HasHeaders a) => a -> m ()
extractContext :: forall (m :: * -> *) a.
(MonadIO m, MonadTracer m, HasHeaders a) =>
a -> m ()
extractContext a
a = do
  Context
context <- m Context
forall (m :: * -> *). MonadIO m => m Context
getContext
  Propagator Context [Header] [Header]
propagator <- m (Propagator Context [Header] [Header])
forall (m :: * -> *).
MonadTracer m =>
m (Propagator Context [Header] [Header])
getPropagator
  Context
updatedContext <- Propagator Context [Header] [Header]
-> [Header] -> Context -> m Context
forall (m :: * -> *) context i o.
MonadIO m =>
Propagator context i o -> i -> context -> m context
extract Propagator Context [Header] [Header]
propagator (a
a a -> Getting [Header] a [Header] -> [Header]
forall s a. s -> Getting a s a -> a
^. Getting [Header] a [Header]
forall a. HasHeaders a => Lens' a [Header]
Lens' a [Header]
headersL) Context
context
  m (Maybe Context) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Maybe Context) -> m ()) -> m (Maybe Context) -> m ()
forall a b. (a -> b) -> a -> b
$ Context -> m (Maybe Context)
forall (m :: * -> *). MonadIO m => Context -> m (Maybe Context)
attachContext Context
updatedContext

-- | Inject our trace context into the given item's headers
injectContext
  :: (MonadIO m, MonadTracer m, HasHeaders a) => a -> m a
injectContext :: forall (m :: * -> *) a.
(MonadIO m, MonadTracer m, HasHeaders a) =>
a -> m a
injectContext a
a = do
  Context
context <- m Context
forall (m :: * -> *). MonadIO m => m Context
getContext
  Propagator Context [Header] [Header]
propagator <- m (Propagator Context [Header] [Header])
forall (m :: * -> *).
MonadTracer m =>
m (Propagator Context [Header] [Header])
getPropagator
  [Header]
headers <- Propagator Context [Header] [Header]
-> Context -> [Header] -> m [Header]
forall (m :: * -> *) context i o.
MonadIO m =>
Propagator context i o -> context -> o -> m o
inject Propagator Context [Header] [Header]
propagator Context
context ([Header] -> m [Header]) -> [Header] -> m [Header]
forall a b. (a -> b) -> a -> b
$ a
a a -> Getting [Header] a [Header] -> [Header]
forall s a. s -> Getting a s a -> a
^. Getting [Header] a [Header]
forall a. HasHeaders a => Lens' a [Header]
Lens' a [Header]
headersL
  a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$ a
a a -> (a -> a) -> a
forall a b. a -> (a -> b) -> b
& ([Header] -> Identity [Header]) -> a -> Identity a
forall a. HasHeaders a => Lens' a [Header]
Lens' a [Header]
headersL (([Header] -> Identity [Header]) -> a -> Identity a)
-> [Header] -> a -> a
forall s t a b. ASetter s t a b -> b -> s -> t
.~ [Header]
headers

getPropagator :: MonadTracer m => m (Propagator Context [Header] [Header])
getPropagator :: forall (m :: * -> *).
MonadTracer m =>
m (Propagator Context [Header] [Header])
getPropagator =
  TracerProvider -> Propagator Context [Header] [Header]
getTracerProviderPropagators (TracerProvider -> Propagator Context [Header] [Header])
-> (Tracer -> TracerProvider)
-> Tracer
-> Propagator Context [Header] [Header]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tracer -> TracerProvider
getTracerTracerProvider (Tracer -> Propagator Context [Header] [Header])
-> m Tracer -> m (Propagator Context [Header] [Header])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Tracer
forall (m :: * -> *). MonadTracer m => m Tracer
getTracer

-- | Process an item (a request, a Job, etc) in a top-level span and context
processWithContext
  :: ( MonadUnliftIO m
     , MonadMask m
     , MonadTracer m
     , HasHeaders a
     , HasCallStack
     )
  => Text
  -- ^ Span name
  -> SpanArguments
  -> a
  -> (a -> m b)
  -> m b
processWithContext :: forall (m :: * -> *) a b.
(MonadUnliftIO m, MonadMask m, MonadTracer m, HasHeaders a,
 HasCallStack) =>
Text -> SpanArguments -> a -> (a -> m b) -> m b
processWithContext Text
name SpanArguments
args a
a a -> m b
f = do
  a -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadTracer m, HasHeaders a) =>
a -> m ()
extractContext a
a
  Text -> SpanArguments -> m b -> m b
forall (m :: * -> *) a.
(MonadUnliftIO m, MonadTracer m, HasCallStack) =>
Text -> SpanArguments -> m a -> m a
inSpan Text
name SpanArguments
args (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ do
    a
a' <- a -> m a
forall (m :: * -> *) a.
(MonadIO m, MonadTracer m, HasHeaders a) =>
a -> m a
injectContext a
a
    m b -> m b
forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
withTraceContext (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ a -> m b
f a
a'