pgmq-effectful-0.2.0.0: test/TracedInterpreterSpec.hs
{-# LANGUAGE OverloadedStrings #-}
module TracedInterpreterSpec (tests) where
import Control.Monad (unless)
import Control.Monad.IO.Class (liftIO)
import Data.HashMap.Strict qualified as HM
import Data.IORef (IORef, readIORef)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Vector qualified as V
import Effectful (runEff)
import Effectful.Error.Static (runError)
import Hasql.Pool qualified as Pool
import OpenTelemetry.Attributes qualified as Attrs
import OpenTelemetry.Context qualified as Ctxt
import OpenTelemetry.Context.ThreadLocal qualified as CtxtLocal
import OpenTelemetry.Exporter.InMemory.Span qualified as InMemoryExporter
import OpenTelemetry.Propagator (extract)
import OpenTelemetry.Propagator.W3CTraceContext qualified as W3C
import OpenTelemetry.Trace.Core qualified as OTel
import OpenTelemetry.Trace.Id.Generator.Default (defaultIdGenerator)
import OpenTelemetry.Util (appendOnlyBoundedCollectionValues)
import Pgmq.Effectful
( MessageBody (..),
MessageId (..),
MessageQuery (..),
PgmqRuntimeError (..),
QueueName,
ReadMessage (..),
SendMessage (..),
createQueue,
deleteMessage,
parseQueueName,
queueNameToText,
readMessage,
runPgmqTraced,
sendMessage,
sendMessageTraced,
)
import Pgmq.Effectful.Telemetry (jsonToTraceHeaders)
import Pgmq.Types qualified as Pgmq
import System.Random (randomRIO)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase)
tests :: Pool.Pool -> TestTree
tests pool =
testGroup
"Traced interpreter"
[ testGroup
"Error propagation"
[ testCase "statement error surfaces PgmqSessionError via Error channel" $ do
(tracer, _provider, _spansRef) <- setupTracer
bogus <- mkQueueName "queue_that_does_not_exist_xyz2"
result <-
runEff
. runError @PgmqRuntimeError
. runPgmqTraced pool tracer
$ deleteMessage
MessageQuery
{ queueName = bogus,
messageId = MessageId 1
}
case result of
Left (_cs, PgmqSessionError _) -> pure ()
Left (_cs, other) ->
assertFailure $
"Expected PgmqSessionError, got " <> show other
Right _ ->
assertFailure
"Expected an error, got success"
],
testGroup
"v1.24 semantic conventions"
[ testCase "publish emits messaging.* + db.* attributes" $ do
(tracer, _provider, spansRef) <- setupTracer
queue <- mkUniqueQueue "publish_attrs"
result <-
runEff . runError @PgmqRuntimeError . runPgmqTraced pool tracer $ do
createQueue queue
sendMessage
SendMessage
{ queueName = queue,
messageBody = MessageBody "hello",
delay = Nothing
}
_ <- assertRight result
spans <- readIORef spansRef
let publishSpans = filter ((==) "publish" . firstWord . OTel.spanName) spans
unless (length publishSpans == 1) $
assertFailure $
"Expected exactly one publish span, got "
<> show (map OTel.spanName spans)
s <- singleSpan publishSpans "publish"
assertEqual
"span name carries destination"
("publish " <> queueNameToText queue)
(OTel.spanName s)
assertSpanKindProducer s
assertAttrText s "messaging.system" "pgmq"
assertAttrText s "messaging.operation" "publish"
assertAttrText s "messaging.destination.name" (queueNameToText queue)
assertAttrText s "db.system" "postgresql"
assertAttrText s "db.operation" "pgmq.send",
testCase "receive emits messaging.* + db.* attributes" $ do
(tracer, _provider, spansRef) <- setupTracer
queue <- mkUniqueQueue "receive_attrs"
result <-
runEff . runError @PgmqRuntimeError . runPgmqTraced pool tracer $ do
createQueue queue
_ <-
sendMessage
SendMessage
{ queueName = queue,
messageBody = MessageBody "one",
delay = Nothing
}
readMessage
ReadMessage
{ queueName = queue,
delay = 30,
batchSize = Just 1,
conditional = Nothing
}
_ <- assertRight result
spans <- readIORef spansRef
let receiveSpans = filter ((==) "receive" . firstWord . OTel.spanName) spans
unless (length receiveSpans == 1) $
assertFailure $
"Expected exactly one receive span, got "
<> show (map OTel.spanName spans)
s <- singleSpan receiveSpans "receive"
assertEqual
"span name carries destination"
("receive " <> queueNameToText queue)
(OTel.spanName s)
assertSpanKindConsumer s
assertAttrText s "messaging.system" "pgmq"
assertAttrText s "messaging.operation" "receive"
assertAttrText s "messaging.destination.name" (queueNameToText queue)
assertAttrText s "db.system" "postgresql"
assertAttrText s "db.operation" "pgmq.read",
testCase "error path records exception event and Error status" $ do
(tracer, _provider, spansRef) <- setupTracer
bogus <- mkQueueName "queue_that_does_not_exist_xyz3"
_ <-
runEff
. runError @PgmqRuntimeError
. runPgmqTraced pool tracer
$ deleteMessage
MessageQuery {queueName = bogus, messageId = MessageId 1}
spans <- readIORef spansRef
s <- case spans of
[only] -> pure only
_ ->
assertFailure
( "Expected exactly one span, got "
<> show (map OTel.spanName spans)
)
>> fail ""
let events =
V.toList
(appendOnlyBoundedCollectionValues (OTel.spanEvents s))
hasExceptionEvent =
any ((==) "exception" . OTel.eventName) events
assertBool "span has exception event" hasExceptionEvent
case OTel.spanStatus s of
OTel.Error desc ->
assertBool
( "span status description should be a short non-PII label, got "
<> T.unpack desc
)
("pool.session.statement" `T.isPrefixOf` desc)
other ->
assertFailure $ "expected Error status, got " <> show other,
testCase "W3C traceparent round-trips through message headers" $ do
(tracer, provider, _spansRef) <- setupTracer
queue <- mkUniqueQueue "w3c_roundtrip"
parentSpan <-
OTel.createSpan
tracer
Ctxt.empty
"parent"
OTel.defaultSpanArguments
parentTraceId <- OTel.traceId <$> OTel.getSpanContext parentSpan
_ <- CtxtLocal.attachContext (Ctxt.insertSpan parentSpan Ctxt.empty)
result <-
runEff . runError @PgmqRuntimeError . runPgmqTraced pool tracer $ do
createQueue queue
_ <- sendMessageTraced provider queue (MessageBody "payload") Nothing
readMessage
ReadMessage
{ queueName = queue,
delay = 30,
batchSize = Just 1,
conditional = Nothing
}
OTel.endSpan parentSpan Nothing
msgs <- assertRight result
msg <- case V.toList msgs of
[m] -> pure m
other ->
assertFailure ("expected 1 message, got " <> show (length other))
>> fail ""
case Pgmq.headers msg of
Nothing ->
assertFailure
"expected trace headers on received message, got none"
Just hdrs -> do
let propagator = OTel.getTracerProviderPropagators provider
ctx <-
liftIO $
extract propagator (jsonToTraceHeaders hdrs) Ctxt.empty
case Ctxt.lookupSpan ctx of
Nothing ->
assertFailure $
"propagator did not extract a span from headers: " <> show hdrs
Just extracted -> do
extractedTraceId <- OTel.traceId <$> OTel.getSpanContext extracted
assertEqual
"trace id round-trips via message headers"
parentTraceId
extractedTraceId
]
]
-- ---------------------------------------------------------------------
-- Helpers
-- ---------------------------------------------------------------------
-- | Build a tracer, tracer provider (with in-memory list processor and
-- W3C Trace Context propagator), and a handle to the exported spans.
setupTracer :: IO (OTel.Tracer, OTel.TracerProvider, IORef [OTel.ImmutableSpan])
setupTracer = do
(processor, spansRef) <- InMemoryExporter.inMemoryListExporter
let opts =
OTel.emptyTracerProviderOptions
{ OTel.tracerProviderOptionsIdGenerator = defaultIdGenerator,
OTel.tracerProviderOptionsPropagators = W3C.w3cTraceContextPropagator
}
provider <- OTel.createTracerProvider [processor] opts
let tracer =
OTel.makeTracer provider "pgmq-effectful-test" OTel.tracerOptions
pure (tracer, provider, spansRef)
mkQueueName :: String -> IO QueueName
mkQueueName raw =
case parseQueueName (T.pack raw) of
Right q -> pure q
Left err ->
assertFailure
("could not build queue name " <> show raw <> ": " <> show err)
>> fail ""
mkUniqueQueue :: String -> IO QueueName
mkUniqueQueue prefix = do
suffix <- randomRIO (100000 :: Int, 999999)
mkQueueName (prefix <> "_" <> show suffix)
firstWord :: Text -> Text
firstWord = T.takeWhile (/= ' ')
assertRight :: (Show a) => Either a b -> IO b
assertRight = \case
Right b -> pure b
Left e ->
assertFailure ("expected Right, got Left " <> show e) >> fail ""
singleSpan :: [OTel.ImmutableSpan] -> String -> IO OTel.ImmutableSpan
singleSpan spans label = case spans of
[s] -> pure s
_ ->
assertFailure
( "expected exactly one "
<> label
<> " span, got "
<> show (map OTel.spanName spans)
)
>> fail ""
-- | 'OTel.SpanKind' has no 'Eq' instance, so assert via pattern.
assertSpanKindProducer :: OTel.ImmutableSpan -> IO ()
assertSpanKindProducer s = case OTel.spanKind s of
OTel.Producer -> pure ()
other -> assertFailure ("expected Producer span kind, got " <> show other)
assertSpanKindConsumer :: OTel.ImmutableSpan -> IO ()
assertSpanKindConsumer s = case OTel.spanKind s of
OTel.Consumer -> pure ()
other -> assertFailure ("expected Consumer span kind, got " <> show other)
-- | Assert a span carries the given attribute as a 'TextAttribute' with
-- the given value. Produces a readable failure message with the full
-- attribute map when the assertion fails.
assertAttrText :: OTel.ImmutableSpan -> Text -> Text -> IO ()
assertAttrText s key expected = do
let attrMap = Attrs.getAttributeMap (OTel.spanAttributes s)
case HM.lookup key attrMap of
Just (Attrs.AttributeValue (Attrs.TextAttribute actual)) ->
assertEqual ("attribute " <> T.unpack key) expected actual
other ->
assertFailure $
"expected text attribute "
<> T.unpack key
<> "="
<> T.unpack expected
<> "; got "
<> show other
<> " in "
<> show attrMap