diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,14 @@ +# Changelog + +## 0.1.0.0 — 2026-05-23 + +### New Features + +* Initial release. +* `kirokuAdapter` — create a Shibuya `Adapter es RecordedEvent` from a + Kiroku store handle and subscription configuration. +* `KirokuAdapterConfig` — subscription name, target, batch size, and + TBQueue buffer size. +* Ack semantics: AckOk/AckRetry/AckDeadLetter are no-ops (checkpoint + managed by Kiroku); AckHalt cancels the subscription. +* `Shibuya.Adapter.Kiroku.Convert` — RecordedEvent to Envelope mapping. diff --git a/shibuya-kiroku-adapter.cabal b/shibuya-kiroku-adapter.cabal new file mode 100644 --- /dev/null +++ b/shibuya-kiroku-adapter.cabal @@ -0,0 +1,87 @@ +cabal-version: 3.0 +name: shibuya-kiroku-adapter +version: 0.1.0.0 +synopsis: + Kiroku event store adapter for the Shibuya queue processing framework + +description: + A Shibuya adapter that integrates with Kiroku, a PostgreSQL event store. + Wraps Kiroku's push-based subscriptions into Shibuya's pull-based Adapter + interface, enabling supervised multi-subscription processing with failure + isolation, per-subscription metrics, and coordinated graceful shutdown. + . + Events flow through a bounded TBQueue bridge (provided by + @kiroku-store@'s @subscriptionStream@) and are lifted into Shibuya's + effect stack via @Stream.morphInner@. Checkpoint management is handled + internally by Kiroku's subscription worker — the adapter's ack semantics + are no-op for AckOk\/AckRetry\/AckDeadLetter and trigger subscription + cancellation for AckHalt. + +author: Nadeem Bitar +maintainer: nadeem@gmail.com +license: BSD-3-Clause +build-type: Simple +category: Concurrency, Database, Eventing +extra-doc-files: CHANGELOG.md +homepage: https://github.com/shinzui/kiroku +bug-reports: https://github.com/shinzui/kiroku/issues + +source-repository head + type: git + location: https://github.com/shinzui/kiroku.git + +common common + default-language: GHC2024 + default-extensions: + DeriveAnyClass + DuplicateRecordFields + OverloadedLabels + OverloadedStrings + + ghc-options: -Wall + +library + import: common + exposed-modules: + Shibuya.Adapter.Kiroku + Shibuya.Adapter.Kiroku.Convert + + build-depends: + , aeson >=2.1 && <2.3 + , base >=4.18 && <5 + , effectful-core >=2.4 && <2.7 + , kiroku-store ^>=0.1 + , shibuya-core >=0.5 && <0.6 + , streamly-core >=0.3 && <0.4 + , text >=2.0 && <2.2 + , unordered-containers >=0.2 && <0.3 + , uuid >=1.3 && <1.4 + + hs-source-dirs: src + +test-suite shibuya-kiroku-adapter-test + import: common + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + , aeson >=2.1 && <2.3 + , base >=4.18 && <5 + , containers >=0.6 && <0.8 + , directory + , effectful >=2.4 && <2.7 + , ephemeral-pg >=0.2 && <0.3 + , generic-lens >=2.2 && <2.4 + , hasql >=1.10 && <1.11 + , hasql-pool >=1.2 && <1.5 + , hspec >=2.10 && <2.12 + , kiroku-store ^>=0.1 + , kiroku-test-support + , lens >=5.2 && <5.4 + , shibuya-core >=0.5 && <0.6 + , shibuya-kiroku-adapter + , stm >=2.5 && <2.6 + , text >=2.0 && <2.2 + , time >=1.12 && <1.15 + , uuid >=1.3 && <1.4 diff --git a/src/Shibuya/Adapter/Kiroku.hs b/src/Shibuya/Adapter/Kiroku.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Kiroku.hs @@ -0,0 +1,209 @@ +{- | Kiroku event store adapter for the Shibuya queue processing framework. + +This adapter wraps Kiroku's push-based subscriptions into Shibuya's +pull-based 'Adapter' interface. Events are bridged through a bounded +'TBQueue' (via @kiroku-store@'s 'subscriptionStream') and lifted into the +effectful stack with @Stream.morphInner@. + +== Example + +@ +import Effectful (runEff) +import Kiroku.Store (withStore, defaultConnectionSettings) +import Shibuya.Adapter.Kiroku +import Shibuya.App +import Shibuya.Core.Ack (AckDecision (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) + +main :: IO () +main = withStore settings $ \\store -> + runEff $ runTracingNoop $ do + adapter <- kirokuAdapter store + KirokuAdapterConfig + { subscriptionName = SubscriptionName \"my-projection\" + , subscriptionTarget = AllStreams + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + + let handler ingested = do + -- process ingested.envelope.payload :: RecordedEvent + pure AckOk + + Right appHandle <- runApp IgnoreFailures 100 + [(ProcessorId \"my-projection\", mkProcessor adapter handler)] + + waitApp appHandle +@ + +== Consumer-Group Example (size 4) + +A consumer group splits one logical subscription across @N@ members. Each +originating stream is deterministically assigned to exactly one member (by a +hash computed in PostgreSQL), so same-stream events stay ordered while distinct +streams are processed in parallel. To run all four members in one process, build +one adapter per member index and wire each into its own processor: + +@ +main :: IO () +main = withStore settings $ \\store -> + runEff $ runTracingNoop $ do + let mkMemberAdapter m = + kirokuAdapter store + KirokuAdapterConfig + { subscriptionName = SubscriptionName \"orders-projection\" + , subscriptionTarget = Category (CategoryName \"orders\") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Just (ConsumerGroup { member = m, size = 4 }) + } + + adapters <- mapM mkMemberAdapter [0, 1, 2, 3] + + let processors = + [ (ProcessorId (\"orders-\" <> T.pack (show m)), mkProcessor (adapters !! m) handler) + | m <- [0 .. 3] + ] + + Right appHandle <- runApp IgnoreFailures 100 processors + waitApp appHandle + where + handler ingested = do + -- process ingested.envelope.payload :: RecordedEvent + pure AckOk +@ + +To run members across separate processes instead, give each process one adapter +with its own 'member' index and the same 'subscriptionName'. Kiroku's per-member +checkpoint (keyed by @(subscriptionName, member)@) lets each process resume from +its own position after a restart. Exactly one live process must own each member +index at a time. + +== Ack Semantics + +Kiroku subscriptions are fundamentally different from message queues: +events are immutable and persistent, and checkpoint advancement is managed +internally by the subscription worker. + +* 'AckOk' — no-op (the normal case; checkpoint advances automatically). +* 'AckRetry' — no-op (events cannot be redelivered; they are always available). +* 'AckDeadLetter' — no-op (there is no dead-letter concept for an event log). +* 'AckHalt' — cancels the underlying Kiroku subscription. + +== Backpressure + +The @bufferSize@ field in 'KirokuAdapterConfig' controls the 'TBQueue' +capacity. When the queue is full, the Kiroku subscription worker blocks +until the Shibuya handler drains events, providing natural backpressure. +-} +module Shibuya.Adapter.Kiroku ( + -- * Adapter + kirokuAdapter, + + -- * Configuration + KirokuAdapterConfig (..), + + -- * Re-exports from kiroku-store + SubscriptionName (..), + SubscriptionTarget (..), + ConsumerGroup (..), +) where + +import Data.Int (Int32) +import Effectful (Eff, IOE, liftIO, (:>)) +import Kiroku.Store.Connection (KirokuStore) +import Kiroku.Store.Subscription.Stream (subscriptionStream) +import Kiroku.Store.Subscription.Types ( + ConsumerGroup (..), + OverflowPolicy (..), + SubscriptionConfigM (..), + SubscriptionName (..), + SubscriptionResult (..), + SubscriptionTarget (..), + defaultSubscriptionConfig, + ) +import Kiroku.Store.Types (RecordedEvent) +import Numeric.Natural (Natural) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Adapter.Kiroku.Convert (toIngested) +import Streamly.Data.Stream qualified as Stream + +{- | Configuration for creating a Kiroku adapter. + +@subscriptionName@ must be unique across all active subscriptions — it +identifies the checkpoint row in the @subscriptions@ table. + +@bufferSize@ controls backpressure: the Kiroku worker blocks when the +internal queue is full, throttling database polling to match the handler's +consumption rate. +-} +data KirokuAdapterConfig = KirokuAdapterConfig + { subscriptionName :: !SubscriptionName + -- ^ Unique subscription identifier (checkpoint key) + , subscriptionTarget :: !SubscriptionTarget + -- ^ 'AllStreams' or @'Category' categoryName@ + , batchSize :: !Int32 + -- ^ Events per database fetch during catch-up + , bufferSize :: !Natural + -- ^ TBQueue capacity (backpressure threshold) + , consumerGroup :: !(Maybe ConsumerGroup) + {- ^ Optional consumer-group membership for this adapter instance. + 'Nothing' (the default) = ordinary single-consumer subscription. + @'Just' ('ConsumerGroup' { member = m, size = n })@ = this adapter is + member @m@ of a group of size @n@, receiving only the events whose + originating stream hashes to slot @m@ (in global-position order). To run a + full size-@n@ group, create @n@ adapters with the same 'subscriptionName' + and distinct 'member' indices, each backed by its own Shibuya processor. + + The validity invariant (@size >= 1@, @0 <= member < size@) is enforced by + the underlying 'Kiroku.Store.Subscription.subscribe' call, which throws + 'Kiroku.Store.Subscription.Types.InvalidConsumerGroup' on violation. + -} + } + +{- | Create a Shibuya 'Adapter' backed by a Kiroku subscription. + +The adapter: + +1. Calls 'subscriptionStream' to start a Kiroku subscription with a + bounded 'TBQueue' bridge. +2. Lifts the @Stream IO RecordedEvent@ to @Stream (Eff es)@ via + @Stream.morphInner liftIO@. +3. Wraps each 'RecordedEvent' into an 'Ingested' value with an + 'Envelope' (mapping event ID → message ID, global position → cursor) + and a no-op 'AckHandle' (except 'AckHalt' which cancels the + subscription). + +The returned adapter's @shutdown@ action cancels the underlying +subscription and flushes the sentinel through the queue so any +blocked stream reader terminates. +-} +kirokuAdapter :: + (IOE :> es) => + KirokuStore -> + KirokuAdapterConfig -> + Eff es (Adapter es RecordedEvent) +kirokuAdapter store (KirokuAdapterConfig subName subTarget bs buf cg) = do + -- Build from 'defaultSubscriptionConfig' and override only the non-default + -- fields. Using the smart constructor (rather than a full record literal) + -- means any future field added to 'SubscriptionConfigM' is inherited at its + -- default automatically — e.g. EP-2's 'consumerGroupGuard', left 'False' here. + let subConfig = + (defaultSubscriptionConfig subName subTarget (\_ -> pure Continue)) + { batchSize = bs + , queueCapacity = 16 + , overflowPolicy = DropSubscription + , consumerGroup = cg + } + + (ioStream, cancelAction) <- liftIO $ subscriptionStream store subConfig buf + + let ingestedStream = fmap (toIngested cancelAction) (Stream.morphInner liftIO ioStream) + + pure + Adapter + { adapterName = "kiroku" + , source = ingestedStream + , shutdown = liftIO cancelAction + } diff --git a/src/Shibuya/Adapter/Kiroku/Convert.hs b/src/Shibuya/Adapter/Kiroku/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Adapter/Kiroku/Convert.hs @@ -0,0 +1,96 @@ +{- | Conversion from Kiroku's 'RecordedEvent' to Shibuya's 'Ingested' and +'Envelope' types. + +== Envelope Mapping + +@ +RecordedEvent field → Envelope field +───────────────────────────────────────── +eventId (UUID) → messageId (Text) +globalPosition → cursor (CursorInt) +createdAt → enqueuedAt +metadata.traceparent → traceContext +(the event itself) → payload +(none) → partition = Nothing +@ + +The adapter preserves W3C trace-context metadata when @metadata@ is a JSON +object containing a string @traceparent@ key. A string @tracestate@ key is +included when present. +-} +module Shibuya.Adapter.Kiroku.Convert ( + -- * Conversion + toIngested, + toEnvelope, +) where + +import Data.Aeson (Value (..)) +import Data.Aeson.Key qualified as Key +import Data.Aeson.KeyMap qualified as KM +import Data.HashMap.Strict qualified as HashMap +import Data.Text qualified as T +import Data.Text.Encoding qualified as TE +import Data.UUID qualified as UUID +import Effectful (IOE, liftIO, (:>)) +import Kiroku.Store.Types ( + EventId (..), + GlobalPosition (..), + RecordedEvent (..), + ) +import Shibuya.Core.Ack (AckDecision (..)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), TraceHeaders) + +{- | Wrap a 'RecordedEvent' into an 'Ingested' value suitable for Shibuya +handlers. + +The 'AckHandle' semantics: + +* 'AckOk', 'AckRetry', 'AckDeadLetter' — no-op (checkpoint is managed by + the Kiroku subscription worker, not by the handler). +* 'AckHalt' — invokes the provided cancel action, stopping the underlying + Kiroku subscription. +-} +toIngested :: (IOE :> es) => IO () -> RecordedEvent -> Ingested es RecordedEvent +toIngested cancelAction event = + Ingested + { envelope = toEnvelope event + , ack = + AckHandle + { finalize = \case + AckHalt _ -> liftIO cancelAction + _ -> pure () + } + , lease = Nothing + } + +{- | Convert a 'RecordedEvent' to a Shibuya 'Envelope'. + +The event's UUID is formatted as text for the 'MessageId', and the +global position is used as an integer 'Cursor' for ordering. +-} +toEnvelope :: RecordedEvent -> Envelope RecordedEvent +toEnvelope event = + let RecordedEvent{eventId = EventId uuid, globalPosition = GlobalPosition pos, createdAt = ts, metadata = meta} = event + in Envelope + { messageId = MessageId (T.pack (UUID.toString uuid)) + , cursor = Just (CursorInt (fromIntegral pos)) + , partition = Nothing + , enqueuedAt = Just ts + , traceContext = metadataTraceContext meta + , attempt = Nothing + , attributes = HashMap.empty + , payload = event + } + +metadataTraceContext :: Maybe Value -> Maybe TraceHeaders +metadataTraceContext (Just (Object metadata)) = do + String traceparent <- KM.lookup (Key.fromString "traceparent") metadata + let traceparentHeader = ("traceparent", TE.encodeUtf8 traceparent) + traceHeaders = + case KM.lookup (Key.fromString "tracestate") metadata of + Just (String tracestate) -> [traceparentHeader, ("tracestate", TE.encodeUtf8 tracestate)] + _ -> [traceparentHeader] + pure traceHeaders +metadataTraceContext _ = Nothing diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,503 @@ +module Main where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.STM (atomically, newTVarIO, readTVar, registerDelay, writeTVar) +import Control.Concurrent.STM qualified as STM +import Control.Lens ((^.)) +import Control.Monad (unless) +import Control.Monad.IO.Class (liftIO) +import Data.Aeson (Value) +import Data.Aeson qualified as Aeson +import Data.Generics.Labels () +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.Int (Int64) +import Data.List (sort) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as T +import Data.Time (UTCTime (..), fromGregorian) +import Data.UUID qualified as UUID +import Effectful (runEff) +import Kiroku.Store +import Kiroku.Test.Postgres (withMigratedTestDatabase, withSharedMigratedPostgres) +import Shibuya.Adapter.Kiroku (KirokuAdapterConfig (..), kirokuAdapter) +import Shibuya.Adapter.Kiroku.Convert (toEnvelope) +import Shibuya.App ( + ProcessorId (..), + SupervisionStrategy (..), + getAppMetrics, + mkProcessor, + runApp, + stopApp, + stopAppGracefully, + ) +import Shibuya.App qualified as Shibuya +import Shibuya.Core.Ack (AckDecision (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Core.Types (Envelope (..)) +import Shibuya.Runner.Metrics (ProcessorState (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import Test.Hspec + +main :: IO () +main = withSharedMigratedPostgres $ hspec $ do + describe "toEnvelope" $ do + it "copies W3C trace metadata into Shibuya trace headers" $ do + let traceparent :: Text + traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + tracestate :: Text + tracestate = "rojo=00f067aa0ba902b7" + Envelope{traceContext} = + toEnvelope + ( makeRecordedEvent + ( Just $ + Aeson.object + [ "traceparent" Aeson..= traceparent + , "tracestate" Aeson..= tracestate + , "other" Aeson..= ("preserved" :: Text) + ] + ) + ) + + traceContext + `shouldBe` Just + [ ("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01") + , ("tracestate", "rojo=00f067aa0ba902b7") + ] + + it "omits trace headers when traceparent is absent or not a string" $ do + let Envelope{traceContext = missingTraceparent} = + toEnvelope (makeRecordedEvent (Just (Aeson.object ["tracestate" Aeson..= ("state" :: Text)]))) + Envelope{traceContext = nonStringTraceparent} = + toEnvelope (makeRecordedEvent (Just (Aeson.object ["traceparent" Aeson..= Aeson.Number 1]))) + + missingTraceparent `shouldBe` Nothing + nonStringTraceparent `shouldBe` Nothing + + around withTestStore $ do + describe "kirokuAdapter" $ do + it "delivers catch-up events through Shibuya pipeline" $ \store -> do + let events = map (\i -> makeEvent ("CU" <> T.pack (show i)) (Aeson.object [])) [1 .. 10 :: Int] + Right _ <- runStoreIO store $ appendToStream (StreamName "shibuya-catchup-1") NoStream events + threadDelay 200_000 + + ref <- newIORef ([] :: [RecordedEvent]) + countVar <- newTVarIO (0 :: Int) + + runEff $ runTracingNoop $ do + adapter <- + kirokuAdapter + store + KirokuAdapterConfig + { subscriptionName = SubscriptionName "shibuya-catchup-test" + , subscriptionTarget = AllStreams + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + + let handler ingested = do + liftIO $ do + modifyIORef' ref (envelopePayload ingested :) + atomically $ do + c <- readTVar countVar + writeTVar countVar (c + 1) + pure AckOk + + res <- runApp IgnoreFailures 100 [(ProcessorId "catchup", mkProcessor adapter handler)] + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ waitForCount countVar 10 10_000_000 + stopApp appHandle + + collected <- readIORef ref + length collected `shouldBe` 10 + + it "delivers live events through Shibuya pipeline" $ \store -> do + ref <- newIORef ([] :: [RecordedEvent]) + countVar <- newTVarIO (0 :: Int) + + runEff $ runTracingNoop $ do + adapter <- + kirokuAdapter + store + KirokuAdapterConfig + { subscriptionName = SubscriptionName "shibuya-live-test" + , subscriptionTarget = AllStreams + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + + let handler ingested = do + liftIO $ do + modifyIORef' ref (envelopePayload ingested :) + atomically $ do + c <- readTVar countVar + writeTVar countVar (c + 1) + pure AckOk + + res <- runApp IgnoreFailures 100 [(ProcessorId "live", mkProcessor adapter handler)] + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ threadDelay 200_000 + liftIO $ do + let events = map (\i -> makeEvent ("Live" <> T.pack (show i)) (Aeson.object [])) [1 .. 5 :: Int] + Right _ <- runStoreIO store $ appendToStream (StreamName "shibuya-live-1") NoStream events + pure () + liftIO $ waitForCount countVar 5 10_000_000 + stopApp appHandle + + collected <- readIORef ref + length collected `shouldBe` 5 + + it "runs multiple category subscriptions concurrently" $ \store -> do + Right _ <- runStoreIO store $ appendToStream (StreamName "orders-s1") NoStream [makeEvent "OrderCreated" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "orders-s2") NoStream [makeEvent "OrderShipped" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "payments-s1") NoStream [makeEvent "PaymentReceived" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "inventory-s1") NoStream [makeEvent "StockUpdated" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "inventory-s2") NoStream [makeEvent "StockDepleted" (Aeson.object [])] + threadDelay 200_000 + + ordersRef <- newIORef ([] :: [RecordedEvent]) + paymentsRef <- newIORef ([] :: [RecordedEvent]) + inventoryRef <- newIORef ([] :: [RecordedEvent]) + ordersCount <- newTVarIO (0 :: Int) + paymentsCount <- newTVarIO (0 :: Int) + inventoryCount <- newTVarIO (0 :: Int) + + runEff $ runTracingNoop $ do + let mkCatAdapter nm cat = + kirokuAdapter + store + KirokuAdapterConfig + { subscriptionName = SubscriptionName nm + , subscriptionTarget = Category (CategoryName cat) + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + let mkHandler ref' cVar ingested = do + liftIO $ do + modifyIORef' ref' (envelopePayload ingested :) + atomically $ do + c <- readTVar cVar + writeTVar cVar (c + 1) + pure AckOk + + ordersAdapter <- mkCatAdapter "orders-proj" "orders" + paymentsAdapter <- mkCatAdapter "payments-proj" "payments" + inventoryAdapter <- mkCatAdapter "inventory-proj" "inventory" + + res <- + runApp + IgnoreFailures + 100 + [ (ProcessorId "orders", mkProcessor ordersAdapter (mkHandler ordersRef ordersCount)) + , (ProcessorId "payments", mkProcessor paymentsAdapter (mkHandler paymentsRef paymentsCount)) + , (ProcessorId "inventory", mkProcessor inventoryAdapter (mkHandler inventoryRef inventoryCount)) + ] + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ do + waitForCount ordersCount 2 10_000_000 + waitForCount paymentsCount 1 10_000_000 + waitForCount inventoryCount 2 10_000_000 + stopApp appHandle + + ordersCollected <- readIORef ordersRef + paymentsCollected <- readIORef paymentsRef + inventoryCollected <- readIORef inventoryRef + length ordersCollected `shouldBe` 2 + length paymentsCollected `shouldBe` 1 + length inventoryCollected `shouldBe` 2 + + it "isolates a failing subscription from healthy ones" $ \store -> do + Right _ <- runStoreIO store $ appendToStream (StreamName "good-s1") NoStream (map (\i -> makeEvent ("Good" <> T.pack (show i)) (Aeson.object [])) [1 .. 3 :: Int]) + Right _ <- runStoreIO store $ appendToStream (StreamName "bad-s1") NoStream [makeEvent "Bad1" (Aeson.object []), makeEvent "Bad2" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "other-s1") NoStream [makeEvent "Other1" (Aeson.object []), makeEvent "Other2" (Aeson.object [])] + threadDelay 200_000 + + goodRef <- newIORef ([] :: [RecordedEvent]) + otherRef <- newIORef ([] :: [RecordedEvent]) + goodCount <- newTVarIO (0 :: Int) + otherCount <- newTVarIO (0 :: Int) + + runEff $ runTracingNoop $ do + goodAdapter <- + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "good-proj" + , subscriptionTarget = Category (CategoryName "good") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + badAdapter <- + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "bad-proj" + , subscriptionTarget = Category (CategoryName "bad") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + otherAdapter <- + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "other-proj" + , subscriptionTarget = Category (CategoryName "other") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + + let goodHandler ingested = do + liftIO $ do + modifyIORef' goodRef (envelopePayload ingested :) + atomically $ do + c <- readTVar goodCount + writeTVar goodCount (c + 1) + pure AckOk + + let badHandler _ingested = do + liftIO $ error "handler crash!" + + let otherHandler ingested = do + liftIO $ do + modifyIORef' otherRef (envelopePayload ingested :) + atomically $ do + c <- readTVar otherCount + writeTVar otherCount (c + 1) + pure AckOk + + res <- + runApp + IgnoreFailures + 100 + [ (ProcessorId "good", mkProcessor goodAdapter goodHandler) + , (ProcessorId "bad", mkProcessor badAdapter badHandler) + , (ProcessorId "other", mkProcessor otherAdapter otherHandler) + ] + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ do + waitForCount goodCount 3 10_000_000 + waitForCount otherCount 2 10_000_000 + threadDelay 500_000 + metrics <- getAppMetrics appHandle + case Map.lookup (ProcessorId "bad") metrics of + Just m -> case m ^. #state of + Failed _ _ -> pure () + other -> liftIO $ expectationFailure ("Expected Failed state, got: " <> show other) + Nothing -> liftIO $ expectationFailure "No metrics for bad processor" + stopApp appHandle + + goodCollected <- readIORef goodRef + otherCollected <- readIORef otherRef + length goodCollected `shouldBe` 3 + length otherCollected `shouldBe` 2 + + it "shuts down all subscriptions coordinately" $ \store -> do + Right _ <- runStoreIO store $ appendToStream (StreamName "shut-a") NoStream [makeEvent "A" (Aeson.object [])] + Right _ <- runStoreIO store $ appendToStream (StreamName "shut-b") NoStream [makeEvent "B" (Aeson.object [])] + threadDelay 200_000 + + countA <- newTVarIO (0 :: Int) + countB <- newTVarIO (0 :: Int) + + runEff $ runTracingNoop $ do + adapterA <- + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "shut-a-proj" + , subscriptionTarget = Category (CategoryName "shut") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + adapterB <- + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "shut-b-proj" + , subscriptionTarget = Category (CategoryName "shut") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Nothing + } + + let handlerA _ingested = do + liftIO $ atomically $ do + c <- readTVar countA + writeTVar countA (c + 1) + pure AckOk + + let handlerB _ingested = do + liftIO $ atomically $ do + c <- readTVar countB + writeTVar countB (c + 1) + pure AckOk + + res <- + runApp + IgnoreFailures + 100 + [ (ProcessorId "a", mkProcessor adapterA handlerA) + , (ProcessorId "b", mkProcessor adapterB handlerB) + ] + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ do + waitForCount countA 2 10_000_000 + waitForCount countB 2 10_000_000 + drained <- stopAppGracefully Shibuya.defaultShutdownConfig appHandle + liftIO $ drained `shouldBe` True + + describe "consumer groups" $ do + it "four-member group delivers a disjoint, complete partition of the stream" $ \store -> do + -- 20 streams × 2 events = 40 events, global positions 1..40, category "cg". + let streams = ["cg-" <> T.pack (show i) | i <- [1 .. 20 :: Int]] + mapM_ + ( \sn -> do + let evs = [makeEvent ("EV" <> T.pack (show k)) (Aeson.object []) | k <- [1 .. 2 :: Int]] + Right _ <- runStoreIO store $ appendToStream (StreamName sn) NoStream evs + pure () + ) + streams + threadDelay 200_000 + + -- One IORef + one TVar counter per member. + refs <- mapM (const (newIORef ([] :: [RecordedEvent]))) [0 .. 3 :: Int] + cvars <- mapM (const (newTVarIO (0 :: Int))) [0 .. 3 :: Int] + + runEff $ runTracingNoop $ do + -- Build one adapter per member index; each carries the same + -- subscriptionName and a distinct member of a size-4 group. + adapters <- + mapM + ( \m -> + kirokuAdapter store $ + KirokuAdapterConfig + { subscriptionName = SubscriptionName "cg-shibuya-group" + , subscriptionTarget = Category (CategoryName "cg") + , batchSize = 100 + , bufferSize = 256 + , consumerGroup = Just (ConsumerGroup{member = m, size = 4}) + } + ) + [0, 1, 2, 3] + + let mkHandler ref' cvar ingested = do + liftIO $ do + modifyIORef' ref' (envelopePayload ingested :) + atomically $ do + c <- readTVar cvar + writeTVar cvar (c + 1) + pure AckOk + + processors = + [ ( ProcessorId ("cg-member-" <> T.pack (show m)) + , mkProcessor (adapters !! m) (mkHandler (refs !! m) (cvars !! m)) + ) + | m <- [0 .. 3 :: Int] + ] + + res <- runApp IgnoreFailures 100 processors + case res of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) + Right appHandle -> do + liftIO $ waitForTotal cvars 40 15_000_000 + stopApp appHandle + + collected <- mapM readIORef refs + -- Disjoint + complete: the union of every member's delivered global + -- positions is exactly [1..40] — no event dropped, none delivered twice. + let allPositions = sort (concatMap (map globalPos) collected) + allPositions `shouldBe` [1 .. 40] + -- No member starved: 20 streams over 4 members gives each ≥ 1 stream. + mapM_ (\c -> length c `shouldSatisfy` (>= 1)) collected + +-- Helpers + +envelopePayload :: Ingested es RecordedEvent -> RecordedEvent +envelopePayload ing = let Ingested{envelope = env} = ing in env ^. #payload + +-- | The raw global position of a recorded event (for disjoint/complete checks). +globalPos :: RecordedEvent -> Int64 +globalPos e = case e ^. #globalPosition of GlobalPosition p -> p + +-- | Wait until the sum of all TVar counts reaches the target or the timeout fires. +waitForTotal :: [STM.TVar Int] -> Int -> Int -> IO () +waitForTotal vars target timeoutMicros = do + timeoutVar <- registerDelay timeoutMicros + result <- + atomically $ + ( do + total <- sum <$> mapM readTVar vars + STM.check (total >= target) + pure True + ) + `STM.orElse` ( do + t <- readTVar timeoutVar + STM.check t + pure False + ) + unless result $ do + actual <- atomically $ sum <$> mapM readTVar vars + expectationFailure ("Timed out waiting for total " <> show target <> ", got " <> show actual) + +waitForCount :: STM.TVar Int -> Int -> Int -> IO () +waitForCount countVar target timeoutMicros = do + timeoutVar <- registerDelay timeoutMicros + result <- + atomically $ + ( do + c <- readTVar countVar + STM.check (c >= target) + pure True + ) + `STM.orElse` ( do + t <- readTVar timeoutVar + STM.check t + pure False + ) + unless result $ do + actual <- atomically $ readTVar countVar + expectationFailure ("Timed out waiting for count " <> show target <> ", got " <> show actual) + +makeEvent :: Text -> Value -> EventData +makeEvent typ p = + EventData + { eventId = Nothing + , eventType = EventType typ + , payload = p + , metadata = Nothing + , causationId = Nothing + , correlationId = Nothing + } + +makeRecordedEvent :: Maybe Value -> RecordedEvent +makeRecordedEvent meta = + RecordedEvent + { eventId = EventId UUID.nil + , eventType = EventType "TraceEvent" + , streamVersion = StreamVersion 1 + , globalPosition = GlobalPosition 1 + , originalStreamId = StreamId 1 + , originalVersion = StreamVersion 1 + , payload = Aeson.object [] + , metadata = meta + , causationId = Nothing + , correlationId = Nothing + , createdAt = UTCTime (fromGregorian 2026 5 16) 0 + } + +withTestStore :: (KirokuStore -> IO ()) -> IO () +withTestStore action = + withMigratedTestDatabase $ \connStr -> + withStore (defaultConnectionSettings connStr) action