module Network.Wai.Middleware.OpenTelemetry
  ( newOpenTelemetryWaiMiddleware
  , openTelemetryMiddleware
  ) where

import Freckle.App.Prelude

import qualified Data.ByteString.Char8 as BS8
import Freckle.App.OpenTelemetry
import Network.Wai
import Network.Wai.Middleware.AddHeaders
import qualified OpenTelemetry.Instrumentation.Wai as Trace

newOpenTelemetryWaiMiddleware :: IO Middleware
newOpenTelemetryWaiMiddleware :: IO Middleware
newOpenTelemetryWaiMiddleware = do
  Middleware
otel <- IO Middleware
HasCallStack => IO Middleware
Trace.newOpenTelemetryWaiMiddleware
  Middleware -> IO Middleware
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Middleware -> IO Middleware) -> Middleware -> IO Middleware
forall a b. (a -> b) -> a -> b
$ Middleware
otel Middleware -> Middleware -> Middleware
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Middleware
openTelemetryMiddleware

-- | Add 'TraceId' information to context and responses
--
-- - Adds @trace_id@ to the logging context
-- - Adds @X-Datadog-TraceId@ to response headers
--
-- This is added automatically by our 'newOpenTelemetryWaiMiddleware'.
openTelemetryMiddleware :: Middleware
openTelemetryMiddleware :: Middleware
openTelemetryMiddleware Application
app Request
request Response -> IO ResponseReceived
respond =
  IO ResponseReceived -> IO ResponseReceived
forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
withTraceIdContext (IO ResponseReceived -> IO ResponseReceived)
-> IO ResponseReceived -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ Middleware
addTraceIdHeader Application
app Request
request Response -> IO ResponseReceived
respond

addTraceIdHeader :: Middleware
addTraceIdHeader :: Middleware
addTraceIdHeader Application
app Request
request Response -> IO ResponseReceived
respond = do
  Maybe Word64
mTraceId <- IO (Maybe Word64)
forall (m :: * -> *). MonadIO m => m (Maybe Word64)
getCurrentTraceIdAsDatadog
  Middleware -> (Word64 -> Middleware) -> Maybe Word64 -> Middleware
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Middleware
forall a. a -> a
id Word64 -> Middleware
forall {a}. Show a => a -> Middleware
addHeader Maybe Word64
mTraceId Application
app Request
request Response -> IO ResponseReceived
respond
 where
  addHeader :: a -> Middleware
addHeader a
traceId =
    [(ByteString, ByteString)] -> Middleware
addHeaders [(ByteString
"X-Datadog-Trace-Id", String -> ByteString
BS8.pack (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ a -> String
forall a. Show a => a -> String
show a
traceId)]