Ensure Kafka consumers include trace_id in log context

CHANGELOG.md

@@ -1,4 +1,11 @@
-## [_Unreleased_](https://github.com/freckle/freckle-app/compare/v1.10.0.0...main)
+## [_Unreleased_](https://github.com/freckle/freckle-app/compare/v1.10.1.0...main)
+
+## [v1.10.0.1](https://github.com/freckle/freckle-app/compare/v1.10.0.0...v1.10.1.0)
+
+- Use `withTraceIdContext` in `Freckle.App.Kafka.Consumer.runConsumer`, ensuring
+  all logging contains the `trace_id` in context.
+- Add `getCurrentTraceIdAsDatadog` and `withTraceIdContext` to
+  `Freckle.App.OpenTelemetry`.
 
 ## [v1.10.0.0](https://github.com/freckle/freckle-app/compare/v1.9.5.1...v1.10.0.0)
 

freckle-app.cabal

@@ -5,7 +5,7 @@ cabal-version: 1.18
 -- see: https://github.com/sol/hpack
 
 name:           freckle-app
-version:        1.10.0.0
+version:        1.10.1.0
 synopsis:       Haskell application toolkit used at Freckle
 description:    Please see README.md
 category:       Utils

library/Freckle/App/Kafka/Consumer.hs

@@ -20,6 +20,7 @@ import qualified Env
 import Freckle.App.Async
 import Freckle.App.Env
 import Freckle.App.Kafka.Producer (envKafkaBrokerAddresses)
+import Freckle.App.OpenTelemetry (withTraceIdContext)
 import Kafka.Consumer hiding
   ( Timeout
   , closeConsumer
@@ -153,14 +154,15 @@ runConsumer
   => Timeout
   -> (a -> m ())
   -> m ()
-runConsumer pollTimeout onMessage = immortalCreateLogged $ do
-  consumer <- view kafkaConsumerL
-  eMessage <-
-    pollMessage consumer $ Kafka.Timeout $ timeoutMs pollTimeout
-  case eMessage of
-    Left (KafkaResponseError RdKafkaRespErrTimedOut) -> logDebug "Polling timeout"
-    Left err -> logError $ "Error polling for message from Kafka" :# ["error" .= show err]
-    Right ConsumerRecord {..} -> for_ crValue $ \bs ->
-      case eitherDecodeStrict bs of
-        Left err -> logError $ "Could not decode message value" :# ["error" .= err]
-        Right a -> onMessage a
+runConsumer pollTimeout onMessage =
+  withTraceIdContext $ immortalCreateLogged $ do
+    consumer <- view kafkaConsumerL
+    eMessage <-
+      pollMessage consumer $ Kafka.Timeout $ timeoutMs pollTimeout
+    case eMessage of
+      Left (KafkaResponseError RdKafkaRespErrTimedOut) -> logDebug "Polling timeout"
+      Left err -> logError $ "Error polling for message from Kafka" :# ["error" .= show err]
+      Right ConsumerRecord {..} -> for_ crValue $ \bs ->
+        case eitherDecodeStrict bs of
+          Left err -> logError $ "Could not decode message value" :# ["error" .= err]
+          Right a -> onMessage a

library/Freckle/App/OpenTelemetry.hs

@@ -34,7 +34,9 @@ module Freckle.App.OpenTelemetry
   , defaultSpanArguments
 
     -- * Querying
+  , withTraceIdContext
   , getCurrentTraceId
+  , getCurrentTraceIdAsDatadog
   , getCurrentSpanContext
 
     -- * Setup
@@ -47,8 +49,13 @@ module Freckle.App.OpenTelemetry
 
 import Freckle.App.Prelude
 
+import Blammo.Logging (MonadMask, withThreadContext, (.=))
+import Data.Word (Word64)
 import OpenTelemetry.Context (lookupSpan)
 import OpenTelemetry.Context.ThreadLocal (getContext)
+import OpenTelemetry.Propagator.Datadog
+  ( convertOpenTelemetryTraceIdToDatadogTraceId
+  )
 import OpenTelemetry.Trace hiding (inSpan)
 import OpenTelemetry.Trace.Core (getSpanContext)
 import qualified OpenTelemetry.Trace.Core as Trace (SpanContext (..))
@@ -65,7 +72,16 @@ withTracerProvider =
 getCurrentTraceId :: MonadIO m => m (Maybe TraceId)
 getCurrentTraceId = fmap Trace.traceId <$> getCurrentSpanContext
 
+getCurrentTraceIdAsDatadog :: MonadIO m => m (Maybe Word64)
+getCurrentTraceIdAsDatadog =
+  fmap convertOpenTelemetryTraceIdToDatadogTraceId <$> getCurrentTraceId
+
 getCurrentSpanContext :: MonadIO m => m (Maybe SpanContext)
 getCurrentSpanContext = do
   mSpan <- lookupSpan <$> getContext
   traverse getSpanContext mSpan
+
+withTraceIdContext :: (MonadIO m, MonadMask m) => m a -> m a
+withTraceIdContext f = do
+  mTraceId <- getCurrentTraceIdAsDatadog
+  maybe f (\traceId -> withThreadContext ["trace_id" .= traceId] f) mTraceId

library/Network/Wai/Middleware/OpenTelemetry.hs

@@ -5,15 +5,11 @@ module Network.Wai.Middleware.OpenTelemetry
 
 import Freckle.App.Prelude
 
-import Blammo.Logging (withThreadContext, (.=))
 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
-import OpenTelemetry.Propagator.Datadog
-  ( convertOpenTelemetryTraceIdToDatadogTraceId
-  )
 
 newOpenTelemetryWaiMiddleware :: IO Middleware
 newOpenTelemetryWaiMiddleware = do
@@ -27,18 +23,13 @@ newOpenTelemetryWaiMiddleware = do
 --
 -- This is added automatically by our 'newOpenTelemetryWaiMiddleware'.
 openTelemetryMiddleware :: Middleware
-openTelemetryMiddleware app request respond = do
-  mTraceId <- getCurrentTraceId
-  case mTraceId of
-    Nothing -> app request respond
-    Just traceId -> do
-      let
-        traceIdInt = convertOpenTelemetryTraceIdToDatadogTraceId traceId
-        traceIdIntBS = BS8.pack $ show traceIdInt
+openTelemetryMiddleware app request respond =
+  withTraceIdContext $ addTraceIdHeader app request respond
 
-      withThreadContext ["trace_id" .= traceIdInt] $
-        addHeaders
-          [("X-Datadog-Trace-Id", traceIdIntBS)]
-          app
-          request
-          respond
+addTraceIdHeader :: Middleware
+addTraceIdHeader app request respond = do
+  mTraceId <- getCurrentTraceIdAsDatadog
+  maybe id addHeader mTraceId app request respond
+ where
+  addHeader traceId =
+    addHeaders [("X-Datadog-Trace-Id", BS8.pack $ show traceId)]

package.yaml

@@ -1,5 +1,5 @@
 name: freckle-app
-version: 1.10.0.0
+version: 1.10.1.0
 maintainer: Freckle Education
 category: Utils
 github: freckle/freckle-app