diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,40 @@ +# Changelog + +All notable changes to `keiro-pgmq` are recorded here. The format follows +[Keep a Changelog](https://keepachangelog.com/), and the package follows the +[Haskell Package Versioning Policy](https://pvp.haskell.org/). + +## [Unreleased] + +_No unreleased changes._ + +## 0.2.0.0 — 2026-07-13 + +No user-facing changes. `keiro-pgmq` is released at 0.2.0.0 to stay in lockstep +with the rest of the Keiro package set, and is rebuilt against `keiro-core` +0.2.0.0. Consumers should note that the stricter event-stream validation and the +relocation of Keiro's framework tables into the dedicated `keiro` PostgreSQL +schema — both described in the `keiro-core` and `keiro` changelogs — apply to any +application that wires this package's workqueue and dispatch workers. + +## 0.1.0.0 — 2026-07-05 + +Initial Hackage release. + +### New Features + +- Added typed PGMQ job definitions, payload codecs, runtime workers, job worker + context, retry and dead-letter policies, and direct job draining. +- Added queue provisioning, FIFO ordered delivery via message groups, producer + headers, batch enqueue support, queue metrics, archive/retention APIs, DLQ + inspection, and redrive helpers. +- Added trace propagation through PGMQ job producer and worker paths. + +### Bug Fixes + +- Classified job decode failures, validated retry tuning, disambiguated long + queue names, and isolated integration-test databases. + +### Other Changes + +- Tightened the internal `keiro-core` dependency bound to `^>=0.1.0.0`. diff --git a/keiro-pgmq.cabal b/keiro-pgmq.cabal new file mode 100644 --- /dev/null +++ b/keiro-pgmq.cabal @@ -0,0 +1,97 @@ +cabal-version: 3.0 +name: keiro-pgmq +version: 0.2.0.0 +synopsis: PostgreSQL job-queue (PGMQ) integration for Keiro +description: + A typed background-job queue for Keiro applications on top of PGMQ (the + PostgreSQL-native message queue) and shibuya (a Broadway-style worker + framework). Declare a @Job@ value bundling a queue, a payload codec, and a + retry/dead-letter policy, then write a plain domain handler of type + @p -> Eff es JobOutcome@. + +license: BSD-3-Clause +author: Nadeem Bitar +maintainer: nadeem@gmail.com +copyright: 2026 Nadeem Bitar +category: Control +homepage: https://github.com/shinzui/keiro#readme +bug-reports: https://github.com/shinzui/keiro/issues +build-type: Simple +tested-with: GHC >=9.12 && <9.13 +extra-doc-files: CHANGELOG.md + +source-repository head + type: git + location: https://github.com/shinzui/keiro.git + +common warnings + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-record-updates + -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints + +common shared + default-language: GHC2024 + default-extensions: + DeriveAnyClass + DuplicateRecordFields + MultilineStrings + OverloadedLabels + OverloadedRecordDot + OverloadedStrings + PackageImports + +library + import: warnings, shared + exposed-modules: + Keiro.PGMQ + Keiro.PGMQ.Codec + Keiro.PGMQ.Dlq + Keiro.PGMQ.Job + Keiro.PGMQ.Metrics + Keiro.PGMQ.Runtime + + hs-source-dirs: src + build-depends: + , aeson >=2.2 + , base >=4.21 && <5 + , effectful-core >=2.6 + , hasql >=1.10 + , hasql-pool >=1.2 + , hs-opentelemetry-api >=1.0 && <1.1 + , keiro-core ^>=0.2.0.0 + , pgmq-config >=0.3 && <0.4 + , pgmq-core >=0.3 && <0.4 + , pgmq-effectful >=0.3 && <0.4 + , pgmq-hasql >=0.3 && <0.4 + , shibuya-core >=0.8.0.1 && <0.9 + , shibuya-pgmq-adapter >=0.11 && <0.12 + , streamly-core >=0.3 + , text >=2.1 + , time >=1.12 + +test-suite keiro-pgmq-test + import: warnings, shared + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + , aeson >=2.2 + , base >=4.21 && <5 + , effectful-core >=2.6 + , hasql >=1.10 + , hasql-pool >=1.2 + , hs-opentelemetry-api >=1.0 && <1.1 + , hs-opentelemetry-propagator-w3c >=1.0 && <1.1 + , hs-opentelemetry-sdk >=1.0 && <1.1 + , hspec >=2.11 + , keiro-core ^>=0.2.0.0 + , keiro-pgmq + , keiro-test-support + , pgmq-config >=0.3 && <0.4 + , pgmq-core >=0.3 && <0.4 + , pgmq-effectful >=0.3 && <0.4 + , pgmq-migration >=0.3 && <0.4 + , shibuya-core >=0.8.0.1 && <0.9 + , shibuya-pgmq-adapter >=0.11 && <0.12 + , text >=2.1 diff --git a/src/Keiro/PGMQ.hs b/src/Keiro/PGMQ.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ.hs @@ -0,0 +1,22 @@ +{- | Umbrella module for @keiro-pgmq@: a typed background-job queue for Keiro +applications on top of PGMQ and shibuya. + +Re-exports the whole public surface — the transport-agnostic runtime +('Keiro.PGMQ.Runtime'), the payload codecs ('Keiro.PGMQ.Codec'), and the typed +'Keiro.PGMQ.Job.Job' ergonomics ('Keiro.PGMQ.Job'), plus DLQ operations +('Keiro.PGMQ.Dlq') and the typed queue-metrics surface ('Keiro.PGMQ.Metrics'). +Import this one module to get everything. +-} +module Keiro.PGMQ ( + module Keiro.PGMQ.Runtime, + module Keiro.PGMQ.Codec, + module Keiro.PGMQ.Job, + module Keiro.PGMQ.Dlq, + module Keiro.PGMQ.Metrics, +) where + +import Keiro.PGMQ.Codec +import Keiro.PGMQ.Dlq +import Keiro.PGMQ.Job +import Keiro.PGMQ.Metrics +import Keiro.PGMQ.Runtime diff --git a/src/Keiro/PGMQ/Codec.hs b/src/Keiro/PGMQ/Codec.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ/Codec.hs @@ -0,0 +1,130 @@ +{- | How a job payload is turned into PGMQ JSON and back. + +PGMQ stores message bodies as JSON ('Data.Aeson.Value'). A 'JobCodec' is the +adapter between a domain payload type @p@ and that JSON. Two are provided: + + * 'aesonJobCodec' - raw @aeson@ encode/decode. A drop-in match for apps that + already use @ToJSON@/@FromJSON@ payloads, so migrating to @keiro-pgmq@ is + mechanical. + * 'keiroJobCodec' - the versioned upgrade. It bridges keiro's + 'Keiro.Codec.Codec' by wrapping payloads in a @{ "v": , "t": + , "data": }@ envelope and replaying the codec's + upcaster chain on decode, giving job payloads the same schema-evolution + story event streams have. + +When raising a 'keiroJobCodec' schema version, deploy upgraded workers before +upgraded producers. A worker that sees an envelope from a future schema version +returns 'JobPayloadFromFuture', and the job runner retries it after the queue's +default retry delay so a rolling deploy can complete. Those retries still +consume delivery attempts, so size @maxRetries * defaultRetryDelay@ to cover +the deploy window. + +Do not switch a non-empty queue directly from 'aesonJobCodec' to +'keiroJobCodec': the wire shape changes from the bare payload to the +@{"v","t","data"}@ envelope. Drain the queue first or use a transitional codec that +accepts both shapes; otherwise old in-flight messages are malformed and will be +dead-lettered. +-} +module Keiro.PGMQ.Codec ( + JobCodec (..), + JobDecodeError (..), + mkJobCodec, + aesonJobCodec, + keiroJobCodec, +) where + +import "aeson" Data.Aeson (FromJSON, ToJSON, Value, object, parseJSON, toJSON, withObject, (.:), (.:?), (.=)) +import "aeson" Data.Aeson.Types (parseEither) +import "base" Data.Bifunctor (first) +import "base" Data.List.NonEmpty qualified as NonEmpty +import "keiro-core" Keiro.Codec (Codec, EventType (..)) +import "keiro-core" Keiro.Codec qualified as Codec +import "text" Data.Text (Text) +import "text" Data.Text qualified as Text + +-- | Why a job payload could not be decoded. +data JobDecodeError + = -- | The payload is malformed for this codec: poison, dead-letter it. + JobPayloadMalformed !Text + | {- | The payload was written by a newer schema version than this worker knows. + Carries payload version, then this codec's version. This is transient + during a rolling deploy: retry it, do not dead-letter it. + -} + JobPayloadFromFuture !Int !Int + deriving stock (Eq, Show) + +-- | How a job payload is turned into PGMQ JSON and back. +data JobCodec p = JobCodec + { encodeJob :: p -> Value + , decodeJob :: Value -> Either JobDecodeError p + } + +-- | Build a 'JobCodec' from an encoder and a legacy text-returning decoder. +mkJobCodec :: (p -> Value) -> (Value -> Either Text p) -> JobCodec p +mkJobCodec encode decode = + JobCodec + { encodeJob = encode + , decodeJob = first JobPayloadMalformed . decode + } + +{- | The default codec: raw @aeson@. Use this when the payload type already has +@ToJSON@/@FromJSON@ instances and you do not need versioned schema evolution. +-} +aesonJobCodec :: (ToJSON p, FromJSON p) => JobCodec p +aesonJobCodec = mkJobCodec toJSON (first Text.pack . parseEither parseJSON) + +{- | Versioned bridge to keiro's 'Keiro.Codec.Codec'. On encode it produces +@{ "v": , "t": , "data": }@; on +decode it reads the version and event-type tag, runs the codec's upcaster chain +up to the current version via 'Keiro.Codec.migrateToCurrent', then runs the +codec's current decoder. Legacy envelopes without @"t"@ still decode for +single-event codecs. Codec migration faults are surfaced as +'JobPayloadMalformed' (via 'show' of the 'Keiro.Codec.CodecError'), except for a +future envelope version, which surfaces as 'JobPayloadFromFuture'. +-} +keiroJobCodec :: Codec p -> JobCodec p +keiroJobCodec codec = + JobCodec + { encodeJob = \p -> + let EventType tag = Codec.eventType codec p + in object + [ "v" .= Codec.schemaVersion codec + , "t" .= tag + , "data" .= Codec.encode codec p + ] + , decodeJob = \value -> do + (version, tag, dataValue) <- parseEnvelope value + selectedType <- resolveEventType codec tag + let currentVersion = Codec.schemaVersion codec + if version > currentVersion + then Left (JobPayloadFromFuture version currentVersion) + else pure () + migrated <- + first (JobPayloadMalformed . Text.pack . show) $ + Codec.migrateToCurrent codec selectedType version dataValue + first JobPayloadMalformed (Codec.decode codec selectedType migrated) + } + +-- | Parse the @{ "v", "t", "data" }@ envelope, surfacing aeson failures as malformed payloads. +parseEnvelope :: Value -> Either JobDecodeError (Int, Maybe EventType, Value) +parseEnvelope = + first (JobPayloadMalformed . Text.pack) . parseEither parser + where + parser = withObject "Keiro.PGMQ.Codec envelope" $ \o -> do + version <- o .: "v" + tag <- fmap EventType <$> o .:? "t" + dataValue <- o .: "data" + pure (version, tag, dataValue) + +resolveEventType :: Codec p -> Maybe EventType -> Either JobDecodeError EventType +resolveEventType _ (Just tag) = Right tag +resolveEventType codec Nothing = + case NonEmpty.toList (Codec.eventTypes codec) of + [tag] -> Right tag + tags -> + Left + . JobPayloadMalformed + $ "missing event-type tag in multi-event keiro job envelope; expected one of " + <> Text.intercalate ", " (fmap eventTypeText tags) + where + eventTypeText (EventType tag) = tag diff --git a/src/Keiro/PGMQ/Dlq.hs b/src/Keiro/PGMQ/Dlq.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ/Dlq.hs @@ -0,0 +1,244 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE LambdaCase #-} +{-# OPTIONS_GHC -Wno-redundant-constraints #-} + +{- | Dead-letter queue inspection and redrive helpers. + +The PGMQ adapter writes DLQ rows as a shibuya wrapper whose required keys are +@original_message@ and @dead_letter_reason@. Keiro's worker path always includes +metadata as well: @original_message_id@, @original_enqueued_at@, @last_read_at@, +@read_count@, and @original_headers@. These helpers parse the required keys and +treat metadata as optional so operators can still inspect legacy or hand-written +rows. + +PGMQ does not expire DLQ rows by itself. The retention model is "archive-then- +purge": 'archiveDlq' retains dead letters by moving them into the archive table +@pgmq.a_@ (preserving @enqueued_at@ / @read_ct@ and stamping @archived_at@) +for audit, while 'purgeDlq' deletes them permanently. An operator who needs an +audit trail runs 'archiveDlq' (retain) and may then 'purgeDlq' (clear the active +table); an operator who does not keeps using 'purgeDlq' alone. Either way, alert +on the DLQ's depth via 'Keiro.PGMQ.Metrics.jobDlqMetrics'. + +Redrive is at-least-once: a crash after sending the original payload back to the +main queue but before deleting the DLQ row leaves the payload in both places. +Handlers must therefore be idempotent. +-} +module Keiro.PGMQ.Dlq ( + DlqEntry (..), + readDlq, + redriveDlq, + purgeDlq, + archiveDlq, + archiveDlqEntry, +) where + +import Keiro.PGMQ.Codec (JobDecodeError (..), decodeJob) +import Keiro.PGMQ.Job (Job (..)) +import Keiro.PGMQ.Runtime (QueueRef (..)) +import "aeson" Data.Aeson (Value (..), withObject, (.:), (.:?)) +import "aeson" Data.Aeson.Types (Parser, parseEither) +import "base" Control.Monad (foldM, void) +import "base" Data.Foldable (toList) +import "base" Data.Int (Int32, Int64) +import "effectful-core" Effectful (Eff, IOE, (:>)) +import "pgmq-effectful" Pgmq.Effectful ( + Message (..), + MessageBody (..), + MessageId, + MessageQuery (..), + Pgmq, + ReadMessage (..), + SendMessage (..), + ) +import "pgmq-effectful" Pgmq.Effectful qualified as Pgmq +import "text" Data.Text (Text) +import "text" Data.Text qualified as Text +import "time" Data.Time (UTCTime) + +-- | A decoded dead-letter entry: shibuya's DLQ wrapper, unwrapped. +data DlqEntry p = DlqEntry + { dlqMessageId :: !MessageId + -- ^ The DLQ row's own PGMQ message id. + , reason :: !Text + -- ^ @poison_pill: ...@, @invalid_payload: ...@, or @max_retries_exceeded@. + , originalPayload :: !(Either JobDecodeError p) + -- ^ The preserved @original_message@ decoded with the job's codec. + , originalMessageId :: !(Maybe Int64) + , originalEnqueuedAt :: !(Maybe UTCTime) + , readCount :: !(Maybe Int64) + , rawBody :: !Value + -- ^ Full DLQ wrapper for forensics. + } + deriving stock (Show) + +data DlqEnvelope = DlqEnvelope + { originalMessage :: !Value + , deadLetterReason :: !Text + , envelopeOriginalMessageId :: !(Maybe Int64) + , envelopeOriginalEnqueuedAt :: !(Maybe UTCTime) + , envelopeReadCount :: !(Maybe Int64) + } + +parseDlqEnvelope :: Value -> Either Text DlqEnvelope +parseDlqEnvelope = + firstText . parseEither parser + where + parser :: Value -> Parser DlqEnvelope + parser = + withObject "DLQ payload" \obj -> do + originalMessage <- obj .: "original_message" + deadLetterReason <- obj .: "dead_letter_reason" + envelopeOriginalMessageId <- obj .:? "original_message_id" + envelopeOriginalEnqueuedAt <- obj .:? "original_enqueued_at" + envelopeReadCount <- obj .:? "read_count" + pure + DlqEnvelope + { originalMessage + , deadLetterReason + , envelopeOriginalMessageId + , envelopeOriginalEnqueuedAt + , envelopeReadCount + } + + firstText = \case + Left err -> Left (Text.pack err) + Right value -> Right value + +{- | Read and decode up to @n@ DLQ entries. The read uses a 30 second visibility +timeout so concurrent inspections do not immediately see the same rows. +-} +readDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int32 -> Eff es [DlqEntry p] +readDlq job n + | n <= 0 = pure [] + | otherwise = do + messages <- + Pgmq.readMessage + ReadMessage + { queueName = job.jobQueue.dlqName + , delay = 30 + , batchSize = Just n + , conditional = Nothing + } + pure (fmap (toEntry job) (toList messages)) + +toEntry :: Job p -> Message -> DlqEntry p +toEntry job message = + let body = unMessageBody message.body + in case parseDlqEnvelope body of + Left err -> + DlqEntry + { dlqMessageId = message.messageId + , reason = "malformed_dlq_payload: " <> err + , originalPayload = Left (JobPayloadMalformed err) + , originalMessageId = Nothing + , originalEnqueuedAt = Nothing + , readCount = Nothing + , rawBody = body + } + Right envelope -> + DlqEntry + { dlqMessageId = message.messageId + , reason = envelope.deadLetterReason + , originalPayload = decodeJob job.jobCodec envelope.originalMessage + , originalMessageId = envelope.envelopeOriginalMessageId + , originalEnqueuedAt = envelope.envelopeOriginalEnqueuedAt + , readCount = envelope.envelopeReadCount + , rawBody = body + } + +{- | Move up to @n@ DLQ rows back to the main queue. Redriven messages start a +fresh PGMQ @read_ct@ on the main queue. Malformed DLQ wrappers are left in the +DLQ for inspection. +-} +redriveDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int -> Eff es Int +redriveDlq job n + | n <= 0 = pure 0 + | otherwise = loop 0 + where + loop moved + | moved >= n = pure moved + | otherwise = do + messages <- + Pgmq.readMessage + ReadMessage + { queueName = job.jobQueue.dlqName + , delay = 30 + , batchSize = Just (fromIntegral (min 100 (n - moved))) + , conditional = Nothing + } + if null messages + then pure moved + else do + movedInBatch <- foldM redriveOne 0 messages + if movedInBatch == 0 + then pure moved + else loop (moved + movedInBatch) + + redriveOne count message = + case parseDlqEnvelope (unMessageBody message.body) of + Left _err -> + pure count + Right envelope -> do + _ <- + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody envelope.originalMessage + , delay = Nothing + } + void $ + Pgmq.deleteMessage + MessageQuery + { queueName = job.jobQueue.dlqName + , messageId = message.messageId + } + pure (count + 1) + +-- | Delete all rows currently in the DLQ. +purgeDlq :: (Pgmq :> es, IOE :> es) => Job p -> Eff es () +purgeDlq job = + void (Pgmq.deleteAllMessagesFromQueue job.jobQueue.dlqName) + +{- | Archive (retain) up to @n@ DLQ rows: move each out of the active DLQ table +@pgmq.q_@ into the archive table @pgmq.a_@, preserving @enqueued_at@ / +@read_ct@ and stamping @archived_at@. Returns the number archived. This is the +audit-retention counterpart to the delete-only 'purgeDlq'. At-most-once per row +per call; a crash before archiving leaves the row in the active DLQ for a re-run. +-} +archiveDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int -> Eff es Int +archiveDlq job n + | n <= 0 = pure 0 + | otherwise = loop 0 + where + loop archived + | archived >= n = pure archived + | otherwise = do + messages <- + Pgmq.readMessage + ReadMessage + { queueName = job.jobQueue.dlqName + , delay = 30 + , batchSize = Just (fromIntegral (min 100 (n - archived))) + , conditional = Nothing + } + if null messages + then pure archived + else do + archivedInBatch <- foldM archiveOne 0 messages + if archivedInBatch == 0 + then pure archived + else loop (archived + archivedInBatch) + + archiveOne count message = do + moved <- archiveDlqEntry job message.messageId + pure (if moved then count + 1 else count) + +-- | Archive one specific DLQ row by message id. 'True' if a row was moved. +archiveDlqEntry :: (Pgmq :> es, IOE :> es) => Job p -> MessageId -> Eff es Bool +archiveDlqEntry job msgId = + Pgmq.archiveMessage + MessageQuery + { queueName = job.jobQueue.dlqName + , messageId = msgId + } diff --git a/src/Keiro/PGMQ/Job.hs b/src/Keiro/PGMQ/Job.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ/Job.hs @@ -0,0 +1,920 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeApplications #-} +-- @enqueue@/@enqueueWithDelay@ carry an @IOE :> es@ constraint that the @Pgmq@ +-- send operation does not strictly require. It is kept deliberately: it is part +-- of the published @keiro-pgmq@ contract (mirrored in the MasterPlan's +-- Integration Points and depended on by the two consumer migrations), and it +-- keeps the producer signatures uniform with the processor/runner ones. We +-- therefore silence the otherwise-correct redundant-constraint warning here. +{-# OPTIONS_GHC -Wno-redundant-constraints #-} + +{- | Layer 2 of @keiro-pgmq@: the typed-'Job' ergonomics built on top of +'Keiro.PGMQ.Runtime'. This is the payoff layer that absorbs the boilerplate two +real apps wrote by hand. + +An application declares a 'Job' value bundling a queue ('Keiro.PGMQ.Runtime.QueueRef'), +a payload codec ('Keiro.PGMQ.Codec.JobCodec'), and a 'RetryPolicy'; then writes +a plain domain handler of type @p -> Eff es 'JobOutcome'@ that never touches +shibuya's @Ingested@/@AckDecision@ or PGMQ's wire types. The package provides: + + * 'enqueue' / 'enqueueWithDelay' — producers. + * 'ensureJobQueue' — idempotent main-queue + DLQ creation. + * 'jobProcessor' — build a shibuya processor from a 'Job' plus a handler. + * 'runJobWorkers' — continuous, multi-processor supervised run (the @rei@ cadence). + * 'runJobOnce' — one-shot drain of up to @n@ messages (the @hospital-capacity@ cadence). + +== Delivery and crash semantics + +Delivery is at-least-once. A handler must be idempotent because the same message +can be delivered again after a worker crash, a handler exception, or a visibility +timeout expiry. Crash redelivery cadence is the active visibility timeout, not +the 'RetryPolicy' delay; the policy delay applies only to explicit 'Retry' and +'RetryDefault' outcomes. Every visibility-timeout expiry consumes one PGMQ +@read_ct@ attempt, and messages whose read count exceeds 'maxRetries' are +dead-lettered before the handler sees them. + +Dead-lettering sends a DLQ row and then deletes the main-queue row, so a crash +between those two statements can leave the message in both places. 'redriveDlq' +has the same at-least-once window in the other direction: it sends the preserved +payload back to the main queue and then deletes the DLQ row. + +Transient database errors during PGMQ polling are retried by the adapter, and a +polling failure that exhausts that retry policy is propagated visibly through +shibuya supervision rather than completing the worker silently. +-} +module Keiro.PGMQ.Job ( + -- * Job declaration + JobOutcome (..), + RetryDelay (..), + RetryPolicy (..), + RetryPolicyConfigError (..), + mkRetryPolicy, + defaultRetryPolicy, + Job (..), + JobPolling (..), + JobOrdering (..), + JobTuning (..), + JobTuningConfigError (..), + mkJobTuning, + defaultJobTuning, + withOrdering, + + -- * Message metadata + MessageHeaders (..), + + -- * Producing work + enqueue, + enqueueWithDelay, + enqueueWithHeaders, + enqueueWithHeadersAndDelay, + enqueueBatch, + enqueueBatchWithDelay, + enqueueBatchWithHeaders, + enqueueTraced, + enqueueTracedWithDelay, + enqueueToGroup, + enqueueToGroupWithDelay, + + -- * Queue lifecycle + QueueKind (..), + PartitionSpec (..), + QueueProvision (..), + standardProvision, + unloggedProvision, + partitionedProvision, + withFifoIndexProvision, + queueProvisionConfigs, + ensureJobQueue, + ensureJobQueueWith, + ensureFifoIndex, + ensureOrderedJobQueue, + + -- * Consuming work + JobContext (..), + jobProcessorWithContext, + jobProcessor, + runJobWorkers, + runJobOnceWithContext, + runJobOnce, +) where + +import Keiro.PGMQ.Codec (JobCodec, JobDecodeError (..), decodeJob, encodeJob) +import Keiro.PGMQ.Runtime (QueueRef (..)) +import "aeson" Data.Aeson (Value, object, (.=)) +import "base" Control.Exception (Exception, SomeException, throwIO) +import "base" Control.Monad (foldM, void) +import "base" Data.Int (Int32, Int64) +import "effectful-core" Effectful (Eff, IOE, liftIO, (:>)) +import "effectful-core" Effectful.Error.Static (Error) +import "effectful-core" Effectful.Exception qualified as EffException +import "effectful-core" Effectful.Reader.Static (Reader, ask) +import "hs-opentelemetry-api" OpenTelemetry.Context.ThreadLocal (getContext) +import "hs-opentelemetry-api" OpenTelemetry.Trace.Core (TracerProvider) +import "pgmq-config" Pgmq.Config.Effectful (ensureQueuesEff) +import "pgmq-config" Pgmq.Config.Types qualified as Config +import "pgmq-effectful" Pgmq.Effectful ( + BatchSendMessage (..), + BatchSendMessageWithHeaders (..), + Message (..), + MessageBody (..), + MessageHeaders (..), + MessageId, + MessageQuery (..), + Pgmq, + PgmqRuntimeError, + ReadMessage (..), + SendMessage (..), + SendMessageWithHeaders (..), + VisibilityTimeoutQuery (..), + injectTraceContext, + mergeTraceHeaders, + ) +import "pgmq-effectful" Pgmq.Effectful qualified as Pgmq +import "pgmq-effectful" Pgmq.Effectful.Effect (readGrouped, readGroupedRoundRobin) +import "pgmq-hasql" Pgmq.Hasql.Statements.Types (ReadGrouped (..)) +import "shibuya-core" Shibuya.App ( + AppConfig (..), + AppError, + AppHandle, + ProcessorId (..), + QueueProcessor, + SupervisionStrategy, + mkProcessor, + runApp, + ) +import "shibuya-core" Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), RetryDelay (..)) +import "shibuya-core" Shibuya.Core.Ingested qualified as Shibuya +import "shibuya-core" Shibuya.Core.Lease (Lease (..)) +import "shibuya-core" Shibuya.Core.Types (Attempt (..), Envelope (..)) +import "shibuya-core" Shibuya.Telemetry.Effect (Tracing) +import "shibuya-pgmq-adapter" Shibuya.Adapter.Pgmq ( + FifoConfig (..), + FifoReadStrategy (..), + PgmqAdapterConfig (..), + PgmqAdapterEnv, + PgmqConfigError, + PollingConfig (..), + defaultConfig, + directDeadLetter, + pgmqAdapter, + ) +import "shibuya-pgmq-adapter" Shibuya.Adapter.Pgmq.Convert ( + mkDlqPayload, + pgmqMessageToEnvelope, + ) +import "text" Data.Text (Text) +import "time" Data.Time (NominalDiffTime, nominalDiffTimeToSeconds) + +-- | What a job handler decides. Never exposes shibuya/PGMQ wire types to the caller. +data JobOutcome + = -- | Processed successfully; delete the message from the queue. + Done + | -- | Leave the message on the queue; redeliver after the delay. + Retry !RetryDelay + | -- | Leave the message on the queue; redeliver after the policy's default retry delay. + RetryDefault + | -- | Poison message; route to the dead-letter queue when enabled, otherwise archive it, with this reason. + Dead !Text + deriving stock (Show) + +{- | How a queue retries and dead-letters. + +The raw constructor is exported for advanced/manual configuration, but it is +not validated. Prefer 'mkRetryPolicy': @maxRetries <= 0@ dead-letters every +message before the handler runs because PGMQ's @read_ct@ is 1 on first delivery +and the adapter auto-dead-letters when @read_ct > maxRetries@. Negative retry +delays can create immediate redelivery storms. + +'maxRetries' is the number of +deliveries PGMQ allows before auto-dead-lettering; 'defaultRetryDelay' is a +convenience default a handler can reach for; 'useDeadLetter' decides whether a +DLQ is created and routed to at all. +-} +data RetryPolicy = RetryPolicy + { maxRetries :: !Int64 + , defaultRetryDelay :: !RetryDelay + , useDeadLetter :: !Bool + } + deriving stock (Eq, Show) + +data RetryPolicyConfigError + = NonPositiveMaxRetries !Int64 + | NegativeRetryDelay !RetryDelay + deriving stock (Eq, Show) + +mkRetryPolicy :: Int64 -> RetryDelay -> Bool -> Either RetryPolicyConfigError RetryPolicy +mkRetryPolicy maxRetries defaultRetryDelay useDeadLetter + | maxRetries < 1 = Left (NonPositiveMaxRetries maxRetries) + | retryDelaySeconds defaultRetryDelay < 0 = Left (NegativeRetryDelay defaultRetryDelay) + | otherwise = + Right + RetryPolicy + { maxRetries + , defaultRetryDelay + , useDeadLetter + } + +-- | Five deliveries, a 60-second default retry delay, and a DLQ enabled. +defaultRetryPolicy :: RetryPolicy +defaultRetryPolicy = + RetryPolicy + { maxRetries = 5 + , defaultRetryDelay = RetryDelay 60 + , useDeadLetter = True + } + +data JobPolling + = -- | Sleep this long between empty polls. + PollEvery !NominalDiffTime + | -- | Long-poll inside the database: max seconds to wait, then check interval in milliseconds. + LongPoll !Int32 !Int32 + deriving stock (Eq, Show) + +{- | How a consumer orders deliveries. + +'Unordered' is the historical behavior: PGMQ's plain @read@, FIFO only in +selection order (@msg_id@ ascending), with NO per-key delivery-order guarantee +under concurrent workers, retries, or visibility-timeout expiry. + +'FifoThroughput' and 'FifoRoundRobin' enable strict per-group ordering via PGMQ +message groups (the reserved @x-pgmq-group@ header). Within one group, messages +are delivered in strict send order; distinct groups proceed in parallel. +'FifoThroughput' fills a batch from the oldest eligible group first (SQS-style, +@read_grouped@); 'FifoRoundRobin' interleaves fairly across groups +(@read_grouped_rr@). Delivery is still at-least-once and there is no +deduplication, so handlers must be idempotent. +-} +data JobOrdering + = Unordered + | FifoThroughput + | FifoRoundRobin + deriving stock (Eq, Show) + +{- | How a consumer reads the queue. + +The raw constructor is exported but not validated. Prefer 'mkJobTuning' so +visibility timeouts, batch sizes, and polling intervals are positive. +-} +data JobTuning = JobTuning + { visibilityTimeout :: !Int32 + , batchSize :: !Int32 + , polling :: !JobPolling + , ordering :: !JobOrdering + } + deriving stock (Eq, Show) + +-- | 30 s visibility timeout, batch of 1, 1 s standard polling, unordered reads. +defaultJobTuning :: JobTuning +defaultJobTuning = + JobTuning + { visibilityTimeout = 30 + , batchSize = 1 + , polling = PollEvery 1 + , ordering = Unordered + } + +data JobTuningConfigError + = NonPositiveVisibilityTimeout !Int32 + | NonPositiveBatchSize !Int32 + | NonPositivePollInterval + deriving stock (Eq, Show) + +mkJobTuning :: Int32 -> Int32 -> JobPolling -> Either JobTuningConfigError JobTuning +mkJobTuning visibilityTimeout batchSize polling + | visibilityTimeout < 1 = Left (NonPositiveVisibilityTimeout visibilityTimeout) + | batchSize < 1 = Left (NonPositiveBatchSize batchSize) + | not (validPolling polling) = Left NonPositivePollInterval + | otherwise = Right JobTuning{visibilityTimeout, batchSize, polling, ordering = Unordered} + +{- | Set the FIFO read strategy on an existing tuning, e.g. +@withOrdering FifoThroughput defaultJobTuning@. Every 'JobOrdering' value is +valid, so this is a plain record update rather than a validating constructor. +-} +withOrdering :: JobOrdering -> JobTuning -> JobTuning +withOrdering o tuning = tuning{ordering = o} + +validPolling :: JobPolling -> Bool +validPolling (PollEvery interval) = interval > 0 +validPolling (LongPoll maxPollSeconds pollIntervalMs) = + maxPollSeconds > 0 && pollIntervalMs > 0 + +toPollingConfig :: JobPolling -> PollingConfig +toPollingConfig (PollEvery interval) = StandardPolling interval +toPollingConfig (LongPoll maxPollSeconds pollIntervalMs) = LongPolling maxPollSeconds pollIntervalMs + +-- | Map an ordering choice to the shibuya adapter's FIFO read config (worker path). +toFifoConfig :: JobOrdering -> Maybe FifoConfig +toFifoConfig Unordered = Nothing +toFifoConfig FifoThroughput = Just (FifoConfig ThroughputOptimized) +toFifoConfig FifoRoundRobin = Just (FifoConfig RoundRobin) + +retryDelaySeconds :: RetryDelay -> NominalDiffTime +retryDelaySeconds (RetryDelay seconds) = seconds + +nominalToSeconds :: NominalDiffTime -> Int32 +nominalToSeconds dt = + let seconds :: Double + seconds = realToFrac (nominalDiffTimeToSeconds dt) + maxSec :: Double + maxSec = fromIntegral (maxBound :: Int32) + minSec :: Double + minSec = fromIntegral (minBound :: Int32) + clamped = max minSec (min maxSec seconds) + in ceiling clamped + +{- | A declarative job: a queue, a payload codec, and a retry policy, named for +telemetry. Construct one and pair it with a handler of type +@p -> Eff es 'JobOutcome'@. +-} +data Job p = Job + { jobName :: !Text + -- ^ Used as the shibuya 'ProcessorId' and telemetry label. + , jobQueue :: !QueueRef + , jobCodec :: !(JobCodec p) + , jobPolicy :: !RetryPolicy + } + +-- | Per-delivery capabilities handed to context-aware handlers. +data JobContext es = JobContext + { extendLease :: !(NominalDiffTime -> Eff es ()) + -- ^ Push the message's visibility timeout further into the future. + , attempt :: !(Maybe Word) + -- ^ Zero-based delivery attempt; @Just 0@ is the first delivery. + , headers :: !(Maybe Value) + {- ^ Drain path: the raw PGMQ message header object (@Just@ when the + message carried headers, @Nothing@ otherwise). Worker path: always + @Nothing@, because the shibuya adapter's @Envelope@ does not surface + arbitrary headers (only the trace context, which shibuya itself uses + to continue the trace). + -} + } + +-- | Producer: encode @p@ with the job's codec and send it to the queue, no delay. +enqueue :: (Pgmq :> es, IOE :> es) => Job p -> p -> Eff es MessageId +enqueue job p = + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody (encodeJob job.jobCodec p) + , delay = Nothing + } + +{- | Producer with an explicit visibility delay (in seconds, PGMQ's @Delay@ is +@Int32@) before first delivery. +-} +enqueueWithDelay :: (Pgmq :> es, IOE :> es) => Job p -> Int32 -> p -> Eff es MessageId +enqueueWithDelay job d p = + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody (encodeJob job.jobCodec p) + , delay = Just d + } + +{- | Producer that attaches caller-supplied message headers (an arbitrary JSON +object) alongside the encoded payload. Headers ride in PGMQ's @headers@ column +and are readable by the consumer (see 'JobContext'\'s @headers@ field on the +drain path). + +The headers are passed through verbatim. In particular the reserved FIFO group +key @x-pgmq-group@ is neither reserved, injected, stripped, nor rewritten, so a +caller (or a sibling plan building ordered delivery) may set it freely. +-} +enqueueWithHeaders :: + (Pgmq :> es, IOE :> es) => Job p -> MessageHeaders -> p -> Eff es MessageId +enqueueWithHeaders job hdrs p = + Pgmq.sendMessageWithHeaders + SendMessageWithHeaders + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody (encodeJob job.jobCodec p) + , messageHeaders = hdrs + , delay = Nothing + } + +{- | 'enqueueWithHeaders' with an explicit visibility delay (in seconds) before +first delivery. +-} +enqueueWithHeadersAndDelay :: + (Pgmq :> es, IOE :> es) => Job p -> Int32 -> MessageHeaders -> p -> Eff es MessageId +enqueueWithHeadersAndDelay job d hdrs p = + Pgmq.sendMessageWithHeaders + SendMessageWithHeaders + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody (encodeJob job.jobCodec p) + , messageHeaders = hdrs + , delay = Just d + } + +{- | Batch producer: encode and enqueue many payloads in a single database +round-trip, returning one 'MessageId' per payload in order. An empty input +short-circuits to @[]@ and issues no statement. +-} +enqueueBatch :: (Pgmq :> es, IOE :> es) => Job p -> [p] -> Eff es [MessageId] +enqueueBatch _ [] = pure [] +enqueueBatch job ps = + Pgmq.batchSendMessage + BatchSendMessage + { queueName = job.jobQueue.physicalName + , messageBodies = map (MessageBody . encodeJob job.jobCodec) ps + , delay = Nothing + } + +{- | 'enqueueBatch' with a single visibility delay (in seconds) applied to every +message in the batch. +-} +enqueueBatchWithDelay :: + (Pgmq :> es, IOE :> es) => Job p -> Int32 -> [p] -> Eff es [MessageId] +enqueueBatchWithDelay _ _ [] = pure [] +enqueueBatchWithDelay job d ps = + Pgmq.batchSendMessage + BatchSendMessage + { queueName = job.jobQueue.physicalName + , messageBodies = map (MessageBody . encodeJob job.jobCodec) ps + , delay = Just d + } + +{- | Batch producer that attaches a distinct header object to each payload. The +input pairs each payload with its headers so the body and header lists cannot be +desynchronized. An empty input short-circuits to @[]@. +-} +enqueueBatchWithHeaders :: + (Pgmq :> es, IOE :> es) => Job p -> [(MessageHeaders, p)] -> Eff es [MessageId] +enqueueBatchWithHeaders _ [] = pure [] +enqueueBatchWithHeaders job pairs = + Pgmq.batchSendMessageWithHeaders + BatchSendMessageWithHeaders + { queueName = job.jobQueue.physicalName + , messageBodies = map (MessageBody . encodeJob job.jobCodec . snd) pairs + , messageHeaders = map fst pairs + , delay = Nothing + } + +{- | Producer that propagates the current OpenTelemetry trace context onto the +enqueued message so the handler runs inside the same trace. The current +thread-local context is injected to carrier headers via the provider's +configured propagator (W3C @traceparent@ by default) and additively merged onto +@extraHeaders@ — any key already present in @extraHeaders@ wins, so a +caller-set @x-pgmq-group@ survives. Pass @MessageHeaders (object [])@ to inject +only the trace. +-} +enqueueTraced :: + (Pgmq :> es, IOE :> es) => + TracerProvider -> Job p -> MessageHeaders -> p -> Eff es MessageId +enqueueTraced provider job extraHeaders p = do + ctx <- liftIO getContext + traceHeaders <- injectTraceContext provider ctx + let merged = MessageHeaders (mergeTraceHeaders traceHeaders (Just extraHeaders.unMessageHeaders)) + enqueueWithHeaders job merged p + +{- | 'enqueueTraced' with an explicit visibility delay (in seconds) before first +delivery. +-} +enqueueTracedWithDelay :: + (Pgmq :> es, IOE :> es) => + TracerProvider -> Job p -> Int32 -> MessageHeaders -> p -> Eff es MessageId +enqueueTracedWithDelay provider job d extraHeaders p = do + ctx <- liftIO getContext + traceHeaders <- injectTraceContext provider ctx + let merged = MessageHeaders (mergeTraceHeaders traceHeaders (Just extraHeaders.unMessageHeaders)) + enqueueWithHeadersAndDelay job d merged p + +{- | Enqueue a payload into the FIFO group named by @groupKey@. The group key is +written under the reserved @x-pgmq-group@ JSONB header, which PGMQ's grouped +reads and the shibuya adapter use to order deliveries per group. Consume with an +ordered 'JobTuning' (see 'withOrdering') to honor the order; within one group, +messages are handled in strict send order while distinct groups proceed in +parallel. +-} +enqueueToGroup :: + (Pgmq :> es, IOE :> es) => Job p -> Text -> p -> Eff es MessageId +enqueueToGroup job groupKey p = + enqueueWithHeaders job (groupHeader groupKey) p + +-- | 'enqueueToGroup' with an explicit first-delivery delay (in seconds). +enqueueToGroupWithDelay :: + (Pgmq :> es, IOE :> es) => Job p -> Int32 -> Text -> p -> Eff es MessageId +enqueueToGroupWithDelay job d groupKey p = + enqueueWithHeadersAndDelay job d (groupHeader groupKey) p + +-- | The reserved FIFO group header for a group key. +groupHeader :: Text -> MessageHeaders +groupHeader k = MessageHeaders (object ["x-pgmq-group" .= k]) + +{- | The three PostgreSQL storage shapes a job's main queue can take. + + * 'StandardKind' — a normal write-ahead-logged queue table (today's default). + * 'UnloggedKind' — an /unlogged/ table: writes skip the WAL (faster) but the + table is truncated to empty on a database crash. For transient, regenerable + work. + * 'PartitionedKind' — storage split across child tables by time or message-id + range, managed by the PostgreSQL extension @pg_partman@. Requires a + @pg_partman@-enabled server (see 'partitionedProvision'). +-} +data QueueKind + = StandardKind + | UnloggedKind + | PartitionedKind !PartitionSpec + deriving stock (Eq, Show) + +{- | Partition interval + retention interval for a partitioned queue. Both are +PostgreSQL/@pg_partman@ duration or integer strings — e.g. @"daily"@ or +@"10000"@ for the interval, @"7 days"@ or @"100000"@ for the retention. +-} +data PartitionSpec = PartitionSpec + { partitionInterval :: !Text + , retentionInterval :: !Text + } + deriving stock (Eq, Show) + +{- | The provisioning choice for a job's /main/ queue: which storage shape, and +whether to create the FIFO GIN index. The DLQ (when the policy enables one) is +always a plain standard queue with no FIFO index. +-} +data QueueProvision = QueueProvision + { provisionKind :: !QueueKind + , provisionFifoIndex :: !Bool + } + deriving stock (Eq, Show) + +-- | A standard main queue with no FIFO index — exactly today's behavior. +standardProvision :: QueueProvision +standardProvision = QueueProvision{provisionKind = StandardKind, provisionFifoIndex = False} + +-- | An unlogged main queue with no FIFO index. +unloggedProvision :: QueueProvision +unloggedProvision = QueueProvision{provisionKind = UnloggedKind, provisionFifoIndex = False} + +-- | A partitioned main queue (no FIFO index) with the given interval/retention. +partitionedProvision :: PartitionSpec -> QueueProvision +partitionedProvision spec = + QueueProvision{provisionKind = PartitionedKind spec, provisionFifoIndex = False} + +-- | Turn on FIFO-index creation for a provisioning choice. +withFifoIndexProvision :: QueueProvision -> QueueProvision +withFifoIndexProvision provision = provision{provisionFifoIndex = True} + +{- | Pure: the list of @pgmq-config@ 'Config.QueueConfig's that +'ensureJobQueueWith' will reconcile — the main queue first (with its chosen kind +and optional FIFO index), then the DLQ (always a standard queue) when the policy +enables one. Exposed so the partitioned path is testable without a +@pg_partman@-enabled database. +-} +queueProvisionConfigs :: QueueProvision -> Job p -> [Config.QueueConfig] +queueProvisionConfigs provision job = + mainConfig : dlqConfigs + where + mainBase = + case provision.provisionKind of + StandardKind -> Config.standardQueue job.jobQueue.physicalName + UnloggedKind -> Config.unloggedQueue job.jobQueue.physicalName + PartitionedKind spec -> + Config.partitionedQueue + job.jobQueue.physicalName + Config.PartitionConfig + { Config.partitionInterval = spec.partitionInterval + , Config.retentionInterval = spec.retentionInterval + } + mainConfig + | provision.provisionFifoIndex = Config.withFifoIndex mainBase + | otherwise = mainBase + dlqConfigs + | job.jobPolicy.useDeadLetter = [Config.standardQueue job.jobQueue.dlqName] + | otherwise = [] + +{- | Idempotent: create the job's main queue with the chosen storage kind and +(optionally) its FIFO index, plus the DLQ (always a standard queue) when the +policy uses one. Routes through @pgmq-config@'s additive reconciler, which lists +existing queues first and only creates what is missing, so this is safe to call +at every worker startup. +-} +ensureJobQueueWith :: (Pgmq :> es) => QueueProvision -> Job p -> Eff es () +ensureJobQueueWith provision job = + ensureQueuesEff (queueProvisionConfigs provision job) + +{- | Idempotent: create the main queue, and the DLQ too when the policy uses +one. Unchanged behavior: @ensureJobQueueWith standardProvision@. Safe to call at +every worker startup. +-} +ensureJobQueue :: (Pgmq :> es) => Job p -> Eff es () +ensureJobQueue = ensureJobQueueWith standardProvision + +{- | Create the FIFO GIN index on the job's /main/ queue's @headers@ column — +the index PGMQ's grouped/ordered reads (@read_grouped@/@read_grouped_rr@) match +against. Idempotent: the index step is always re-applied and the underlying SQL +is @CREATE INDEX IF NOT EXISTS@, so a second call is a harmless no-op. Routing +through @pgmq-config@'s reconciler (which lists existing queues first) means +calling this on an already-provisioned queue does not recreate the queue. This +is the artifact the FIFO ordered-delivery plan +(@docs/plans/77-add-fifo-ordered-delivery-via-message-groups-to-keiro-pgmq.md@) +consumes for ordered jobs. +-} +ensureFifoIndex :: (Pgmq :> es) => Job p -> Eff es () +ensureFifoIndex job = + ensureQueuesEff + [Config.withFifoIndex (Config.standardQueue job.jobQueue.physicalName)] + +{- | Provision an ordered job's queue: create the main queue (and the DLQ when +the policy uses one) plus the FIFO GIN index that grouped reads need. Composes +'ensureJobQueue' and 'ensureFifoIndex'; both are idempotent, so this is safe to +call at every startup. +-} +ensureOrderedJobQueue :: (Pgmq :> es) => Job p -> Eff es () +ensureOrderedJobQueue job = do + ensureJobQueue job + ensureFifoIndex job + +{- | The PGMQ adapter rejected the config derived from a job's tuning. Job tuning +is validated at construction ('mkJobTuning') and 'adapterConfigFor' derives the +adapter config deterministically, so this indicates an internal inconsistency +rather than a recoverable condition; it is surfaced as an exception. +-} +newtype JobAdapterConfigInvalid = JobAdapterConfigInvalid PgmqConfigError + deriving stock (Show) + deriving anyclass (Exception) + +{- | Build the shibuya PGMQ adapter config from a job's queue and policy: route +to the DLQ via the adapter's @directDeadLetter@ path when the policy enables it. +-} +adapterConfigFor :: JobTuning -> Job p -> PgmqAdapterConfig +adapterConfigFor tuning job = + (defaultConfig job.jobQueue.physicalName) + { visibilityTimeout = tuning.visibilityTimeout + , batchSize = tuning.batchSize + , polling = toPollingConfig tuning.polling + , fifoConfig = toFifoConfig tuning.ordering + , maxRetries = job.jobPolicy.maxRetries + , deadLetterConfig = + if job.jobPolicy.useDeadLetter + then Just (directDeadLetter job.jobQueue.dlqName True) + else Nothing + } + +{- | The boilerplate this package absorbs once: decode the raw JSON payload with +the job's codec, run the domain handler, and translate its 'JobOutcome' into a +shibuya 'AckDecision'. A payload the codec rejects is dead-lettered. +-} +wrapHandler :: + Job p -> + (JobContext es -> p -> Eff es JobOutcome) -> + (Shibuya.Message es Value -> Eff es AckDecision) +wrapHandler job handle ingested = + case decodeJob job.jobCodec ingested.envelope.payload of + Left (JobPayloadFromFuture _payloadVersion _workerVersion) -> + pure (AckRetry job.jobPolicy.defaultRetryDelay) + Left (JobPayloadMalformed err) -> + pure (AckDeadLetter (InvalidPayload err)) + Right p -> toAck <$> handle (contextFor ingested) p + where + contextFor message = + JobContext + { extendLease = maybe (\_ -> pure ()) (.leaseExtend) message.lease + , attempt = fmap (.unAttempt) message.envelope.attempt + , headers = Nothing + } + + toAck Done = AckOk + toAck (Retry d) = AckRetry d + toAck RetryDefault = AckRetry job.jobPolicy.defaultRetryDelay + toAck (Dead why) = AckDeadLetter (PoisonPill why) + +{- | Build a shibuya processor for a job with explicit tuning and a context-aware +handler. The handler must finish, or call 'extendLease', before +'visibilityTimeout' expires; otherwise PGMQ may redeliver the message +concurrently and each redelivery consumes one retry attempt. After a worker +crash, redelivery happens when the visibility timeout expires; the 'RetryPolicy' +delay only governs explicit 'Retry' and 'RetryDefault' outcomes. +-} +jobProcessorWithContext :: + ( Pgmq :> es + , Error PgmqRuntimeError :> es + , Reader PgmqAdapterEnv :> es + , IOE :> es + , Tracing :> es + ) => + JobTuning -> + Job p -> + (JobContext es -> p -> Eff es JobOutcome) -> + Eff es (ProcessorId, QueueProcessor es) +jobProcessorWithContext tuning job handle = do + env <- ask + adapter <- + pgmqAdapter env (adapterConfigFor tuning job) + >>= either (liftIO . throwIO . JobAdapterConfigInvalid) pure + pure (ProcessorId job.jobName, mkProcessor adapter (wrapHandler job handle)) + +{- | Build a shibuya processor for a job using 'defaultJobTuning': a PGMQ adapter +configured from the job's policy, paired with the wrapped handler. Pass the +result to 'runJobWorkers'. The same visibility-timeout and crash-redelivery +rules documented on 'jobProcessorWithContext' apply here. +-} +jobProcessor :: + ( Pgmq :> es + , Error PgmqRuntimeError :> es + , Reader PgmqAdapterEnv :> es + , IOE :> es + , Tracing :> es + ) => + Job p -> + (p -> Eff es JobOutcome) -> + Eff es (ProcessorId, QueueProcessor es) +jobProcessor job handle = + jobProcessorWithContext defaultJobTuning job (\_context p -> handle p) + +{- | Continuous, multi-processor run (the @rei@ cadence): run a supervised app +over several processors built with 'jobProcessor'. Returns the app handle; the +caller decides whether to block on it. The inbox size is clamped to at least 1. +-} +runJobWorkers :: + (Pgmq :> es, Reader PgmqAdapterEnv :> es, IOE :> es, Tracing :> es) => + SupervisionStrategy -> + Int -> + [Eff es (ProcessorId, QueueProcessor es)] -> + Eff es (Either AppError (AppHandle es)) +runJobWorkers strategy inboxSize procs = do + ps <- sequence procs + runApp AppConfig{strategy = strategy, inboxSize = max 1 inboxSize} ps + +{- | One-shot drain of up to @n@ messages with explicit tuning and a +context-aware handler. This reads directly from PGMQ and returns when the queue +is empty or @n@ messages have been acknowledged/retried/dead-lettered, +whichever comes first. + +If a handler throws, the message is left on the main queue and remains invisible +until the active visibility timeout expires; the drain keeps processing the rest +of the batch and does not count that message in the returned total. +-} +runJobOnceWithContext :: + (Pgmq :> es, IOE :> es, Tracing :> es) => + JobTuning -> + Int -> + Job p -> + (JobContext es -> p -> Eff es JobOutcome) -> + Eff es Int +runJobOnceWithContext tuning n job handle + | n <= 0 = pure 0 + | otherwise = drain 0 + where + drain handled + | handled >= n = pure handled + | otherwise = do + let qty = nextBatchSize (n - handled) + messages <- case tuning.ordering of + Unordered -> + Pgmq.readMessage + ReadMessage + { queueName = job.jobQueue.physicalName + , delay = tuning.visibilityTimeout + , batchSize = Just qty + , conditional = Nothing + } + FifoThroughput -> + readGrouped + ReadGrouped + { queueName = job.jobQueue.physicalName + , visibilityTimeout = tuning.visibilityTimeout + , qty = qty + } + FifoRoundRobin -> + readGroupedRoundRobin + ReadGrouped + { queueName = job.jobQueue.physicalName + , visibilityTimeout = tuning.visibilityTimeout + , qty = qty + } + if null messages + then pure handled + else do + handledInBatch <- foldM step 0 messages + drain (handled + handledInBatch) + + nextBatchSize remaining = + fromIntegral (min remaining (fromIntegral tuning.batchSize :: Int)) + + step count message = do + disposed <- processMessage message + pure $ + if disposed + then count + 1 + else count + + processMessage message + | message.readCount > job.jobPolicy.maxRetries = do + ackMessage message (AckDeadLetter MaxRetriesExceeded) + pure True + | otherwise = + case decodeJob job.jobCodec (messagePayload message) of + Left (JobPayloadFromFuture _payloadVersion _workerVersion) -> do + ackMessage message (AckRetry job.jobPolicy.defaultRetryDelay) + pure True + Left (JobPayloadMalformed err) -> do + ackMessage message (AckDeadLetter (InvalidPayload err)) + pure True + Right p -> do + outcome <- EffException.try @SomeException (handle (contextFor message) p) + case outcome of + Left _handlerException -> + pure False + Right jobOutcome -> do + ackMessage message (outcomeToAck jobOutcome) + pure True + + messagePayload = (.payload) . pgmqMessageToEnvelope + + contextFor message = + let envelope = pgmqMessageToEnvelope message + in JobContext + { extendLease = \duration -> + void $ + Pgmq.changeVisibilityTimeout + VisibilityTimeoutQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + , visibilityTimeoutOffset = nominalToSeconds duration + } + , attempt = fmap (.unAttempt) envelope.attempt + , headers = message.headers + } + + outcomeToAck Done = AckOk + outcomeToAck (Retry d) = AckRetry d + outcomeToAck RetryDefault = AckRetry job.jobPolicy.defaultRetryDelay + outcomeToAck (Dead why) = AckDeadLetter (PoisonPill why) + + ackMessage message AckOk = + void $ + Pgmq.deleteMessage + MessageQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + } + ackMessage message (AckRetry delay) = + void $ + Pgmq.changeVisibilityTimeout + VisibilityTimeoutQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + , visibilityTimeoutOffset = nominalToSeconds (retryDelaySeconds delay) + } + ackMessage message (AckDeadLetter reason) + | job.jobPolicy.useDeadLetter = do + sendDlq message reason + void $ + Pgmq.deleteMessage + MessageQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + } + | otherwise = + void $ + Pgmq.archiveMessage + MessageQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + } + ackMessage message (AckHalt _reason) = + void $ + Pgmq.changeVisibilityTimeout + VisibilityTimeoutQuery + { queueName = job.jobQueue.physicalName + , messageId = message.messageId + , visibilityTimeoutOffset = 3600 + } + + sendDlq message reason = + case message.headers of + Just headers -> + void $ + Pgmq.sendMessageWithHeaders + SendMessageWithHeaders + { queueName = job.jobQueue.dlqName + , messageBody = mkDlqPayload message reason True + , messageHeaders = MessageHeaders headers + , delay = Nothing + } + Nothing -> + void $ + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.dlqName + , messageBody = mkDlqPayload message reason True + , delay = Nothing + } + +{- | One-shot drain of up to @n@ messages (the @hospital-capacity@ cadence): +read directly from PGMQ with 'defaultJobTuning', run the handler on each +available message, and return promptly when the queue is empty. +-} +runJobOnce :: + (Pgmq :> es, IOE :> es, Tracing :> es) => + Int -> + Job p -> + (p -> Eff es JobOutcome) -> + Eff es () +runJobOnce n job handle = + void $ + runJobOnceWithContext + defaultJobTuning + n + job + (\_context p -> handle p) diff --git a/src/Keiro/PGMQ/Metrics.hs b/src/Keiro/PGMQ/Metrics.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ/Metrics.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE DataKinds #-} + +{- | Typed, 'Job'-keyed queue metrics for @keiro-pgmq@. + +PGMQ stores each queue as a table @pgmq.q_@; its @metrics()@ function reports +depth and message age. These helpers fetch that 'QueueMetrics' for a job's MAIN +queue and its DEAD-LETTER queue without the caller deriving any physical name. + +Use 'jobDlqMetrics' (its 'queueLength') for the depth alerting that +'Keiro.PGMQ.Dlq' recommends; pair it with 'Keiro.PGMQ.Dlq.archiveDlq' / +'Keiro.PGMQ.Dlq.purgeDlq' for retention. +-} +module Keiro.PGMQ.Metrics ( + QueueMetrics (..), + jobQueueMetrics, + jobDlqMetrics, + queueDepth, + allJobMetrics, +) where + +import Keiro.PGMQ.Job (Job (..)) +import Keiro.PGMQ.Runtime (QueueRef (..)) +import "base" Data.Int (Int64) +import "effectful-core" Effectful (Eff, (:>)) +import "pgmq-effectful" Pgmq.Effectful (Pgmq, QueueMetrics (..)) +import "pgmq-effectful" Pgmq.Effectful qualified as Pgmq + +-- | The metrics for a job's MAIN queue (depth, visible depth, oldest/newest age, throughput). +jobQueueMetrics :: (Pgmq :> es) => Job p -> Eff es QueueMetrics +jobQueueMetrics job = Pgmq.queueMetrics job.jobQueue.physicalName + +-- | The metrics for a job's DEAD-LETTER queue. Use its 'queueLength' for depth alerting. +jobDlqMetrics :: (Pgmq :> es) => Job p -> Eff es QueueMetrics +jobDlqMetrics job = Pgmq.queueMetrics job.jobQueue.dlqName + +-- | The main queue's immediately-readable depth (PGMQ 'queueVisibleLength'): "work waiting". +queueDepth :: (Pgmq :> es) => Job p -> Eff es Int64 +queueDepth job = do + metrics <- jobQueueMetrics job + pure metrics.queueVisibleLength + +-- | Every queue's metrics (passthrough over 'Pgmq.Effectful.allQueueMetrics'); not Job-keyed. +allJobMetrics :: (Pgmq :> es) => Eff es [QueueMetrics] +allJobMetrics = Pgmq.allQueueMetrics diff --git a/src/Keiro/PGMQ/Runtime.hs b/src/Keiro/PGMQ/Runtime.hs new file mode 100644 --- /dev/null +++ b/src/Keiro/PGMQ/Runtime.hs @@ -0,0 +1,232 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +{- | Layer 1 of @keiro-pgmq@: transport-agnostic plumbing shared by every PGMQ +integration (and, in the future, by case B — PGMQ as an integration-event +transport). This layer owns the two things every PGMQ user repeats: + + 1. Turning an arbitrary logical queue name into a PGMQ-legal physical name + plus a matching dead-letter-queue name ('QueueRef' / 'queueRef'). + 2. Running the @Pgmq : Tracing : Error PgmqRuntimeError : IOE@ effect stack + against a connection pool with an optional OpenTelemetry tracer + ('JobRuntime' / 'withJobRuntime' / 'runJobEff'). + +Nothing here knows about jobs, codecs, or retry policies — that is layer 2 +('Keiro.PGMQ.Job'). +-} +module Keiro.PGMQ.Runtime ( + -- * Queue-name derivation + QueueRef (..), + queueRef, + + -- * Runtime handle and runners + JobRuntime (..), + withJobRuntime, + runJobEff, + + -- * Re-exports + PgmqRuntimeError, +) where + +import "base" Control.Exception (bracket) +import "base" Data.Bits (xor) +import "base" Data.Char (ord) +import "base" Data.Word (Word64) +import "base" Numeric (showHex) +import "effectful-core" Effectful (Eff, IOE, runEff) +import "effectful-core" Effectful.Error.Static (Error, runErrorNoCallStack) +import "effectful-core" Effectful.Reader.Static (Reader, runReader) +import "hasql" Hasql.Connection.Settings qualified as Conn +import "hasql-pool" Hasql.Pool (Pool) +import "hasql-pool" Hasql.Pool qualified as Pool +import "hasql-pool" Hasql.Pool.Config qualified as Pool.Config +import "pgmq-core" Pgmq.Types (QueueName) +import "pgmq-core" Pgmq.Types qualified as Pgmq +import "pgmq-effectful" Pgmq.Effectful (Pgmq, PgmqRuntimeError, runPgmq, runPgmqTraced) +import "shibuya-core" Shibuya.Telemetry.Effect (Tracer, Tracing, runTracing, runTracingNoop) +import "shibuya-pgmq-adapter" Shibuya.Adapter.Pgmq (PgmqAdapterEnv, mkPgmqAdapterEnv) +import "text" Data.Text (Text) +import "text" Data.Text qualified as Text + +{- | A logical queue identity plus the PGMQ-legal physical names derived from +it. PGMQ caps queue names at 47 characters and rejects dots and other +non-@[a-z0-9_]@ characters, so a caller-facing logical name (which may contain +dots, e.g. @"hospital_capacity.reservation_work"@) must be sanitized once, +centrally, into a valid physical name plus a @"_dlq"@-suffixed dead-letter name. +-} +data QueueRef = QueueRef + { logicalName :: !Text + -- ^ Caller-facing name; may contain dots or otherwise-illegal characters. + , physicalName :: !QueueName + -- ^ Derived, PGMQ-valid name used for the main queue. + , dlqName :: !QueueName + -- ^ Derived @"_dlq"@ name used for the dead-letter queue. + } + deriving stock (Eq, Show) + +{- | Derive a 'QueueRef' from a logical name. Total: it sanitizes rather than +failing. It lower-cases, replaces every character that is not @[a-z0-9_]@ with +@'_'@, collapses repeated underscores (also trimming leading/trailing ones), +and guarantees a leading letter. + +Short sanitized names are used byte-for-byte unless they end in @"_dlq"@. +Sanitization equivalence is intentional: @"a.b"@ and @"a_b"@ name the same +queue, so distinct logical queues must differ after lower-casing and replacing +illegal characters with underscores. + +When the sanitized base exceeds 43 characters, or the base ends in @"_dlq"@, +the physical main-queue name is @_<16 hex chars>@ where the hex +suffix is FNV-1a-64 over the full original logical name. This keeps the derived +DLQ name inside PGMQ's 47-character ceiling and establishes the invariant that +physical main-queue names never end in @"_dlq"@ while derived DLQ names always +do. + +Migration note: pre-existing deployments whose sanitized logical name exceeded +43 characters, or ended in @"_dlq"@, derive a different physical queue after +this change. Messages in the old physical queue are not lost, but new workers +will not read them. Drain the old queue before upgrading or temporarily run a +worker against the old physical name. + +For example, @queueRef "hospital_capacity.reservation_work"@ yields +@physicalName == "hospital_capacity_reservation_work"@ and +@dlqName == "hospital_capacity_reservation_work_dlq"@. +-} +queueRef :: Text -> QueueRef +queueRef logical = + QueueRef + { logicalName = logical + , physicalName = forceQueueName base + , dlqName = forceQueueName (base <> "_dlq") + } + where + base = physicalBase logical + +-- | Reserve 4 characters for the @"_dlq"@ suffix below the 47-char ceiling. +maxBaseLength :: Int +maxBaseLength = 43 + +hashedPrefixLength :: Int +hashedPrefixLength = 26 + +physicalBase :: Text -> Text +physicalBase logical = + if Text.length base <= maxBaseLength && not ("_dlq" `Text.isSuffixOf` base) + then base + else hashedBase logical base + where + base = sanitize logical + +sanitize :: Text -> Text +sanitize = + ensureLeadingLetter + . collapseUnderscores + . Text.map toLegal + . Text.toLower + where + toLegal c + | (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_' = c + | otherwise = '_' + +{- | Replace runs of underscores with a single underscore and drop leading and +trailing underscores. +-} +collapseUnderscores :: Text -> Text +collapseUnderscores = + Text.intercalate "_" . filter (not . Text.null) . Text.splitOn "_" + +{- | Guarantee the name begins with an ASCII letter (PGMQ-derived table names +must start with a letter). Falls back to a bare @"q"@ for an empty result. +-} +ensureLeadingLetter :: Text -> Text +ensureLeadingLetter t = + case Text.uncons t of + Nothing -> "q" + Just (c, _) + | c >= 'a' && c <= 'z' -> t + | otherwise -> Text.cons 'q' t + +hashedBase :: Text -> Text -> Text +hashedBase logical base = + prefix <> "_" <> fnv1a64Hex logical + where + trimmedPrefix = Text.dropWhileEnd (== '_') (Text.take hashedPrefixLength base) + prefix + | Text.null trimmedPrefix = "q" + | otherwise = trimmedPrefix + +fnv1a64Hex :: Text -> Text +fnv1a64Hex logical = + Text.pack (replicate (16 - length rendered) '0' <> rendered) + where + rendered = showHex (Text.foldl' step offset logical) "" + offset :: Word64 + offset = 0xcbf29ce484222325 + prime :: Word64 + prime = 0x100000001b3 + step hash c = (hash `xor` fromIntegral (ord c)) * prime + +{- | Build a 'QueueName' from an already-sanitized 'Text'. The sanitizer should +never produce an invalid name; if it somehow does, that is a programmer error, +so we fail loudly with a clear message. +-} +forceQueueName :: Text -> QueueName +forceQueueName t = + case Pgmq.parseQueueName t of + Right qn -> qn + Left err -> + error + ( "Keiro.PGMQ.Runtime.queueRef: derived an invalid PGMQ queue name " + <> show t + <> ": " + <> show err + ) + +{- | Opaque runtime: a Hasql connection pool plus an optional OpenTelemetry +tracer. Construct it with 'withJobRuntime' (which manages the pool lifecycle). +-} +data JobRuntime = JobRuntime + { runtimePool :: !Pool + , runtimeTracer :: !(Maybe Tracer) + } + +{- | Acquire a Hasql pool from a libpq connection string, run @action@ with the +resulting 'JobRuntime', and release the pool afterwards (even on exception). +-} +withJobRuntime :: Text -> Maybe Tracer -> (JobRuntime -> IO a) -> IO a +withJobRuntime connStr tracer action = + bracket acquire Pool.release $ \pool -> + action JobRuntime{runtimePool = pool, runtimeTracer = tracer} + where + acquire = + Pool.acquire $ + Pool.Config.settings + [ Pool.Config.staticConnectionSettings (Conn.connectionString connStr) + ] + +{- | Run a @Reader PgmqAdapterEnv : Pgmq : Tracing : Error PgmqRuntimeError : IOE@ +effect action against the runtime, surfacing PGMQ errors as a @Left@. The +'Reader' supplies the 'PgmqAdapterEnv' (built from the pool) that the +shibuya-pgmq adapter needs for its transactional ack/dead-letter path, so job +processors can construct their adapter without the pool being threaded by hand. +The tracer (if any) is threaded into both the shibuya 'Tracing' effect and the +@pgmq@ interpreter: + + * @Nothing@ → 'runTracingNoop' + 'runPgmq' + * @Just tr@ → 'runTracing' + 'runPgmqTraced' + +The interpreter order matters: @runPgmq@/@runPgmqTraced@ require +@Error PgmqRuntimeError :> es@ and @IOE :> es@, so we peel @Reader@ first (it is +pure), then @Pgmq@, then @Tracing@, then @Error@, then @IOE@. +-} +runJobEff :: + JobRuntime -> + Eff '[Reader PgmqAdapterEnv, Pgmq, Tracing, Error PgmqRuntimeError, IOE] a -> + IO (Either PgmqRuntimeError a) +runJobEff rt act = + runEff $ + runErrorNoCallStack @PgmqRuntimeError $ + case rt.runtimeTracer of + Nothing -> runTracingNoop (runPgmq rt.runtimePool (runReader env act)) + Just tr -> runTracing tr (runPgmqTraced rt.runtimePool tr (runReader env act)) + where + env = mkPgmqAdapterEnv rt.runtimePool diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,981 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | End-to-end integration test for @keiro-pgmq@. + +Uses @keiro-test-support@ to start one suite-level PostgreSQL server, installs +the PGMQ schema into the migrated template database, then gives every example a +fresh cloned database. The tests drive the package's public API against that +isolated database: 'enqueue' puts work on a queue, 'runJobOnce' drains it, and +we read the queue back through @pgmq-effectful@'s 'queueMetrics' to prove that a +@Done@ handler deletes the message, a @Retry@ handler leaves it, and a @Dead@ +handler (or an undecodable payload) routes it to the dead-letter queue. +-} +module Main (main) where + +import Control.Concurrent (threadDelay) +import Control.Exception (bracket, throwIO) +import Data.Aeson (FromJSON, ToJSON, Value (..), object, parseJSON, toJSON, (.=)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Aeson.Types (parseEither) +import Data.Either (isRight) +import Data.Foldable (toList, traverse_) +import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) +import Data.Int (Int32, Int64) +import Data.List (find) +import Data.List.NonEmpty (NonEmpty (..)) +import Data.Text (Text) +import Data.Text qualified as Text +import Effectful (Eff, IOE, liftIO, (:>)) +import Effectful.Error.Static (Error) +import Effectful.Reader.Static (Reader) +import GHC.Generics (Generic) +import Hasql.Connection.Settings qualified as Conn +import Hasql.Decoders qualified as Decoders +import Hasql.Encoders qualified as Encoders +import Hasql.Pool (Pool) +import Hasql.Pool qualified as Pool +import Hasql.Pool.Config qualified as Pool.Config +import Hasql.Session qualified as Session +import Hasql.Statement qualified as Statement +import Keiro.Codec (EventType (..)) +import Keiro.Codec qualified as CoreCodec +import Keiro.PGMQ +import Keiro.Test.Postgres qualified as Postgres +import OpenTelemetry.Context qualified as Ctxt +import OpenTelemetry.Context.ThreadLocal qualified as CtxtLocal +import OpenTelemetry.Propagator.W3CTraceContext qualified as W3C +import OpenTelemetry.Trace.Core qualified as OTel +import OpenTelemetry.Trace.Id.Generator.Default (defaultIdGenerator) +import Pgmq.Config.Types qualified as Config +import Pgmq.Effectful (Message (..), MessageBody (..), Pgmq, QueueMetrics (..), ReadMessage (..), SendMessage (..)) +import Pgmq.Effectful qualified as Pgmq +import Pgmq.Migration qualified as Migration +import Pgmq.Types (QueueName, parseQueueName, queueNameToText) +import Shibuya.Adapter.Pgmq (PgmqAdapterEnv) +import Shibuya.App (AppHandle, ShutdownConfig (..), SupervisionStrategy (IgnoreFailures), stopAppGracefully) +import Shibuya.Telemetry.Effect (Tracing) +import System.Timeout (timeout) +import Test.Hspec + +-- | A sample job payload defined entirely in the test. +data Ping = Ping + { message :: Text + , count :: Int + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON, FromJSON) + +-- | The effect stack 'runJobEff' interprets. +type Stack = '[Reader PgmqAdapterEnv, Pgmq, Tracing, Error PgmqRuntimeError, IOE] + +main :: IO () +main = + Postgres.withMigratedSuiteWith installPgmq \fixture -> + hspec $ + describe "Keiro.PGMQ" $ + around (Postgres.withFreshDatabase fixture) spec + +-- | Install the PGMQ schema into the ephemeral database via @pgmq-migration@. +installPgmq :: Text -> IO () +installPgmq connStr = + withPool connStr $ \pool -> do + result <- Pool.use pool Migration.migrate + case result of + Left usageErr -> fail ("pgmq pool error during migration: " <> show usageErr) + Right (Left migErr) -> fail ("pgmq migration error: " <> show migErr) + Right (Right ()) -> pure () + +withPool :: Text -> (Pool -> IO a) -> IO a +withPool connStr = + bracket + ( Pool.acquire $ + Pool.Config.settings + [Pool.Config.staticConnectionSettings (Conn.connectionString connStr)] + ) + Pool.release + +{- | Count the rows currently in a DLQ's archive table @pgmq.a_@ via a +raw @hasql@ session. PGMQ exposes no "read the archive" function, so retention is +proven with plain SQL. The queue name is sanitized to @[a-z0-9_]@ by 'queueRef', +so interpolating it into the table identifier is safe here. +-} +archiveCount :: Text -> Text -> IO Int64 +archiveCount connStr dlqPhysical = + withPool connStr $ \pool -> do + let sql = "SELECT count(*) FROM pgmq.a_" <> dlqPhysical + session = + Session.statement () $ + Statement.preparable + sql + Encoders.noParams + (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int8))) + result <- Pool.use pool session + either (\e -> fail ("archive count failed: " <> show e)) pure result + +{- | Run a @Stack@ action against a fresh 'JobRuntime' (no tracer), failing the +test on any PGMQ runtime error. +-} +runDb :: Text -> Eff Stack a -> IO a +runDb connStr act = + withJobRuntime connStr Nothing $ \rt -> do + res <- runJobEff rt act + either (\e -> fail ("PGMQ runtime error: " <> show e)) pure res + +-- | A job over 'Ping' with a distinct queue name per test (avoids collisions). +mkJob :: Text -> Job Ping +mkJob name = + Job + { jobName = name + , jobQueue = queueRef name + , jobCodec = aesonJobCodec + , jobPolicy = defaultRetryPolicy + } + +versionedPingCodec :: CoreCodec.Codec Ping +versionedPingCodec = + CoreCodec.Codec + { eventTypes = EventType "ping" :| [] + , eventType = \_ -> EventType "ping" + , schemaVersion = 2 + , encode = toJSON + , decode = \_ value -> + case parseEither parseJSON value of + Left err -> Left (Text.pack err) + Right ping -> Right ping + , upcasters = + [ + ( 1 + , \_ value -> + case value of + String msg -> + Right $ + object + [ "message" .= msg + , "count" .= (1 :: Int) + ] + _ -> Left "expected v1 string payload" + ) + ] + } + +-- | Total number of messages currently on a queue (visible or not). +queueLen :: QueueName -> Eff Stack Int64 +queueLen q = do + metrics <- Pgmq.queueMetrics q + pure metrics.queueLength + +{- | Look up a queue by physical name in a 'Pgmq.listQueues' result and report +whether it is unlogged. 'Nothing' means the queue was not found. +-} +queueIsUnlogged :: QueueName -> [Pgmq.Queue] -> Maybe Bool +queueIsUnlogged qn queues = + fmap (.isUnlogged) (find (\q -> q.name == qn) queues) + +readOneIsEmpty :: QueueName -> Eff Stack Bool +readOneIsEmpty q = do + messages <- + Pgmq.readMessage + ReadMessage + { queueName = q + , delay = 30 + , batchSize = Just 1 + , conditional = Nothing + } + pure (null messages) + +-- | Read up to @n@ messages back off a queue (making them invisible for 30 s). +readMessages :: QueueName -> Int32 -> Eff Stack [Message] +readMessages q n = do + messages <- + Pgmq.readMessage + ReadMessage + { queueName = q + , delay = 30 + , batchSize = Just n + , conditional = Nothing + } + pure (toList messages) + +-- | Look up a single key in a @Maybe Value@ header object. +headerKey :: Text -> Maybe Value -> Maybe Value +headerKey k = \case + Just (Object o) -> KeyMap.lookup (Key.fromText k) o + _ -> Nothing + +{- | A real tracer provider with the W3C Trace Context propagator and a +non-dummy id generator, so an active span produces a @traceparent@ on injection. +No span processors are needed — the test inspects propagated headers, not +exported spans. +-} +setupW3CProvider :: IO OTel.TracerProvider +setupW3CProvider = + OTel.createTracerProvider + [] + OTel.emptyTracerProviderOptions + { OTel.tracerProviderOptionsIdGenerator = defaultIdGenerator + , OTel.tracerProviderOptionsPropagators = W3C.w3cTraceContextPropagator + } + +stopAppQuickly :: (IOE :> es) => AppHandle es -> Eff es () +stopAppQuickly app = do + _ <- stopAppGracefully ShutdownConfig{drainTimeout = 1} app + pure () + +waitUntil :: IO Bool -> IO Bool +waitUntil predicate = + maybe False id <$> timeout 10_000_000 loop + where + loop = do + ok <- predicate + if ok + then pure True + else threadDelay 100_000 >> loop + +spec :: SpecWith Text +spec = do + it "round-trips a payload through aesonJobCodec" $ \_connStr -> do + let codec = aesonJobCodec :: JobCodec Ping + sample = Ping "hello" 7 + decodeJob codec (encodeJob codec sample) `shouldBe` Right sample + + it "round-trips a payload through keiroJobCodec's versioned envelope" $ \_connStr -> do + let codec = keiroJobCodec versionedPingCodec + sample = Ping "hello" 7 + decodeJob codec (encodeJob codec sample) `shouldBe` Right sample + + it "decodes old keiroJobCodec payloads through the upcaster chain" $ \_connStr -> do + let codec = keiroJobCodec versionedPingCodec + v1Envelope = + object + [ "v" .= (1 :: Int) + , "data" .= String "legacy" + ] + decodeJob codec v1Envelope `shouldBe` Right (Ping "legacy" 1) + + it "classifies future keiroJobCodec payloads as retryable" $ \_connStr -> do + let codec = keiroJobCodec versionedPingCodec + futureEnvelope = + object + [ "v" .= (99 :: Int) + , "data" .= object [] + ] + decodeJob codec futureEnvelope `shouldBe` Left (JobPayloadFromFuture 99 2) + + it "classifies malformed keiroJobCodec envelopes as malformed payloads" $ \_connStr -> do + let codec = keiroJobCodec versionedPingCodec + decodeJob codec (String "not an envelope") `shouldSatisfy` \case + Left (JobPayloadMalformed _) -> True + _ -> False + + it "validates retry policies" $ \_connStr -> do + mkRetryPolicy 0 (RetryDelay 60) True + `shouldBe` Left (NonPositiveMaxRetries 0) + mkRetryPolicy 1 (RetryDelay (-1)) True + `shouldBe` Left (NegativeRetryDelay (RetryDelay (-1))) + mkRetryPolicy 1 (RetryDelay 0) True + `shouldBe` Right (RetryPolicy 1 (RetryDelay 0) True) + + it "validates job tuning" $ \_connStr -> do + mkJobTuning 0 1 (PollEvery 1) + `shouldBe` Left (NonPositiveVisibilityTimeout 0) + mkJobTuning 30 0 (PollEvery 1) + `shouldBe` Left (NonPositiveBatchSize 0) + mkJobTuning 30 1 (PollEvery 0) + `shouldBe` Left NonPositivePollInterval + mkJobTuning 30 1 (LongPoll 0 100) + `shouldBe` Left NonPositivePollInterval + mkJobTuning 30 1 (PollEvery 1) + `shouldBe` Right defaultJobTuning + + it "derives distinct physical names for long logical queue names" $ \_connStr -> do + let commonPrefix = Text.replicate 43 "a" + first = queueRef (commonPrefix <> "x") + second = queueRef (commonPrefix <> "y") + first.physicalName `shouldNotBe` second.physicalName + Text.length (queueNameToText first.physicalName) `shouldBe` 43 + Text.length (queueNameToText second.physicalName) `shouldBe` 43 + + it "disambiguates logical names ending in _dlq from derived DLQ names" $ \_connStr -> do + let foo = queueRef "foo" + masquerading = queueRef "foo_dlq" + masqueradingPhysical = queueNameToText masquerading.physicalName + masquerading.physicalName `shouldNotBe` foo.dlqName + masqueradingPhysical `shouldNotSatisfy` Text.isSuffixOf "_dlq" + + it "keeps short physical queue names unchanged" $ \_connStr -> do + queueNameToText (queueRef "hospital_capacity.reservation_work").physicalName + `shouldBe` "hospital_capacity_reservation_work" + + it "always derives PGMQ-parseable queue names" $ \_connStr -> do + let logicalNames = + [ "" + , "!!!" + , Text.replicate 100 "x" + , "___trailing___" + , "foo_dlq" + , "9starts.with.digit" + ] + traverse_ + ( \logical -> do + let ref = queueRef logical + physical = queueNameToText ref.physicalName + dlq = queueNameToText ref.dlqName + parseQueueName physical `shouldSatisfy` isRight + parseQueueName dlq `shouldSatisfy` isRight + physical `shouldNotSatisfy` Text.isSuffixOf "_dlq" + ) + logicalNames + + it "Done deletes the message" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.done" + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "do" 1) + runJobOnce 1 job (\_ -> pure Done) + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + it "Retry redelivers the message" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.retry" + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "again" 2) + runJobOnce 1 job (\_ -> pure (Retry (RetryDelay 0))) + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 1 + + it "RetryDefault redelivers after the policy default delay" $ \connStr -> do + let job = + (mkJob "keiro_pgmq_test.retry_default") + { jobPolicy = RetryPolicy 5 (RetryDelay 5) True + } + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "again by default" 2) + runJobOnce 1 job (\_ -> pure RetryDefault) + len <- runDb connStr (queueLen job.jobQueue.physicalName) + emptyImmediateRead <- runDb connStr (readOneIsEmpty job.jobQueue.physicalName) + len `shouldBe` 1 + emptyImmediateRead `shouldBe` True + + it "runJobOnceWithContext returns promptly when n exceeds the queue length" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.once_short_queue" + result <- + timeout 2_000_000 $ + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "only" 1) + runJobOnceWithContext defaultJobTuning 5 job \_ctx _payload -> + pure Done + result `shouldBe` Just 1 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + it "runJobOnceWithContext drains messages in batches greater than one" $ \connStr -> do + handled <- newIORef (0 :: Int) + let job = mkJob "keiro_pgmq_test.once_batch" + tuning = + either (error . show) id $ + mkJobTuning 30 2 (PollEvery 1) + drained <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "first" 1) + _ <- enqueue job (Ping "second" 2) + _ <- enqueue job (Ping "third" 3) + runJobOnceWithContext tuning 3 job \_ctx _payload -> do + liftIO $ modifyIORef' handled (+ 1) + pure Done + drained `shouldBe` 3 + readIORef handled `shouldReturn` 3 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + it "runJobOnceWithContext Retry delay hides the message until the delay expires" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.once_retry_delay" + drained <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "later" 1) + runJobOnceWithContext defaultJobTuning 1 job \_ctx _payload -> + pure (Retry (RetryDelay 5)) + drained `shouldBe` 1 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + emptyImmediateRead <- runDb connStr (readOneIsEmpty job.jobQueue.physicalName) + len `shouldBe` 1 + emptyImmediateRead `shouldBe` True + + it "runJobOnceWithContext leaves thrown-handler messages invisible and continues" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.once_throw" + tuning = + either (error . show) id $ + mkJobTuning 2 2 (PollEvery 1) + drained <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "throw" 1) + _ <- enqueue job (Ping "ok" 2) + runJobOnceWithContext tuning 2 job \_ctx payload -> + if payload.message == "throw" + then liftIO $ throwIO (userError "handler failed") + else pure Done + drained `shouldBe` 1 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + emptyImmediateRead <- runDb connStr (readOneIsEmpty job.jobQueue.physicalName) + len `shouldBe` 1 + emptyImmediateRead `shouldBe` True + threadDelay 2_200_000 + drainedAfterVisibilityTimeout <- + runDb connStr $ + runJobOnceWithContext tuning 1 job \_ctx _payload -> + pure Done + drainedAfterVisibilityTimeout `shouldBe` 1 + + it "runJobOnceWithContext auto-routes max-retry messages to the DLQ before rerunning the handler" $ \connStr -> do + callCount <- newIORef (0 :: Int) + let job = + (mkJob "keiro_pgmq_test.once_max_retries") + { jobPolicy = RetryPolicy 1 (RetryDelay 0) True + } + firstDrain <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "retry-limit" 1) + runJobOnceWithContext defaultJobTuning 1 job \_ctx _payload -> do + liftIO $ modifyIORef' callCount (+ 1) + pure (Retry (RetryDelay 0)) + secondDrain <- + runDb connStr $ + runJobOnceWithContext defaultJobTuning 1 job \_ctx _payload -> do + liftIO $ modifyIORef' callCount (+ 1) + pure Done + firstDrain `shouldBe` 1 + secondDrain `shouldBe` 1 + readIORef callCount `shouldReturn` 1 + mainLen <- runDb connStr (queueLen job.jobQueue.physicalName) + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + mainLen `shouldBe` 0 + dlqLen `shouldBe` 1 + + it "worker-path lease extension prevents redelivery" $ \connStr -> do + callCount <- newIORef (0 :: Int) + handlerDone <- newIORef False + let job = mkJob "keiro_pgmq_test.worker_lease" + tuning = + either (error . show) id $ + mkJobTuning 2 1 (PollEvery 0.2) + processed <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "slow" 4) + result <- + runJobWorkers + IgnoreFailures + 16 + [ jobProcessorWithContext tuning job \ctx _payload -> do + liftIO $ modifyIORef' callCount (+ 1) + ctx.extendLease 30 + liftIO $ threadDelay 4_000_000 + liftIO $ writeIORef handlerDone True + pure Done + ] + case result of + Left err -> liftIO $ fail ("runJobWorkers failed: " <> show err) + Right app -> do + ok <- liftIO $ waitUntil (readIORef handlerDone) + stopAppQuickly app + pure ok + processed `shouldBe` True + readIORef callCount `shouldReturn` 1 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + it "worker-path context exposes the first attempt number" $ \connStr -> do + seenAttempt <- newIORef Nothing + let job = mkJob "keiro_pgmq_test.worker_attempt" + processed <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "attempt" 1) + result <- + runJobWorkers + IgnoreFailures + 16 + [ jobProcessorWithContext defaultJobTuning job \ctx _payload -> do + liftIO $ writeIORef seenAttempt (Just ctx.attempt) + pure Done + ] + case result of + Left err -> liftIO $ fail ("runJobWorkers failed: " <> show err) + Right app -> do + ok <- liftIO $ waitUntil ((/= Nothing) <$> readIORef seenAttempt) + stopAppQuickly app + pure ok + processed `shouldBe` True + readIORef seenAttempt `shouldReturn` Just (Just 0) + + it "runJobWorkers processes an enqueued message" $ \connStr -> do + processedRef <- newIORef False + let job = mkJob "keiro_pgmq_test.worker_smoke" + processed <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "worker" 1) + result <- + runJobWorkers + IgnoreFailures + 16 + [ jobProcessor job \_payload -> do + liftIO $ writeIORef processedRef True + pure Done + ] + case result of + Left err -> liftIO $ fail ("runJobWorkers failed: " <> show err) + Right app -> do + ok <- liftIO $ waitUntil (readIORef processedRef) + stopAppQuickly app + pure ok + processed `shouldBe` True + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + it "runJobWorkers survives a transient database error during polling" $ \_connStr -> + pendingWith "needs a deterministic keiro-pgmq-level transient polling fault injector; EP-1 covers this in upstream shibuya and shibuya-pgmq-adapter tests" + + it "worker-path retry limit auto-routes to the DLQ before the handler reruns" $ \connStr -> do + callCount <- newIORef (0 :: Int) + let job = + (mkJob "keiro_pgmq_test.worker_max_retries") + { jobPolicy = RetryPolicy 1 (RetryDelay 0) True + } + tuning = + either (error . show) id $ + mkJobTuning 30 1 (PollEvery 0.1) + dlqReached <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "retry-limit" 1) + result <- + runJobWorkers + IgnoreFailures + 16 + [ jobProcessorWithContext tuning job \_ctx _payload -> do + liftIO $ modifyIORef' callCount (+ 1) + pure (Retry (RetryDelay 0)) + ] + case result of + Left err -> liftIO $ fail ("runJobWorkers failed: " <> show err) + Right app -> do + ok <- liftIO $ waitUntil do + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + pure (dlqLen == 1) + stopAppQuickly app + pure ok + dlqReached `shouldBe` True + readIORef callCount `shouldReturn` 1 + + it "enqueueWithDelay delays first delivery" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.enqueue_delay" + runDb connStr $ do + ensureJobQueue job + _ <- enqueueWithDelay job 5 (Ping "later" 1) + pure () + len <- runDb connStr (queueLen job.jobQueue.physicalName) + emptyImmediateRead <- runDb connStr (readOneIsEmpty job.jobQueue.physicalName) + len `shouldBe` 1 + emptyImmediateRead `shouldBe` True + + it "Dead routes the message to the DLQ" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dead" + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "poison" 3) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + mainLen <- runDb connStr (queueLen job.jobQueue.physicalName) + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + mainLen `shouldBe` 0 + dlqLen `shouldBe` 1 + + it "readDlq decodes the original dead-lettered payload" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_read" + payload = Ping "poison" 3 + entries <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job payload + runJobOnce 1 job (\_ -> pure (Dead "bad")) + readDlq job 1 + case entries of + [entry] -> do + entry.reason `shouldSatisfy` Text.isPrefixOf "poison_pill" + entry.originalPayload `shouldBe` Right payload + entry.originalMessageId `shouldSatisfy` (/= Nothing) + entry.readCount `shouldBe` Just 1 + _ -> expectationFailure ("expected one DLQ entry, got " <> show (length entries)) + + it "redriveDlq moves dead-lettered payloads back to the main queue" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_redrive" + redriven <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "redrive" 1) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + redriveDlq job 10 + redriven `shouldBe` 1 + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + mainLen <- runDb connStr (queueLen job.jobQueue.physicalName) + dlqLen `shouldBe` 0 + mainLen `shouldBe` 1 + runDb connStr $ + runJobOnce 1 job (\_ -> pure Done) + finalMainLen <- runDb connStr (queueLen job.jobQueue.physicalName) + finalMainLen `shouldBe` 0 + + it "purgeDlq empties the DLQ" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_purge" + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "purge" 1) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + purgeDlq job + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + dlqLen `shouldBe` 0 + + it "readDlq preserves malformed DLQ wrappers as malformed entries" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_malformed" + entries <- + runDb connStr $ do + ensureJobQueue job + _ <- + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.dlqName + , messageBody = MessageBody (String "not a dlq wrapper") + , delay = Nothing + } + readDlq job 1 + case entries of + [entry] -> do + entry.reason `shouldSatisfy` Text.isPrefixOf "malformed_dlq_payload" + entry.originalPayload `shouldSatisfy` \case + Left (JobPayloadMalformed _) -> True + _ -> False + _ -> expectationFailure ("expected one malformed DLQ entry, got " <> show (length entries)) + + it "undecodable payload routes to the DLQ" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.bad" + runDb connStr $ do + ensureJobQueue job + -- Send raw JSON the Ping codec cannot decode, bypassing enqueue. + _ <- + Pgmq.sendMessage + SendMessage + { queueName = job.jobQueue.physicalName + , messageBody = MessageBody (String "not a ping") + , delay = Nothing + } + runJobOnce 1 job (\_ -> pure Done) + mainLen <- runDb connStr (queueLen job.jobQueue.physicalName) + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + mainLen `shouldBe` 0 + dlqLen `shouldBe` 1 + + -- EP-1 M1: header-carrying enqueue and the reserved-key contract. + it "enqueueWithHeaders attaches a header readable on the raw PGMQ message" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.hdr_attach" + msgs <- + runDb connStr $ do + ensureJobQueue job + _ <- + enqueueWithHeaders + job + (MessageHeaders (object ["tenant" .= ("acme" :: Text)])) + (Ping "hdr" 1) + readMessages job.jobQueue.physicalName 1 + case msgs of + [m] -> headerKey "tenant" m.headers `shouldBe` Just (String "acme") + _ -> expectationFailure ("expected one message, got " <> show (length msgs)) + + it "enqueueWithHeaders leaves the x-pgmq-group key untouched" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.hdr_group" + msgs <- + runDb connStr $ do + ensureJobQueue job + _ <- + enqueueWithHeaders + job + (MessageHeaders (object ["x-pgmq-group" .= ("g1" :: Text)])) + (Ping "g" 1) + readMessages job.jobQueue.physicalName 1 + case msgs of + [m] -> headerKey "x-pgmq-group" m.headers `shouldBe` Just (String "g1") + _ -> expectationFailure ("expected one message, got " <> show (length msgs)) + + -- EP-1 M2: batch enqueue. + it "enqueueBatch of three payloads yields three ids and queue depth three" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.batch" + ids <- + runDb connStr $ do + ensureJobQueue job + enqueueBatch job [Ping "a" 1, Ping "b" 2, Ping "c" 3] + length ids `shouldBe` 3 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 3 + + it "enqueueBatchWithHeaders attaches per-message headers" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.batch_headers" + msgs <- + runDb connStr $ do + ensureJobQueue job + _ <- + enqueueBatchWithHeaders + job + [ (MessageHeaders (object ["i" .= (1 :: Int)]), Ping "a" 1) + , (MessageHeaders (object ["i" .= (2 :: Int)]), Ping "b" 2) + ] + readMessages job.jobQueue.physicalName 2 + length msgs `shouldBe` 2 + map (headerKey "i" . (.headers)) msgs + `shouldMatchList` [Just (Number 1), Just (Number 2)] + + -- EP-1 M3: handler-visible headers and trace propagation. + it "drain-path JobContext exposes the enqueued headers" $ \connStr -> do + seen <- newIORef Nothing + let job = mkJob "keiro_pgmq_test.ctx_headers" + runDb connStr $ do + ensureJobQueue job + _ <- + enqueueWithHeaders + job + (MessageHeaders (object ["tenant" .= ("acme" :: Text)])) + (Ping "h" 1) + _ <- + runJobOnceWithContext defaultJobTuning 1 job \ctx _payload -> do + liftIO (writeIORef seen ctx.headers) + pure Done + pure () + captured <- readIORef seen + headerKey "tenant" captured `shouldBe` Just (String "acme") + + it "a traceparent set at enqueue is visible to the drain-path handler" $ \connStr -> do + seen <- newIORef Nothing + provider <- setupW3CProvider + let tracer = OTel.makeTracer provider "keiro-pgmq-test" OTel.tracerOptions + job = mkJob "keiro_pgmq_test.traceparent" + parentSpan <- OTel.createSpan tracer Ctxt.empty "enqueue" OTel.defaultSpanArguments + _ <- CtxtLocal.attachContext (Ctxt.insertSpan parentSpan Ctxt.empty) + runDb connStr $ do + ensureJobQueue job + _ <- enqueueTraced provider job (MessageHeaders (object [])) (Ping "t" 1) + _ <- + runJobOnceWithContext defaultJobTuning 1 job \ctx _payload -> do + liftIO (writeIORef seen ctx.headers) + pure Done + pure () + OTel.endSpan parentSpan Nothing + captured <- readIORef seen + headerKey "traceparent" captured `shouldSatisfy` \case + Just (String _) -> True + _ -> False + + -- EP-2 M1: unlogged vs standard provisioning. + it "ensureJobQueueWith unlogged creates an unlogged queue" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.unlogged" + unlogged <- + runDb connStr $ do + ensureJobQueueWith unloggedProvision job + queues <- Pgmq.listQueues + pure (queueIsUnlogged job.jobQueue.physicalName queues) + unlogged `shouldBe` Just True + + it "ensureJobQueue (standard) creates a logged queue" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.standard_logged" + unlogged <- + runDb connStr $ do + ensureJobQueue job + queues <- Pgmq.listQueues + pure (queueIsUnlogged job.jobQueue.physicalName queues) + unlogged `shouldBe` Just False + + -- EP-2 M2: partitioned config shape (pure) + pending live test. + it "ensureJobQueueWith partitioned builds a partitioned QueueConfig" $ \_connStr -> do + let job = mkJob "keiro_pgmq_test.partitioned" + spec = PartitionSpec{partitionInterval = "daily", retentionInterval = "7 days"} + case queueProvisionConfigs (partitionedProvision spec) job of + (mainCfg : _) -> + case mainCfg.queueType of + Config.PartitionedQueue pc -> do + pc.partitionInterval `shouldBe` "daily" + pc.retentionInterval `shouldBe` "7 days" + mainCfg.queueName `shouldBe` job.jobQueue.physicalName + other -> + expectationFailure + ("expected PartitionedQueue, got " <> show other) + [] -> expectationFailure "expected at least the main queue config" + + it "ensureJobQueueWith partitioned creates a partitioned queue (live)" $ \_connStr -> + pendingWith + "requires a pg_partman-enabled PostgreSQL; the keiro test database installs only \ + \the PGMQ schema via pgmq-migration, which does not load pg_partman" + + -- EP-2 M3: FIFO index idempotence. + it "ensureFifoIndex is idempotent and the queue still accepts reads" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.fifo_index" + roundTripped <- + runDb connStr $ do + ensureJobQueue job + ensureFifoIndex job + ensureFifoIndex job -- second call must not error + _ <- enqueue job (Ping "after-index" 1) + runJobOnce 1 job (\_ -> pure Done) + queueLen job.jobQueue.physicalName + roundTripped `shouldBe` 0 + + -- EP-3 M3: group-keyed producer + ordered queue setup. + it "enqueueToGroup writes the x-pgmq-group header" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.group_header" + msgs <- + runDb connStr $ do + ensureOrderedJobQueue job + _ <- enqueueToGroup job "g1" (Ping "grouped" 1) + readMessages job.jobQueue.physicalName 1 + case msgs of + [m] -> headerKey "x-pgmq-group" m.headers `shouldBe` Just (String "g1") + _ -> expectationFailure ("expected one message, got " <> show (length msgs)) + + it "ensureOrderedJobQueue is idempotent and the queue accepts grouped work" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.ordered_setup" + len <- + runDb connStr $ do + ensureOrderedJobQueue job + ensureOrderedJobQueue job -- second call must not error + _ <- enqueueToGroup job "g1" (Ping "x" 1) + _ <- + runJobOnceWithContext (withOrdering FifoThroughput defaultJobTuning) 1 job \_ctx _p -> + pure Done + queueLen job.jobQueue.physicalName + len `shouldBe` 0 + + -- EP-3 M4: end-to-end ordering proof. + it "FifoThroughput drain preserves strict within-group order and fully drains" $ \connStr -> do + observed <- newIORef ([] :: [Text]) + let job = mkJob "keiro_pgmq_test.fifo_order" + drained <- + runDb connStr $ do + ensureOrderedJobQueue job + _ <- enqueueToGroup job "a" (Ping "a1" 1) + _ <- enqueueToGroup job "b" (Ping "b1" 1) + _ <- enqueueToGroup job "a" (Ping "a2" 2) + _ <- enqueueToGroup job "a" (Ping "a3" 3) + _ <- enqueueToGroup job "b" (Ping "b2" 2) + runJobOnceWithContext (withOrdering FifoThroughput defaultJobTuning) 5 job \_ctx payload -> do + liftIO $ modifyIORef' observed (<> [payload.message]) + pure Done + log' <- readIORef observed + drained `shouldBe` 5 + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + filter (Text.isPrefixOf "a") log' `shouldBe` ["a1", "a2", "a3"] + filter (Text.isPrefixOf "b") log' `shouldBe` ["b1", "b2"] + + it "FifoThroughput worker path preserves within-group order" $ \connStr -> do + observed <- newIORef ([] :: [Text]) + let job = mkJob "keiro_pgmq_test.fifo_worker" + tuning = + withOrdering FifoThroughput $ + either (error . show) id $ + mkJobTuning 30 1 (PollEvery 0.1) + processed <- + runDb connStr $ do + ensureOrderedJobQueue job + _ <- enqueueToGroup job "a" (Ping "a1" 1) + _ <- enqueueToGroup job "a" (Ping "a2" 2) + _ <- enqueueToGroup job "a" (Ping "a3" 3) + result <- + runJobWorkers + IgnoreFailures + 16 + [ jobProcessorWithContext tuning job \_ctx payload -> do + liftIO $ modifyIORef' observed (<> [payload.message]) + pure Done + ] + case result of + Left err -> liftIO $ fail ("runJobWorkers failed: " <> show err) + Right app -> do + ok <- liftIO $ waitUntil ((>= 3) . length <$> readIORef observed) + stopAppQuickly app + pure ok + processed `shouldBe` True + log' <- readIORef observed + filter (Text.isPrefixOf "a") log' `shouldBe` ["a1", "a2", "a3"] + len <- runDb connStr (queueLen job.jobQueue.physicalName) + len `shouldBe` 0 + + -- EP-4 M1: typed metrics surface (main + DLQ). + it "jobQueueMetrics reports main-queue depth after enqueue" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.metrics_depth" + (mainMetrics, dlqMetrics) <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "a" 1) + _ <- enqueue job (Ping "b" 2) + _ <- enqueue job (Ping "c" 3) + mainMetrics <- jobQueueMetrics job + dlqMetrics <- jobDlqMetrics job + pure (mainMetrics, dlqMetrics) + mainMetrics.queueLength `shouldBe` 3 + mainMetrics.queueVisibleLength `shouldBe` 3 + dlqMetrics.queueLength `shouldBe` 0 + + it "jobDlqMetrics reports DLQ depth after a Dead outcome" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.metrics_dlq" + (mainMetrics, dlqMetrics) <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "poison" 1) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + mainMetrics <- jobQueueMetrics job + dlqMetrics <- jobDlqMetrics job + pure (mainMetrics, dlqMetrics) + mainMetrics.queueLength `shouldBe` 0 + dlqMetrics.queueLength `shouldBe` 1 + + -- EP-4 M2: archive/retention API. + it "archiveDlq retains dead-lettered rows in the archive table" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_archive" + (archived, dlqLen) <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "poison" 1) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + archived <- archiveDlq job 10 + dlqMetrics <- jobDlqMetrics job + pure (archived, dlqMetrics.queueLength) + archived `shouldBe` 1 + dlqLen `shouldBe` 0 + retained <- archiveCount connStr (queueNameToText job.jobQueue.dlqName) + retained `shouldBe` 1 + + -- EP-4 M3: end-to-end retention lifecycle. + it "archived DLQ rows survive a purge" $ \connStr -> do + let job = mkJob "keiro_pgmq_test.dlq_archive_purge" + archived <- + runDb connStr $ do + ensureJobQueue job + _ <- enqueue job (Ping "poison" 1) + runJobOnce 1 job (\_ -> pure (Dead "bad")) + archived <- archiveDlq job 10 + purgeDlq job + pure archived + archived `shouldBe` 1 + dlqLen <- runDb connStr (queueLen job.jobQueue.dlqName) + dlqLen `shouldBe` 0 + retained <- archiveCount connStr (queueNameToText job.jobQueue.dlqName) + retained `shouldBe` 1