-- FUTUREWORK(mangoiv):
-- instrument http/2 request similarly to how it was done for http-client here:
-- https://github.com/iand675/hs-opentelemetry/blob/0b3c854a88113fc18df8561202a76357e593a294/instrumentation/http-client/src/OpenTelemetry/Instrumentation/HttpClient/Raw.hs#L60
-- This is non-trivial because http/2 forgets the structure on the out objs.
module Wire.OpenTelemetry
  ( -- * instrumentation helpers
    withTracer,
    withTracerC,

    -- * outbound instrumentation

    -- ** http client
    withClientInstrumentation,
  )
where

import Control.Monad.Codensity (Codensity (Codensity))
import Data.Text (Text)
import Network.HTTP.Client (Request, Response)
import OpenTelemetry.Context.ThreadLocal (getContext)
import OpenTelemetry.Instrumentation.HttpClient.Raw
import OpenTelemetry.Trace
import UnliftIO (MonadUnliftIO, bracket, liftIO)

-- | a tracer for a service like brig, galley, etc.
withTracer :: (MonadUnliftIO m) => (Tracer -> m r) -> m r
withTracer :: forall (m :: * -> *) r. MonadUnliftIO m => (Tracer -> m r) -> m r
withTracer Tracer -> m r
k =
  m TracerProvider
-> (TracerProvider -> m ()) -> (TracerProvider -> m r) -> m r
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket
    (IO TracerProvider -> m TracerProvider
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO TracerProvider
initializeGlobalTracerProvider)
    TracerProvider -> m ()
forall (m :: * -> *). MonadIO m => TracerProvider -> m ()
shutdownTracerProvider
    \TracerProvider
tp -> Tracer -> m r
k (Tracer -> m r) -> Tracer -> m r
forall a b. (a -> b) -> a -> b
$ TracerProvider -> InstrumentationLibrary -> TracerOptions -> Tracer
makeTracer TracerProvider
tp InstrumentationLibrary
"wire-otel" TracerOptions
tracerOptions

-- | like 'withTracer' but in 'Codensity'
withTracerC :: Codensity IO Tracer
withTracerC :: Codensity IO Tracer
withTracerC = (forall b. (Tracer -> IO b) -> IO b) -> Codensity IO Tracer
forall k (m :: k -> *) a.
(forall (b :: k). (a -> m b) -> m b) -> Codensity m a
Codensity (Tracer -> IO b) -> IO b
forall b. (Tracer -> IO b) -> IO b
forall (m :: * -> *) r. MonadUnliftIO m => (Tracer -> m r) -> m r
withTracer

-- | instrument a http client
withClientInstrumentation ::
  (MonadUnliftIO m) =>
  -- | name of the caller
  Text ->
  -- | continuation that takes a continuation that takes a request and a way to respond to a request
  ((Request -> (Request -> m (Response a)) -> m (Response a)) -> m b) ->
  m b
withClientInstrumentation :: forall (m :: * -> *) a b.
MonadUnliftIO m =>
Text
-> ((Request -> (Request -> m (Response a)) -> m (Response a))
    -> m b)
-> m b
withClientInstrumentation Text
info (Request -> (Request -> m (Response a)) -> m (Response a)) -> m b
k = do
  Tracer
tracer <- m Tracer
forall (m :: * -> *). MonadIO m => m Tracer
httpTracerProvider
  Tracer -> Text -> SpanArguments -> m b -> m b
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpan Tracer
tracer Text
info SpanArguments
defaultSpanArguments {kind = Client} do
    Context
otelCtx <- m Context
forall (m :: * -> *). MonadIO m => m Context
getContext
    (Request -> (Request -> m (Response a)) -> m (Response a)) -> m b
k \Request
req Request -> m (Response a)
respond -> do
      Response a
resp <- Request -> m (Response a)
respond (Request -> m (Response a)) -> m Request -> m (Response a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HttpClientInstrumentationConfig -> Context -> Request -> m Request
forall (m :: * -> *).
MonadIO m =>
HttpClientInstrumentationConfig -> Context -> Request -> m Request
instrumentRequest HttpClientInstrumentationConfig
httpClientInstrumentationConfig Context
otelCtx Request
req
      HttpClientInstrumentationConfig -> Context -> Response a -> m ()
forall (m :: * -> *) a.
MonadIO m =>
HttpClientInstrumentationConfig -> Context -> Response a -> m ()
instrumentResponse HttpClientInstrumentationConfig
httpClientInstrumentationConfig Context
otelCtx Response a
resp
      Response a -> m (Response a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Response a
resp