diff --git a/shibuya-pgmq-adapter.cabal b/shibuya-pgmq-adapter.cabal --- a/shibuya-pgmq-adapter.cabal +++ b/shibuya-pgmq-adapter.cabal @@ -1,6 +1,6 @@ cabal-version: 3.14 name: shibuya-pgmq-adapter -version: 0.4.0.0 +version: 0.5.0.0 synopsis: PGMQ adapter for the Shibuya queue processing framework description: A Shibuya adapter that integrates with pgmq (PostgreSQL Message Queue) @@ -46,12 +46,13 @@ pgmq-core ^>=0.2, pgmq-effectful ^>=0.2, pgmq-hasql ^>=0.2, - shibuya-core ^>=0.4.0.0, + shibuya-core ^>=0.5.0.0, stm ^>=2.5, streamly ^>=0.11, streamly-core ^>=0.3, text ^>=2.1.3, time ^>=1.14, + unordered-containers ^>=0.2, vector ^>=0.13, hs-source-dirs: src @@ -113,11 +114,12 @@ pgmq-migration ^>=0.2, quickcheck-instances ^>=0.3, random, - shibuya-core ^>=0.4.0.0, + shibuya-core ^>=0.5.0.0, shibuya-pgmq-adapter, stm, streamly ^>=0.11, streamly-core ^>=0.3, text, time, + unordered-containers ^>=0.2, vector, diff --git a/src/Shibuya/Adapter/Pgmq.hs b/src/Shibuya/Adapter/Pgmq.hs --- a/src/Shibuya/Adapter/Pgmq.hs +++ b/src/Shibuya/Adapter/Pgmq.hs @@ -150,6 +150,7 @@ topicDeadLetter, ) import Shibuya.Adapter.Pgmq.Internal (pgmqSource, pgmqSourceWithPrefetch) +import Shibuya.Telemetry.Effect (Tracing) import Streamly.Data.Stream (Stream) import Streamly.Data.Stream qualified as Stream import Streamly.Data.Stream.Prelude qualified as StreamP @@ -189,7 +190,7 @@ -- adapter <- pgmqAdapter config -- @ pgmqAdapter :: - (Pgmq :> es, IOE :> es) => + (Pgmq :> es, IOE :> es, Tracing :> es) => PgmqAdapterConfig -> Eff es (Adapter es Value) pgmqAdapter config = do diff --git a/src/Shibuya/Adapter/Pgmq/Convert.hs b/src/Shibuya/Adapter/Pgmq/Convert.hs --- a/src/Shibuya/Adapter/Pgmq/Convert.hs +++ b/src/Shibuya/Adapter/Pgmq/Convert.hs @@ -19,6 +19,7 @@ import Data.Aeson (Value (..), object, (.=)) import Data.Aeson.Key qualified as Key import Data.Aeson.KeyMap qualified as KeyMap +import Data.HashMap.Strict qualified as HashMap import Data.Int (Int64) import Data.Text (Text) import Data.Text qualified as Text @@ -81,6 +82,15 @@ -- The payload is the raw JSON Value from pgmq. -- Extracts W3C trace context from headers if present. -- Populates the delivery 'attempt' counter from pgmq's 'readCount'. +-- +-- 'Envelope.attributes' is left empty: pgmq has no spec-defined typed +-- messaging-attribute conventions in OpenTelemetry semantic-conventions +-- v1.27. The framework's @processOne@ already sets the standard +-- @messaging.system="shibuya"@ default plus the spec-aligned +-- @messaging.destination.name@ / @messaging.operation@ / +-- @messaging.message.id@ from the @ProcessorId@ and envelope's +-- @MessageId@. The field is a forward-compatible hook for the day a +-- @messaging.pgmq.*@ convention is defined upstream. pgmqMessageToEnvelope :: Pgmq.Message -> Envelope Value pgmqMessageToEnvelope msg = Envelope @@ -90,6 +100,7 @@ enqueuedAt = Just msg.enqueuedAt, traceContext = extractTraceHeaders msg.headers, attempt = Just (readCountToAttempt msg.readCount), + attributes = HashMap.empty, payload = Pgmq.unMessageBody msg.body } diff --git a/src/Shibuya/Adapter/Pgmq/Internal.hs b/src/Shibuya/Adapter/Pgmq/Internal.hs --- a/src/Shibuya/Adapter/Pgmq/Internal.hs +++ b/src/Shibuya/Adapter/Pgmq/Internal.hs @@ -14,6 +14,7 @@ -- * AckHandle Construction mkAckHandle, + mergeDlqHeaders, -- * Lease Construction mkLease, @@ -32,10 +33,13 @@ import Control.Concurrent (threadDelay) import Control.Monad (when) import Control.Monad.IO.Class (liftIO) -import Data.Aeson (Value) +import Data.Aeson (Value (..)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KeyMap import Data.Function ((&)) import Data.Int (Int32) import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE import Data.Time (NominalDiffTime, nominalDiffTimeToSeconds) import Data.Vector (Vector) import Data.Vector qualified as Vector @@ -85,6 +89,9 @@ import Shibuya.Core.AckHandle (AckHandle (..)) import Shibuya.Core.Ingested (Ingested (..)) import Shibuya.Core.Lease (Lease (..)) +import Shibuya.Core.Types (TraceHeaders) +import Shibuya.Telemetry.Effect (Tracing) +import Shibuya.Telemetry.Propagation (currentTraceHeaders) import Streamly.Data.Stream (Stream) import Streamly.Data.Stream qualified as Stream import Streamly.Data.Stream.Prelude qualified as StreamP @@ -175,8 +182,19 @@ } -- | Create an AckHandle for a message. +-- +-- The 'AckDeadLetter' branch threads the *consumer's* current trace +-- context (looked up via 'currentTraceHeaders' against the active OTel +-- span) into the DLQ message's headers. The original producer's +-- @traceparent@/@tracestate@ are preserved under the +-- @x-shibuya-upstream-traceparent@ / @x-shibuya-upstream-tracestate@ +-- keys so a DLQ post-mortem can walk back to the origin if it wants. +-- When tracing is disabled (or there is no active span at the call +-- site), the original headers are forwarded verbatim — exactly the +-- pre-0.5.0.0 behavior. See plan 1 / Finding F3 in the parent +-- shibuya repo's plan 9. mkAckHandle :: - (Pgmq :> es, IOE :> es) => + (Pgmq :> es, IOE :> es, Tracing :> es) => PgmqAdapterConfig -> Pgmq.Message -> AckHandle es @@ -205,12 +223,17 @@ -- No DLQ configured - just archive the message void $ archiveMessage (MessageQuery queueName msgId) Just dlqConfig -> do - -- Send to DLQ with metadata, preserving trace context if present + -- Build DLQ headers: pull the consumer's current trace + -- context (Nothing if tracing is off or no active span); + -- merge with the original message's headers (consumer's + -- traceparent wins, original preserved under the + -- x-shibuya-upstream-* keys). + consumerHdrs <- currentTraceHeaders let dlqBody = mkDlqPayload msg reason dlqConfig.includeMetadata + dlqHeaders = mergeDlqHeaders consumerHdrs msg.headers case dlqConfig.dlqTarget of DirectQueue dlqQueueName -> - -- Send directly to a specific DLQ - case msg.headers of + case dlqHeaders of Just headers -> void $ sendMessageWithHeaders $ @@ -229,8 +252,7 @@ delay = Nothing } TopicRoute routingKey -> - -- Route via topic pattern matching (pgmq 1.11.0+) - case msg.headers of + case dlqHeaders of Just headers -> void $ sendTopicWithHeaders $ @@ -265,10 +287,62 @@ void :: (Functor f) => f a -> f () void = fmap (const ()) +-- | Merge the consumer's current trace headers with the original +-- message's headers JSON for the DLQ-write path. +-- +-- Rules: +-- +-- * If the consumer has no active span (tracing disabled, or +-- producer-side path runs outside any 'withSpan'), forward the +-- original headers verbatim — matches the pre-0.5.0.0 behavior. +-- * Otherwise, the consumer's @traceparent@ overrides the original's +-- active @traceparent@; the original's @traceparent@ / +-- @tracestate@ (if present) move to +-- @x-shibuya-upstream-traceparent@ / @x-shibuya-upstream-tracestate@. +-- +-- Returns 'Nothing' only if both inputs are empty (no consumer +-- context AND no original headers); in that case the caller falls +-- through to the no-headers @sendMessage@/@sendTopic@ path. +mergeDlqHeaders :: Maybe TraceHeaders -> Maybe Value -> Maybe Value +mergeDlqHeaders Nothing originalHeaders = originalHeaders +mergeDlqHeaders (Just consumerHdrs) originalHeaders = + let originalObj = case originalHeaders of + Just (Object obj) -> obj + _ -> KeyMap.empty + stashedUpstream = stashUpstreamTrace originalObj + consumerEntries = traceHeadersToKeyMap consumerHdrs + merged = stashedUpstream <> consumerEntries + in if KeyMap.null merged + then Nothing + else Just (Object merged) + where + -- Move any active @traceparent@/@tracestate@ on the original + -- headers under the @x-shibuya-upstream-*@ prefix so the + -- consumer's value can take the active slot. Other keys pass + -- through unchanged. + stashUpstreamTrace obj = + foldr + (uncurry (rename obj)) + (KeyMap.delete "traceparent" (KeyMap.delete "tracestate" obj)) + [ ("traceparent", "x-shibuya-upstream-traceparent"), + ("tracestate", "x-shibuya-upstream-tracestate") + ] + rename src srcKey dstKey acc = + case KeyMap.lookup (Key.fromText srcKey) src of + Just v -> KeyMap.insert (Key.fromText dstKey) v acc + Nothing -> acc + + -- Convert TraceHeaders ([(ByteString, ByteString)]) to a JSON object. + traceHeadersToKeyMap hdrs = + KeyMap.fromList + [ (Key.fromText (TE.decodeUtf8 k), String (TE.decodeUtf8 v)) + | (k, v) <- hdrs + ] + -- | Create an Ingested from a pgmq Message. -- Handles auto dead-lettering when maxRetries is exceeded. mkIngested :: - (Pgmq :> es, IOE :> es) => + (Pgmq :> es, IOE :> es, Tracing :> es) => PgmqAdapterConfig -> Pgmq.Message -> Eff es (Maybe (Ingested es Value)) @@ -354,7 +428,7 @@ -- This stream polls pgmq and yields Ingested messages. -- Uses unfoldEach to process ALL messages from each batch, not just the first. pgmqSource :: - (Pgmq :> es, IOE :> es) => + (Pgmq :> es, IOE :> es, Tracing :> es) => PgmqAdapterConfig -> Stream (Eff es) (Ingested es Value) pgmqSource config = @@ -408,7 +482,7 @@ -- source = pgmqSourceWithPrefetch (StreamP.maxBuffer 2) config -- @ pgmqSourceWithPrefetch :: - (Pgmq :> es, IOE :> es) => + (Pgmq :> es, IOE :> es, Tracing :> es) => (StreamP.Config -> StreamP.Config) -> PgmqAdapterConfig -> Stream (Eff es) (Ingested es Value) diff --git a/test/Shibuya/Adapter/Pgmq/InternalSpec.hs b/test/Shibuya/Adapter/Pgmq/InternalSpec.hs --- a/test/Shibuya/Adapter/Pgmq/InternalSpec.hs +++ b/test/Shibuya/Adapter/Pgmq/InternalSpec.hs @@ -1,5 +1,9 @@ +{-# LANGUAGE OverloadedStrings #-} + module Shibuya.Adapter.Pgmq.InternalSpec (spec) where +import Data.Aeson (Value (..), object, (.=)) +import Data.Aeson.KeyMap qualified as KeyMap import Data.Int (Int32) import Data.Time (NominalDiffTime) import Pgmq.Hasql.Statements.Types (ReadGrouped (..), ReadMessage (..), ReadWithPollMessage (..)) @@ -9,7 +13,8 @@ defaultPollingConfig, ) import Shibuya.Adapter.Pgmq.Internal - ( mkReadGrouped, + ( mergeDlqHeaders, + mkReadGrouped, mkReadMessage, mkReadWithPoll, nominalToSeconds, @@ -22,6 +27,7 @@ mkReadMessageSpec mkReadWithPollSpec mkReadGroupedSpec + mergeDlqHeadersSpec -- | Tests for nominalToSeconds nominalToSecondsSpec :: Spec @@ -167,3 +173,83 @@ it "sets qty to batchSize" $ do queryQty `shouldBe` 20 + +-- | Tests for mergeDlqHeaders, the helper that injects the failing +-- consumer's trace context onto a DLQ message while preserving the +-- original producer's trace under x-shibuya-upstream-* keys. +mergeDlqHeadersSpec :: Spec +mergeDlqHeadersSpec = describe "mergeDlqHeaders" $ do + it "with no consumer headers, forwards original headers verbatim" $ do + let original = + Just $ + object + [ "traceparent" .= ("00-producer-trace-id-pid-01" :: String), + "tracestate" .= ("vendor=opaque" :: String), + "custom" .= ("value" :: String) + ] + mergeDlqHeaders Nothing original `shouldBe` original + + it "with no consumer headers and no original, returns Nothing" $ do + mergeDlqHeaders Nothing Nothing `shouldBe` Nothing + + it "consumer's traceparent overrides original's, original moves under x-shibuya-upstream-traceparent" $ do + let original = + Just $ + object + [ "traceparent" .= ("00-producer-trace-id-pid-01" :: String), + "tracestate" .= ("vendor=opaque" :: String), + "custom" .= ("preserved" :: String) + ] + consumerHdrs = + Just + [ ("traceparent", "00-consumer-trace-id-cid-01"), + ("tracestate", "consumer=ok") + ] + case mergeDlqHeaders consumerHdrs original of + Just (Object obj) -> do + KeyMap.lookup "traceparent" obj + `shouldBe` Just (String "00-consumer-trace-id-cid-01") + KeyMap.lookup "tracestate" obj + `shouldBe` Just (String "consumer=ok") + KeyMap.lookup "x-shibuya-upstream-traceparent" obj + `shouldBe` Just (String "00-producer-trace-id-pid-01") + KeyMap.lookup "x-shibuya-upstream-tracestate" obj + `shouldBe` Just (String "vendor=opaque") + KeyMap.lookup "custom" obj + `shouldBe` Just (String "preserved") + other -> expectationFailure $ "expected merged Object, got " <> show other + + it "consumer's headers carry through when original headers are absent" $ do + let consumerHdrs = + Just + [ ("traceparent", "00-consumer-trace-id-cid-01"), + ("tracestate", "consumer=ok") + ] + case mergeDlqHeaders consumerHdrs Nothing of + Just (Object obj) -> do + KeyMap.lookup "traceparent" obj + `shouldBe` Just (String "00-consumer-trace-id-cid-01") + KeyMap.lookup "tracestate" obj + `shouldBe` Just (String "consumer=ok") + KeyMap.lookup "x-shibuya-upstream-traceparent" obj `shouldBe` Nothing + KeyMap.lookup "x-shibuya-upstream-tracestate" obj `shouldBe` Nothing + other -> expectationFailure $ "expected merged Object, got " <> show other + + it "original's tracestate is stashed even when consumer has no tracestate" $ do + let original = + Just $ + object + [ "traceparent" .= ("00-producer-trace-id-pid-01" :: String), + "tracestate" .= ("vendor=opaque" :: String) + ] + consumerHdrs = Just [("traceparent", "00-consumer-trace-id-cid-01")] + case mergeDlqHeaders consumerHdrs original of + Just (Object obj) -> do + KeyMap.lookup "traceparent" obj + `shouldBe` Just (String "00-consumer-trace-id-cid-01") + KeyMap.lookup "tracestate" obj `shouldBe` Nothing + KeyMap.lookup "x-shibuya-upstream-traceparent" obj + `shouldBe` Just (String "00-producer-trace-id-pid-01") + KeyMap.lookup "x-shibuya-upstream-tracestate" obj + `shouldBe` Just (String "vendor=opaque") + other -> expectationFailure $ "expected merged Object, got " <> show other