shibuya-kafka-adapter 0.9.1.0 → 0.12.0.0
raw patch · 12 files changed
+870/−111 lines, 12 filesdep ~shibuya-corePVP ok
version bump matches the API change (PVP)
Dependency ranges changed: shibuya-core
API changes (from Hackage documentation)
- Shibuya.Adapter.Kafka.Internal: dropStaleRecords :: forall (es :: [Effect]). IOE :> es => KafkaAdapterState -> Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString))) -> Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))
+ Shibuya.Adapter.Kafka.Internal: Polled :: !Word64 -> !ConsumerRecord (Maybe ByteString) (Maybe ByteString) -> Polled
+ Shibuya.Adapter.Kafka.Internal: [ackSignal] :: KafkaAdapterState -> !TVar Word64
+ Shibuya.Adapter.Kafka.Internal: [pollSerial] :: Polled -> !Word64
+ Shibuya.Adapter.Kafka.Internal: [record] :: Polled -> !ConsumerRecord (Maybe ByteString) (Maybe ByteString)
+ Shibuya.Adapter.Kafka.Internal: data Polled
+ Shibuya.Adapter.Kafka.Internal: gateRecords :: forall (es :: [Effect]). IOE :> es => KafkaAdapterState -> Stream (Eff es) (Either KafkaError Polled) -> Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))
- Shibuya.Adapter.Kafka.Internal: KafkaAdapterState :: !TVar Bool -> !IORef AckState -> !IORef (Maybe KafkaError) -> !MVar () -> KafkaAdapterState
+ Shibuya.Adapter.Kafka.Internal: KafkaAdapterState :: !TVar Bool -> !IORef AckState -> !TVar Word64 -> !IORef (Maybe KafkaError) -> !MVar () -> KafkaAdapterState
- Shibuya.Adapter.Kafka.Internal: kafkaSource :: forall (es :: [Effect]). (KafkaConsumer :> es, Error KafkaError :> es, IOE :> es) => KafkaAdapterState -> KafkaAdapterConfig -> Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))
+ Shibuya.Adapter.Kafka.Internal: kafkaSource :: forall (es :: [Effect]). (KafkaConsumer :> es, Error KafkaError :> es, IOE :> es) => KafkaAdapterState -> KafkaAdapterConfig -> Stream (Eff es) (Either KafkaError Polled)
Files
- CHANGELOG.md +28/−0
- README.md +8/−2
- shibuya-kafka-adapter.cabal +6/−3
- src/Shibuya/Adapter/Kafka.hs +12/−2
- src/Shibuya/Adapter/Kafka/Internal.hs +193/−53
- test/Kafka/ScriptedLog.hs +103/−0
- test/Kafka/TestEnv.hs +24/−10
- test/Main.hs +4/−0
- test/Shibuya/Adapter/Kafka/AckHandleTest.hs +11/−4
- test/Shibuya/Adapter/Kafka/DeliveryStatusTest.hs +102/−0
- test/Shibuya/Adapter/Kafka/IntegrationTest.hs +110/−37
- test/Shibuya/Adapter/Kafka/SourceGateTest.hs +269/−0
CHANGELOG.md view
@@ -1,5 +1,33 @@ # Changelog +## 0.12.0.0 — 2026-10-26++### Breaking Changes++- Require the published `shibuya-core ^>=0.12.0.0` across the library, tests,+ examples, and benchmarks. Its delivery-status contract lets the runner skip+ superseded deliveries. This also adopts the extended application configuration,+ processor metrics, and paused-state API introduced in core 0.11.0.0 from+ `mori://shinzui/shibuya`.+- In the exposed internal module, replace `dropStaleRecords` with `gateRecords`;+ `kafkaSource` now yields seek-stamped `Polled` records.++### Bug Fixes++- Under serial single-message processing, skip buffered successors before+ handler entry until the retried record resolves. Same-batch work already+ admitted together and later revocation races retain their documented limits.+- Retain replayed successors behind a pending retry so successful processing+ advances stored offsets to the log end without a restart. Missing retry offsets+ resume at the next available record with a warning.++### Other Changes++- Keep all three repository packages on the shared `0.12.0.0` version line.+- Update the Nix library build to use the published Shibuya core 0.12.0.0.+- Run integration tests against the project-local Redpanda cluster, with+ `KAFKA_TEST_BROKER` available to select another broker.+ ## 0.9.1.0 — 2026-09-21 ### Bug Fixes
README.md view
@@ -10,15 +10,17 @@ Offset reset policy belongs to the `Subscription` passed to `runKafkaConsumer`, for example `topics [TopicName "orders"] <> offsetReset Earliest`. `KafkaAdapterConfig.topics` is metadata used for adapter naming and a construction-time warning if it differs from the live subscription. -`AckRetry` seeks the partition back to the failed message and does not store that message's offset. `AckDeadLetter` stores the offset so the consumer group moves on, prints `[shibuya-kafka-adapter] WARNING: dead-lettered message DROPPED` to stderr, and makes the message unrecoverable from the group's committed position. This adapter does not include a DLQ producer.+`AckRetry` seeks the partition back to the failed message and does not store that message's offset. Records polled before the seek are discarded. Replayed successors wait behind the retried record until it is resolved; if retention or compaction removed that offset, the adapter warns, abandons the missing record, and resumes at the next available offset. `AckDeadLetter` stores the offset so the consumer group moves on, prints `[shibuya-kafka-adapter] WARNING: dead-lettered message DROPPED` to stderr, and makes the message unrecoverable from the group's committed position. This adapter does not include a DLQ producer. +Under serial single-message processing, already-buffered successors are skipped before handler entry and redelivered after the retried record resolves. This preserves partition handler order, subject to the missing-offset recovery policy. A retry cannot undo successors already admitted in the same batch; later superseded batches are skipped. Concurrent processing is unsupported and can admit a successor before the retry is recorded.+ Kafka does not expose a per-message delivery-attempt counter through this consumer API, so `Envelope.attempt` is always `Nothing`. Handlers that need bounded retries must use their own store or return `AckHalt` to stop the stream. `AckHalt` pauses the originating partition and stops the processor. Polling then stops; after `max.poll.interval.ms` (librdkafka default: 300000 ms, or 5 minutes) the broker may evict this consumer and rebalance the partition to another group member. A single-member group stalls until restart. On shutdown, the adapter commits offsets stored so far and signals the source stream to stop. Let the surrounding `runKafkaConsumer` scope end normally after stopping the app so the consumer close path can flush offsets stored during the drain window. -For rebalance visibility, create state with `newKafkaAdapterState`, install `Kafka.Consumer.setCallback (Kafka.Consumer.rebalanceCallback (kafkaRebalanceHandler state))` before consumer creation, and pass the same state to `kafkaAdapterWith`. The helper logs rebalance events and clears retry barriers for revoked partitions; in-flight fencing for cooperative rebalances is out of scope.+For rebalance visibility, create state with `newKafkaAdapterState`, install `Kafka.Consumer.setCallback (Kafka.Consumer.rebalanceCallback (kafkaRebalanceHandler state))` before consumer creation, and pass the same state to `kafkaAdapterWith`. The helper logs rebalance events and fences deliveries from revoked assignments. Deliveries already revoked when their status is checked are skipped before handler entry. A revocation after that admission check cannot cancel an admitted handler. ## Packages @@ -62,3 +64,7 @@ ## License MIT. See package cabal files for details.+++The adapter requires the published `shibuya-core ^>=0.12.0.0`, which provides+the delivery-status API used to skip superseded deliveries before handler entry.
shibuya-kafka-adapter.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.12 name: shibuya-kafka-adapter-version: 0.9.1.0+version: 0.12.0.0 synopsis: Kafka adapter for the Shibuya queue processing framework description: A Shibuya adapter that integrates with Apache Kafka via kafka-effectful@@ -67,7 +67,7 @@ hw-kafka-client >=5.3 && <6, hw-kafka-streamly ^>=0.2, kafka-effectful ^>=0.3.1.0,- shibuya-core ^>=0.10.0.0,+ shibuya-core ^>=0.12.0.0, stm ^>=2.5, streamly ^>=0.11, streamly-core ^>=0.3,@@ -101,11 +101,14 @@ QuasiQuotes other-modules:+ Kafka.ScriptedLog Kafka.TestEnv Shibuya.Adapter.Kafka.AckHandleTest Shibuya.Adapter.Kafka.AdapterTest Shibuya.Adapter.Kafka.ConvertTest+ Shibuya.Adapter.Kafka.DeliveryStatusTest Shibuya.Adapter.Kafka.IntegrationTest+ Shibuya.Adapter.Kafka.SourceGateTest build-depends: async ^>=2.2,@@ -119,7 +122,7 @@ kafka-effectful, process, random,- shibuya-core ^>=0.10.0.0,+ shibuya-core ^>=0.12.0.0, shibuya-kafka-adapter, stm, streamly,
src/Shibuya/Adapter/Kafka.hs view
@@ -34,11 +34,18 @@ -- 4. On @AckRetry@, the offset is not stored. The adapter seeks the partition -- back to the earliest unresolved delivery so Kafka can redeliver it. Later -- buffered callbacks cannot replace or commit past that recovery boundary.+-- Pre-seek records are discarded; replayed successors wait until the retry+-- resolves. A missing retry offset is abandoned with a warning.+-- A serial single-message runner skips already-buffered successors before+-- handler entry and redelivers them after the retry. Same-batch successors+-- have already been admitted together; concurrent admission has no such+-- ordering guarantee. -- 5. On @AckDeadLetter@, the offset is stored after a loud stderr warning. -- 6. On @AckHalt@, the partition is paused and offset is not stored. -- -- == Serial Operation Required --+-- Under serial single-message processing a retried record blocks its successors. -- This adapter must be run with serial message processing. librdkafka stores the -- highest offset per partition without gap tracking, so concurrent finalization -- can commit past an earlier message that failed, halted, or requested retry.@@ -83,6 +90,9 @@ -- -- == Rebalance Callback Helper --+-- With the helper installed, deliveries already revoked at the admission check+-- are skipped. Revocation after the check cannot cancel an admitted handler.+-- -- 'kafkaRebalanceHandler' is optional. Install it with -- @Kafka.Consumer.setCallback (Kafka.Consumer.rebalanceCallback (kafkaRebalanceHandler state))@ -- before creating the consumer when you want stderr visibility into assignment@@ -129,7 +139,7 @@ import Kafka.Types (BatchSize (..), BrokerAddress (..), KafkaError (..), Timeout (..), TopicName (..)) import Shibuya.Adapter (Adapter (..)) import Shibuya.Adapter.Kafka.Config (KafkaAdapterConfig (..), defaultConfig)-import Shibuya.Adapter.Kafka.Internal (KafkaAcknowledgementException (..), KafkaAdapterState (..), dropStaleRecords, ingestedStream, kafkaSource, markPartitionsAssigned, markPartitionsRevoked, mkIngested, newKafkaAdapterState, withConsumerLock)+import Shibuya.Adapter.Kafka.Internal (KafkaAcknowledgementException (..), KafkaAdapterState (..), gateRecords, ingestedStream, kafkaSource, markPartitionsAssigned, markPartitionsRevoked, mkIngested, newKafkaAdapterState, withConsumerLock) import System.IO (hPutStrLn, stderr) -- | Create a Kafka adapter with the given configuration.@@ -172,7 +182,7 @@ warnOnSubscriptionMismatch config let messageSource = ingestedStream (mkIngested state config) $- dropStaleRecords state $+ gateRecords state $ kafkaSource state config pure Adapter
src/Shibuya/Adapter/Kafka/Internal.hs view
@@ -11,7 +11,8 @@ -- * Stream Construction kafkaSource,- dropStaleRecords,+ Polled (..),+ gateRecords, ingestedStream, -- * Ingested Construction@@ -24,7 +25,8 @@ import Control.Concurrent (threadDelay) import Control.Concurrent.MVar (MVar, newMVar, putMVar, takeMVar)-import Control.Concurrent.STM (TVar, newTVarIO, readTVarIO)+import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVar, readTVarIO, retry)+import Control.Exception qualified as IOExc import Control.Monad (unless, when) import Data.ByteString (ByteString) import Data.Function ((&))@@ -68,18 +70,24 @@ data RetryBarrier = RetryBarrier { offset :: !Offset,- retryToken :: !DeliveryToken+ retryToken :: !DeliveryToken,+ replayed :: !Bool } data PartitionAckState = PartitionAckState { barrier :: !(Maybe RetryBarrier), validFromToken :: !DeliveryToken, assignmentGeneration :: !AssignmentGeneration,- assigned :: !Bool+ assigned :: !Bool,+ retryRequests :: !Word64,+ seeksCompleted :: !Word64,+ seekedAt :: !Word64,+ halted :: !Bool } data AckState = AckState { nextDeliveryToken :: !Word64,+ seekSerial :: !Word64, partitions :: !(Map PartitionKey PartitionAckState) } @@ -104,16 +112,21 @@ { barrier = Nothing, validFromToken = DeliveryToken 0, assignmentGeneration = AssignmentGeneration 0,- assigned = True+ assigned = True,+ retryRequests = 0,+ seeksCompleted = 0,+ seekedAt = 0,+ halted = False } initialAckState :: AckState-initialAckState = AckState {nextDeliveryToken = 1, partitions = Map.empty}+initialAckState = AckState {nextDeliveryToken = 1, seekSerial = 0, partitions = Map.empty} -- | Mutable state shared by the source stream and ack handles. data KafkaAdapterState = KafkaAdapterState { shutdownVar :: !(TVar Bool), ackState :: !(IORef AckState),+ ackSignal :: !(TVar Word64), fatalError :: !(IORef (Maybe KafkaError)), -- | Serializes every librdkafka consumer operation. Under -- 'Shibuya.App.runApp' the consumer handle is shared between the ingester@@ -134,6 +147,7 @@ KafkaAdapterState <$> newTVarIO False <*> newIORef initialAckState+ <*> newTVarIO 0 <*> newIORef Nothing <*> newMVar () @@ -150,10 +164,11 @@ markPartitionsRevoked state = setPartitionsAssigned state False setPartitionsAssigned :: KafkaAdapterState -> Bool -> [PartitionKey] -> IO ()-setPartitionsAssigned state isAssigned keys =+setPartitionsAssigned state isAssigned keys = IOExc.mask_ $ do atomicModifyIORef' state.ackState $ \ackState' -> let partitions' = foldr (Map.alter (Just . transition . maybe initialPartitionAckState id)) ackState'.partitions keys in (ackState' {partitions = partitions'}, ())+ signalAckState state where transition partitionState | partitionState.assigned == isAssigned =@@ -165,7 +180,9 @@ { barrier = Nothing, validFromToken = DeliveryToken 0, assignmentGeneration = nextGeneration partitionState.assignmentGeneration,- assigned = isAssigned+ assigned = isAssigned,+ seeksCompleted = partitionState.retryRequests,+ halted = False } nextGeneration (AssignmentGeneration generation') = AssignmentGeneration (generation' + 1)@@ -203,7 +220,8 @@ boundedLockTimeout :: Timeout -> Timeout boundedLockTimeout t = Timeout (min (unTimeout t) maxPollHoldMillis) --- | Create a stream of 'ConsumerRecord's by repeatedly polling the broker.+-- | Poll records and stamp each successful record with the completed seek serial+-- under the consumer lock. Feed this stream through 'gateRecords'. -- -- Calls 'pollMessageBatch' in a loop under 'consumerLock', preserving errors as -- @Left@ values. Each poll uses a timeout capped at @maxPollHoldMillis@ so the@@ -216,7 +234,7 @@ (KafkaConsumer :> es, Error KafkaError :> es, IOE :> es) => KafkaAdapterState -> KafkaAdapterConfig ->- Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))+ Stream (Eff es) (Either KafkaError Polled) kafkaSource state config = skipNonFatal $ Stream.unfoldrM step ()@@ -232,26 +250,97 @@ if isShutdown then pure Nothing else do- batch <- withConsumerLock state (pollMessageBatch pollT config.batchSize)+ batch <- withConsumerLock state $ do+ records <- pollMessageBatch pollT config.batchSize+ serial <- Effectful.liftIO $ (.seekSerial) <$> readIORef state.ackState+ pure (map (fmap (Polled serial)) records) pure (Just (batch, ())) --- | Drop records already buffered above a pending retry barrier.-dropStaleRecords ::+-- | A record stamped under the consumer lock with the latest completed seek.+data Polled = Polled+ { pollSerial :: !Word64,+ record :: !(ConsumerRecord (Maybe ByteString) (Maybe ByteString))+ }++data Admission = Emit | EmitAfterGap !Offset | Drop | Hold++-- | Drop pre-seek records; retain post-seek successors until their replay resolves.+-- Holds are interruptible and never own the consumer lock. Missing offsets are+-- abandoned with a warning, fencing callbacks from before the gap.+gateRecords :: (IOE :> es) => KafkaAdapterState ->- Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString))) ->+ Stream (Eff es) (Either KafkaError Polled) -> Stream (Eff es) (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))-dropStaleRecords state =- Stream.filterM $ \case- Left _ -> pure True- Right cr -> do- ackState' <- Effectful.liftIO $ readIORef state.ackState- let partitionState = Map.findWithDefault initialPartitionAckState (partitionKey cr) ackState'.partitions- pure $- partitionState.assigned && case partitionState.barrier of- Nothing -> True- Just retryBarrier -> cr.crOffset <= retryBarrier.offset+gateRecords state =+ Stream.mapMaybeM $ \case+ Left err -> pure (Just (Left err))+ Right polled -> do+ admitted <- Effectful.liftIO $ admitRecord state polled+ pure $ if admitted then Just (Right polled.record) else Nothing +admitRecord :: KafkaAdapterState -> Polled -> IO Bool+admitRecord state polled = do+ ackState' <- readIORef state.ackState+ let partitionState = Map.findWithDefault initialPartitionAckState key ackState'.partitions+ case (partitionState.barrier, isFresh partitionState) of+ (Nothing, True) -> pure True+ _ -> wait+ where+ cr = polled.record+ key = partitionKey cr++ isFresh partitionState =+ partitionState.assigned+ && partitionState.retryRequests == partitionState.seeksCompleted+ && polled.pollSerial >= partitionState.seekedAt++ wait = do+ version <- readTVarIO state.ackSignal+ admission <- atomicModifyIORef' state.ackState decide+ case admission of+ Emit -> pure True+ EmitAfterGap missing -> do+ hPutStrLn stderr $+ "[shibuya-kafka-adapter] WARNING: retried record is no longer in the log; resuming after it: "+ <> show (cr.crTopic, cr.crPartition, missing)+ pure True+ Drop -> pure False+ Hold -> do+ mbFatal <- readIORef state.fatalError+ case mbFatal of+ Just _ -> pure False+ Nothing -> do+ stop <- atomically $ do+ isShutdown <- readTVar state.shutdownVar+ current <- readTVar state.ackSignal+ if isShutdown+ then pure True+ else if current /= version then pure False else retry+ if stop then pure False else wait++ decide ackState' =+ let partitionState = Map.findWithDefault initialPartitionAckState key ackState'.partitions+ update partitionState' = ackState' {partitions = Map.insert key partitionState' ackState'.partitions}+ in if not (isFresh partitionState)+ then (ackState', Drop)+ else case partitionState.barrier of+ Nothing -> (ackState', Emit)+ Just retryBarrier+ | cr.crOffset < retryBarrier.offset -> (ackState', Emit)+ | cr.crOffset == retryBarrier.offset ->+ (update partitionState {barrier = Just retryBarrier {replayed = True}}, Emit)+ | partitionState.halted -> (ackState', Drop)+ | retryBarrier.replayed -> (ackState', Hold)+ | otherwise ->+ ( update+ partitionState+ { barrier = Nothing,+ validFromToken = max partitionState.validFromToken (DeliveryToken ackState'.nextDeliveryToken)+ },+ EmitAfterGap retryBarrier.offset+ )+ -- | Create an 'AckHandle' for a single 'ConsumerRecord'. -- -- Maps 'AckDecision' to Kafka operations:@@ -268,6 +357,16 @@ Eff es (AckHandle es) mkAckHandle state config cr = do attempt <- Effectful.liftIO $ registerDelivery state cr+ ackHandleFor state config attempt cr++ackHandleFor ::+ (KafkaConsumer :> es, Error KafkaError :> es, IOE :> es) =>+ KafkaAdapterState ->+ KafkaAdapterConfig ->+ DeliveryAttempt ->+ ConsumerRecord (Maybe ByteString) (Maybe ByteString) ->+ Eff es (AckHandle es)+ackHandleFor state config attempt cr = do finalizerLock <- Effectful.liftIO $ newMVar () completed <- Effectful.liftIO $ newIORef False pure $ AckHandle $ \decision ->@@ -315,9 +414,9 @@ mbTarget <- Effectful.liftIO $ recordRetry state attempt case mbTarget of Nothing -> pure ()- Just target ->+ Just (target, request) -> ackAttempt state $- withConsumerLock state $+ withConsumerLock state $ do seekPartitions [ TopicPartition { tpTopicName = cr.crTopic,@@ -326,6 +425,7 @@ } ] (boundedLockTimeout config.pollTimeout)+ Effectful.liftIO $ markSeekCompleted state attempt.partition request AckDeadLetter reason -> storeGuarded state attempt cr $ Effectful.liftIO $@@ -335,7 +435,7 @@ <> " reason=" <> Text.unpack (renderDeadLetterReason reason) AckHalt _ -> do- isAccepted <- Effectful.liftIO $ deliveryIsAccepted state attempt+ isAccepted <- Effectful.liftIO $ markHalted state attempt when isAccepted $ ackAttempt state (withConsumerLock state (pausePartitions [(cr.crTopic, cr.crPartition)])) @@ -360,10 +460,11 @@ go (attempt + 1) recordFatalError :: KafkaAdapterState -> KafkaError -> IO ()-recordFatalError state err =+recordFatalError state err = IOExc.mask_ $ do atomicModifyIORef' state.fatalError $ \case Just existing -> (Just existing, ()) Nothing -> (Just err, ())+ signalAckState state storeGuarded :: (KafkaConsumer :> es, Error KafkaError :> es, IOE :> es) =>@@ -378,36 +479,41 @@ beforeStore ackAttempt state (withConsumerLock state (storeOffsetMessage cr)) -recordRetry :: KafkaAdapterState -> DeliveryAttempt -> IO (Maybe Offset)-recordRetry state attempt =- atomicModifyIORef' state.ackState $ \ackState' ->+recordRetry :: KafkaAdapterState -> DeliveryAttempt -> IO (Maybe (Offset, Word64))+recordRetry state attempt = IOExc.mask_ $ do+ result <- atomicModifyIORef' state.ackState $ \ackState' -> let partitionState = Map.findWithDefault initialPartitionAckState attempt.partition ackState'.partitions in if not (attemptIsAccepted partitionState attempt) then (ackState', Nothing) else let retryBarrier = case partitionState.barrier of- Nothing -> RetryBarrier attempt.recordOffset attempt.token+ Nothing -> RetryBarrier attempt.recordOffset attempt.token False Just existing- | attempt.recordOffset < existing.offset -> RetryBarrier attempt.recordOffset attempt.token+ | attempt.recordOffset < existing.offset -> RetryBarrier attempt.recordOffset attempt.token False | attempt.recordOffset == existing.offset && attempt.token > existing.retryToken -> existing {retryToken = attempt.token} | otherwise -> existing- partitionState' = partitionState {barrier = Just retryBarrier}+ partitionState' = partitionState {barrier = Just retryBarrier {replayed = False}, retryRequests = partitionState.retryRequests + 1} ackState'' = ackState' {partitions = Map.insert attempt.partition partitionState' ackState'.partitions}- in (ackState'', Just retryBarrier.offset)+ in (ackState'', Just (retryBarrier.offset, partitionState'.retryRequests)) + case result of+ Nothing -> pure ()+ Just _ -> signalAckState state+ pure result+ claimStore :: KafkaAdapterState -> DeliveryAttempt -> IO Bool-claimStore state attempt =- atomicModifyIORef' state.ackState $ \ackState' ->+claimStore state attempt = IOExc.mask_ $ do+ (accepted, cleared) <- atomicModifyIORef' state.ackState $ \ackState' -> let partitionState = Map.findWithDefault initialPartitionAckState attempt.partition ackState'.partitions in if not (attemptIsAccepted partitionState attempt)- then (ackState', False)+ then (ackState', (False, False)) else case partitionState.barrier of- Nothing -> (ackState', True)+ Nothing -> (ackState', (True, False)) Just retryBarrier- | attempt.recordOffset > retryBarrier.offset -> (ackState', False)- | attempt.recordOffset < retryBarrier.offset -> (ackState', True)- | attempt.token <= retryBarrier.retryToken -> (ackState', False)+ | attempt.recordOffset > retryBarrier.offset -> (ackState', (False, False))+ | attempt.recordOffset < retryBarrier.offset -> (ackState', (True, False))+ | attempt.token <= retryBarrier.retryToken -> (ackState', (False, False)) | otherwise -> let partitionState' = partitionState@@ -415,14 +521,33 @@ validFromToken = max partitionState.validFromToken attempt.token } ackState'' = ackState' {partitions = Map.insert attempt.partition partitionState' ackState'.partitions}- in (ackState'', True)+ in (ackState'', (True, True)) -deliveryIsAccepted :: KafkaAdapterState -> DeliveryAttempt -> IO Bool-deliveryIsAccepted state attempt = do- ackState' <- readIORef state.ackState- let partitionState = Map.findWithDefault initialPartitionAckState attempt.partition ackState'.partitions- pure $ attemptIsAccepted partitionState attempt+ when cleared $ signalAckState state+ pure accepted +signalAckState :: KafkaAdapterState -> IO ()+signalAckState state = atomically $ modifyTVar' state.ackSignal (+ 1)++markSeekCompleted :: KafkaAdapterState -> PartitionKey -> Word64 -> IO ()+markSeekCompleted state key request = IOExc.mask_ $ do+ atomicModifyIORef' state.ackState $ \ledger ->+ let serial = ledger.seekSerial + 1+ partition = Map.findWithDefault initialPartitionAckState key ledger.partitions+ updated = partition {seekedAt = serial, seeksCompleted = max partition.seeksCompleted request}+ in (ledger {seekSerial = serial, partitions = Map.insert key updated ledger.partitions}, ())+ signalAckState state++markHalted :: KafkaAdapterState -> DeliveryAttempt -> IO Bool+markHalted state attempt = IOExc.mask_ $ do+ accepted <- atomicModifyIORef' state.ackState $ \ledger ->+ let partition = Map.findWithDefault initialPartitionAckState attempt.partition ledger.partitions+ in if attemptIsAccepted partition attempt+ then (ledger {partitions = Map.insert attempt.partition partition {halted = True} ledger.partitions}, True)+ else (ledger, False)+ when accepted $ signalAckState state+ pure accepted+ attemptIsAccepted :: PartitionAckState -> DeliveryAttempt -> Bool attemptIsAccepted partitionState attempt = partitionState.assigned@@ -447,11 +572,26 @@ ConsumerRecord (Maybe ByteString) (Maybe ByteString) -> Eff es (Ingested es (Maybe ByteString)) mkIngested state config cr = do- ackHandle <- mkAckHandle state config cr- pure $- Core.mkIngested- (consumerRecordToEnvelope cr)- ackHandle+ attempt <- Effectful.liftIO $ registerDelivery state cr+ ackHandle <- ackHandleFor state config attempt cr+ pure+ (Core.mkIngested (consumerRecordToEnvelope cr) ackHandle)+ { Core.deliveryStatus = Effectful.liftIO (deliveryStatusOf state attempt)+ }++-- Admission is a single ledger read, independent of consumer-lock ownership.+-- A current answer cannot cancel a later revoke or concurrent retry.+deliveryStatusOf :: KafkaAdapterState -> DeliveryAttempt -> IO Core.DeliveryStatus+deliveryStatusOf state attempt = do+ ledger <- readIORef state.ackState+ let partition = Map.findWithDefault initialPartitionAckState attempt.partition ledger.partitions+ current =+ attemptIsAccepted partition attempt && case partition.barrier of+ Nothing -> True+ Just barrier ->+ attempt.recordOffset < barrier.offset+ || (attempt.recordOffset == barrier.offset && attempt.token > barrier.retryToken)+ pure $ if current then Core.DeliveryCurrent else Core.DeliverySuperseded -- | Transform a poll stream of @Either KafkaError ConsumerRecord@ into a -- stream of 'Ingested'.
+ test/Kafka/ScriptedLog.hs view
@@ -0,0 +1,103 @@+-- | A broker-free 'KafkaConsumer' interpreter backed by an in-memory,+-- single-partition log that supports poll, seek, and offset store.+module Kafka.ScriptedLog+ ( ScriptedLog (..),+ newScriptedLog,+ runScriptedConsumer,+ scriptedTopic,+ )+where++import Control.Concurrent (threadDelay)+import Control.Monad (forM_)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef (IORef, atomicModifyIORef', newIORef)+import Data.Int (Int64)+import Effectful (Eff, IOE, liftIO, (:>))+import Effectful.Dispatch.Dynamic (interpret)+import Kafka.Consumer (RdKafkaRespErrT (..))+import Kafka.Consumer.Types (ConsumerRecord (..), Offset (..), PartitionOffset (..), SubscribedPartitions (..), Timestamp (..), TopicPartition (..))+import Kafka.Effectful.Consumer.Effect (KafkaConsumer (..))+import Kafka.Types (BatchSize (..), KafkaError (..), PartitionId (..), Timeout (..), TopicName (..))++data ScriptedLog = ScriptedLog+ { records :: ![ConsumerRecord (Maybe ByteString) (Maybe ByteString)],+ -- | Index of the next record a poll returns.+ position :: !(IORef Int),+ -- | Offsets passed to @storeOffsetMessage@, oldest first.+ stored :: !(IORef [Int64]),+ -- | Offsets passed to @seekPartitions@, oldest first.+ seeks :: !(IORef [Int64]),+ -- | Timeout in milliseconds of every poll the adapter made, oldest first.+ polls :: !(IORef [Int])+ }++scriptedTopic :: TopicName+scriptedTopic = TopicName "orders"++-- | A log holding offsets @0 .. count - 1@ on partition 0 of 'scriptedTopic'.+newScriptedLog :: Int -> IO ScriptedLog+newScriptedLog count =+ ScriptedLog (map recordAt [0 .. fromIntegral count - 1])+ <$> newIORef 0+ <*> newIORef []+ <*> newIORef []+ <*> newIORef []+ where+ recordAt offset =+ ConsumerRecord+ { crTopic = scriptedTopic,+ crPartition = PartitionId 0,+ crOffset = Offset offset,+ crTimestamp = NoTimestamp,+ crHeaders = mempty,+ crKey = Nothing,+ crValue = Just (BS8.pack (show offset))+ }++runScriptedConsumer :: (IOE :> es) => ScriptedLog -> Eff (KafkaConsumer : es) a -> Eff es a+runScriptedConsumer scripted =+ interpret $ \_env -> \case+ PollMessageBatch (Timeout millis) (BatchSize size) -> liftIO $ do+ atomicModifyIORef' scripted.polls (\done -> (done <> [millis], ()))+ batch <-+ atomicModifyIORef' scripted.position $ \current ->+ let batch = take size (drop current scripted.records)+ in (current + length batch, batch)+ -- A real poll blocks for its timeout when the partition is drained.+ if null batch then threadDelay (millis * 1000) else pure ()+ pure (map Right batch)+ SeekPartitions targets _ -> liftIO $+ forM_ targets $ \target -> case target.tpOffset of+ PartitionOffset offset -> do+ atomicModifyIORef' scripted.position (const (fromIntegral offset, ()))+ atomicModifyIORef' scripted.seeks (\done -> (done <> [offset], ()))+ other -> error ("ScriptedLog: unsupported seek target " <> show other)+ StoreOffsetMessage cr -> liftIO $ atomicModifyIORef' scripted.stored (\done -> (done <> [unOffset cr.crOffset], ()))+ Subscription -> pure [(scriptedTopic, SubscribedPartitionsAll)]+ CommitAllOffsets _ -> pure ()+ PausePartitions _ -> pure ()+ PollMessageEither (Timeout millis) -> liftIO $ do+ atomicModifyIORef' scripted.polls (\done -> (done <> [millis], ()))+ next <-+ atomicModifyIORef' scripted.position $ \current ->+ case drop current scripted.records of+ record : _ -> (current + 1, Just record)+ [] -> (current, Nothing)+ case next of+ Just record -> pure (Right record)+ Nothing -> do+ -- A real poll blocks for its timeout when the partition is drained.+ threadDelay (millis * 1000)+ pure (Left (KafkaResponseError RdKafkaRespErrTimedOut))+ PollMessage _ -> error "ScriptedLog: PollMessage not scripted"+ CommitOffsetMessage _ _ -> error "ScriptedLog: CommitOffsetMessage not scripted"+ CommitPartitionsOffsets _ _ -> error "ScriptedLog: CommitPartitionsOffsets not scripted"+ StoreOffsets _ -> error "ScriptedLog: StoreOffsets not scripted"+ Assign _ -> error "ScriptedLog: Assign not scripted"+ ResumePartitions _ -> error "ScriptedLog: ResumePartitions not scripted"+ Committed _ _ -> error "ScriptedLog: Committed not scripted"+ Position _ -> error "ScriptedLog: Position not scripted"+ Assignment -> error "ScriptedLog: Assignment not scripted"+ AskConsumerHandle -> error "ScriptedLog: AskConsumerHandle not scripted"
test/Kafka/TestEnv.hs view
@@ -6,6 +6,7 @@ -- * Producing produceMessages,+ produceMessagesBounded, produceKeyedMessages, producePartitionMessages, @@ -60,7 +61,8 @@ import Shibuya.Core.Types (Envelope (..)) import Streamly.Data.Fold qualified as Fold import Streamly.Data.Stream qualified as Stream-import System.Process (callCommand)+import System.Environment (lookupEnv)+import System.Process (callProcess) import System.Random (randomRIO) -- | Test environment with isolated topic and group.@@ -75,9 +77,10 @@ withTestEnv :: (TestEnv -> IO a) -> IO a withTestEnv f = do prefix <- randomPrefix+ broker <- maybe "127.0.0.1:9092" Text.pack <$> lookupEnv "KAFKA_TEST_BROKER" let env = TestEnv- { testBroker = BrokerAddress "127.0.0.1:9092",+ { testBroker = BrokerAddress broker, testTopic = TopicName (Text.pack (prefix <> "-topic")), testGroupId = ConsumerGroupId (Text.pack (prefix <> "-group")), testPrefix = prefix@@ -97,12 +100,16 @@ -- | Create a topic with the specified number of partitions via rpk. createTopicWithPartitions :: TestEnv -> Int -> IO () createTopicWithPartitions env n =- callCommand $- "rpk topic create "- <> Text.unpack (unTopicName env.testTopic)- <> " -p "- <> show n- <> " 2>/dev/null || true"+ callProcess+ "rpk"+ [ "topic",+ "create",+ Text.unpack (unTopicName env.testTopic),+ "-p",+ show n,+ "-X",+ "brokers=" <> Text.unpack (unBrokerAddress env.testBroker)+ ] -- | Build consumer properties for testing. mkConsumerProps :: TestEnv -> ConsumerProperties@@ -124,9 +131,16 @@ -- | Produce messages with the given payloads to the test topic. produceMessages :: TestEnv -> [ByteString] -> IO ()-produceMessages env payloads = do+produceMessages env = produceMessagesWith (mkProducerProps env) env++-- Bound delivery and producer close when the broker is unavailable.+produceMessagesBounded :: TestEnv -> [ByteString] -> IO ()+produceMessagesBounded env = produceMessagesWith (mkProducerProps env <> P.sendTimeout (Timeout 5000)) env++produceMessagesWith :: ProducerProperties -> TestEnv -> [ByteString] -> IO ()+produceMessagesWith props env payloads = do result <- runEff . runError @KafkaError $ do- runKafkaProducer (mkProducerProps env) $ do+ runKafkaProducer props $ do forM_ payloads $ \payload -> produceMessage ProducerRecord
test/Main.hs view
@@ -3,7 +3,9 @@ import Shibuya.Adapter.Kafka.AckHandleTest qualified as AckHandleTest import Shibuya.Adapter.Kafka.AdapterTest qualified as AdapterTest import Shibuya.Adapter.Kafka.ConvertTest qualified as ConvertTest+import Shibuya.Adapter.Kafka.DeliveryStatusTest qualified as DeliveryStatusTest import Shibuya.Adapter.Kafka.IntegrationTest qualified as IntegrationTest+import Shibuya.Adapter.Kafka.SourceGateTest qualified as SourceGateTest import Test.Tasty (defaultMain, testGroup) main :: IO ()@@ -14,5 +16,7 @@ [ AckHandleTest.tests, AdapterTest.tests, ConvertTest.tests,+ DeliveryStatusTest.tests,+ SourceGateTest.tests, IntegrationTest.tests ]
test/Shibuya/Adapter/Kafka/AckHandleTest.hs view
@@ -1,4 +1,4 @@-module Shibuya.Adapter.Kafka.AckHandleTest (tests) where+module Shibuya.Adapter.Kafka.AckHandleTest (tests, MockState (..), defaultMockState, runFinalizer, runMockConsumer, recordAt, testConfig, finalizeRecord, assertRight) where import Control.Concurrent.Async qualified as Async import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar)@@ -17,7 +17,7 @@ import Shibuya.Adapter (Adapter (..)) import Shibuya.Adapter.Kafka (kafkaRebalanceHandler) import Shibuya.Adapter.Kafka.Config (KafkaAdapterConfig (..))-import Shibuya.Adapter.Kafka.Internal (KafkaAcknowledgementException (..), KafkaAdapterState (..), ingestedStream, kafkaSource, mkAckHandle, mkIngested, newKafkaAdapterState)+import Shibuya.Adapter.Kafka.Internal (KafkaAcknowledgementException (..), KafkaAdapterState (..), gateRecords, ingestedStream, kafkaSource, mkAckHandle, mkIngested, newKafkaAdapterState) import Shibuya.App (ProcessorId (..), defaultAppConfig, getAppMaster, mkProcessor, runApp, waitApp) import Shibuya.Core.Ack (AckDecision (..), HaltReason (..), RetryDelay (..)) import Shibuya.Core.AckHandle (AckHandle (..))@@ -41,6 +41,7 @@ storeFailuresRemaining :: !Int, pauseFailuresRemaining :: !Int, seekFailuresRemaining :: !Int,+ seekBlock :: !(Maybe (MVar (), MVar ())), storeError :: !KafkaError, pauseError :: !KafkaError, seekError :: !KafkaError@@ -316,7 +317,7 @@ result <- runEff . runErrorNoCallStack @KafkaError . runMockConsumer mock $ Stream.fold Fold.drain $- ingestedStream unreachableBuilder (kafkaSource state testConfig)+ ingestedStream unreachableBuilder (gateRecords state (kafkaSource state testConfig)) assertEqual "source error" (Left err) result runFinalizer ::@@ -336,7 +337,12 @@ interpret $ \_env -> \case StoreOffsetMessage cr -> attemptStore mock cr.crOffset PausePartitions _ -> attemptPause mock- SeekPartitions tps seekTimeout -> recordSeek mock tps seekTimeout+ SeekPartitions tps seekTimeout -> do+ block <- liftIO $ (.seekBlock) <$> readIORef mock+ liftIO $ case block of+ Nothing -> pure ()+ Just (started, release) -> putMVar started () >> takeMVar release+ recordSeek mock tps seekTimeout PollMessage _ -> error "AckHandleTest: PollMessage not exercised" PollMessageBatch _ _ -> error "AckHandleTest: PollMessageBatch not exercised" PollMessageEither _ -> error "AckHandleTest: PollMessageEither not exercised"@@ -447,6 +453,7 @@ storeFailuresRemaining = 0, pauseFailuresRemaining = 0, seekFailuresRemaining = 0,+ seekBlock = Nothing, storeError = KafkaResponseError RdKafkaRespErrTransport, pauseError = KafkaResponseError RdKafkaRespErrTransport, seekError = KafkaResponseError RdKafkaRespErrTransport
+ test/Shibuya/Adapter/Kafka/DeliveryStatusTest.hs view
@@ -0,0 +1,102 @@+module Shibuya.Adapter.Kafka.DeliveryStatusTest (tests) where++import Data.IORef (newIORef)+import Effectful (liftIO)+import Kafka.Consumer.Types (ConsumerRecord (..))+import Kafka.Types (PartitionId (..))+import Shibuya.Adapter.Kafka.AckHandleTest hiding (tests)+import Shibuya.Adapter.Kafka.Internal (Polled (..), gateRecords, markPartitionsAssigned, markPartitionsRevoked, mkIngested, newKafkaAdapterState)+import Shibuya.Core.Ack (AckDecision (..), RetryDelay (..))+import Shibuya.Core.AckHandle (AckHandle (..))+import Shibuya.Core.Ingested (DeliveryStatus (..), Ingested (..))+import Streamly.Data.Fold qualified as Fold+import Streamly.Data.Stream qualified as Stream+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertEqual, testCase)++tests :: TestTree+tests =+ testGroup+ "Delivery status"+ [ testCase "retry supersedes the original and successors, preserving predecessors and other partitions" testBarrier,+ testCase "replayed retry supersedes that attempt again" testRepeated,+ testCase "revoke and reassign fence the old generation" testGeneration,+ testCase "missing-offset recovery supersedes old deliveries" testGap+ ]++testBarrier :: IO ()+testBarrier = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer+ mock+ ( do+ before <- mkIngested state testConfig (recordAt 41)+ original <- mkIngested state testConfig (recordAt 42)+ successor <- mkIngested state testConfig (recordAt 43)+ other <- mkIngested state testConfig (recordAt 44) {crPartition = PartitionId 1}+ mapM_ (\i -> i.deliveryStatus >>= liftIO . assertEqual "initially current" DeliveryCurrent) [before, original, successor, other]+ let AckHandle retry = original.ack+ retry (AckRetry (RetryDelay 0))+ statuses <- mapM (.deliveryStatus) [before, original, successor, other]+ liftIO $ assertEqual "barrier admission" [DeliveryCurrent, DeliverySuperseded, DeliverySuperseded, DeliveryCurrent] statuses+ replay <- mkIngested state testConfig (recordAt 42)+ replay.deliveryStatus >>= liftIO . assertEqual "new barrier attempt current" DeliveryCurrent+ let AckHandle finalize = replay.ack+ finalize AckOk+ freshSuccessor <- mkIngested state testConfig (recordAt 43)+ after <- mapM (.deliveryStatus) [original, successor, replay, freshSuccessor]+ liftIO $ assertEqual "token floor persists after success" [DeliverySuperseded, DeliverySuperseded, DeliveryCurrent, DeliveryCurrent] after+ )+ >>= assertRight++testRepeated :: IO ()+testRepeated = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer+ mock+ ( do+ finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))+ replay <- mkIngested state testConfig (recordAt 42)+ let AckHandle finalize = replay.ack+ finalize (AckRetry (RetryDelay 0))+ replay.deliveryStatus >>= liftIO . assertEqual "retried replay superseded" DeliverySuperseded+ next <- mkIngested state testConfig (recordAt 42)+ next.deliveryStatus >>= liftIO . assertEqual "next replay current" DeliveryCurrent+ )+ >>= assertRight++testGeneration :: IO ()+testGeneration = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer+ mock+ ( do+ old <- mkIngested state testConfig (recordAt 42)+ let key = (recordAt 42).crTopic+ liftIO $ markPartitionsRevoked state [(key, PartitionId 0)]+ old.deliveryStatus >>= liftIO . assertEqual "revoked attempt" DeliverySuperseded+ liftIO $ markPartitionsAssigned state [(key, PartitionId 0)]+ old.deliveryStatus >>= liftIO . assertEqual "old generation after reassignment" DeliverySuperseded+ fresh <- mkIngested state testConfig (recordAt 42)+ fresh.deliveryStatus >>= liftIO . assertEqual "new assignment" DeliveryCurrent+ )+ >>= assertRight++testGap :: IO ()+testGap = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer+ mock+ ( do+ old <- mkIngested state testConfig (recordAt 46)+ finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))+ Stream.fold Fold.drain $ gateRecords state $ Stream.fromList [Right (Polled 1 (recordAt 45))]+ old.deliveryStatus >>= liftIO . assertEqual "gap fences old token" DeliverySuperseded+ fresh <- mkIngested state testConfig (recordAt 45)+ fresh.deliveryStatus >>= liftIO . assertEqual "next available offset" DeliveryCurrent+ )+ >>= assertRight
test/Shibuya/Adapter/Kafka/IntegrationTest.hs view
@@ -1,5 +1,6 @@ module Shibuya.Adapter.Kafka.IntegrationTest (tests) where +import Control.Concurrent (threadDelay) import Control.Concurrent.Async qualified as Async import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) import Control.Concurrent.STM (TChan, atomically, newTChanIO, readTChan, writeTChan)@@ -10,11 +11,13 @@ import Data.ByteString.Char8 qualified as BS8 import Data.IORef (atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef) import Data.List (nub, sort)+import Data.Map.Strict qualified as Map import Data.Maybe (mapMaybe) import Data.Text qualified as Text import Effectful (Eff, Limit (..), Persistence (..), UnliftStrategy (..), runEff, withEffToIO) import Effectful.Error.Static (runError)-import Kafka.Consumer.Types (ConsumerRecord (..), OffsetCommit (..), OffsetReset (..), RebalanceEvent (..))+import Effectful.Exception qualified as Exception+import Kafka.Consumer.Types (ConsumerRecord (..), OffsetCommit (..), OffsetReset (..), PartitionOffset (..), RebalanceEvent (..), TopicPartition (..)) import Kafka.Effectful.Consumer ( brokersList, groupId,@@ -25,7 +28,7 @@ setCallback, topics, )-import Kafka.Effectful.Consumer.Effect (commitAllOffsets, pollMessageBatch)+import Kafka.Effectful.Consumer.Effect (assignment, commitAllOffsets, committed, pollMessageBatch) import Kafka.TestEnv ( TestEnv (..), consumeN,@@ -33,6 +36,7 @@ createTopicWithPartitions, produceKeyedMessages, produceMessages,+ produceMessagesBounded, producePartitionMessages, withTestEnv, )@@ -45,7 +49,9 @@ ) import Shibuya.Adapter (Adapter (..)) import Shibuya.Adapter.Kafka (KafkaAdapterConfig (..), kafkaAdapter, kafkaAdapterWith, kafkaRebalanceHandler, newKafkaAdapterState)-import Shibuya.App (ProcessorId (..), defaultAppConfig, mkProcessor, runApp, waitApp)+import Shibuya.Adapter.Kafka.Internal (withConsumerLock)+import Shibuya.Adapter.Kafka.SourceGateTest (afterYields, await, waitUntil)+import Shibuya.App (ProcessorId (..), defaultAppConfig, defaultShutdownConfig, mkProcessor, runApp, stopApp, stopAppGracefully) import Shibuya.Core.Ack (AckDecision (..), RetryDelay (..)) import Shibuya.Core.AckHandle (AckHandle (..)) import Shibuya.Core.Ingested (Ingested (..), Message (..))@@ -73,6 +79,8 @@ testCase "AckRetry is not committed past when session exits" testAckRetryAbandonedSession, testCase "late callback after revocation cannot advance the broker offset" testRevokedCallbackDoesNotCommit, testCase "actual reassignment fences a late callback from the old owner" testActualReassignmentFence,+ testCase "a retried record blocks its successors" (liveRetryScenario True False),+ testCase "buffered retry commits every successor" testBufferedRetryCommitsSuccessors, testCase "Handler exception redelivers instead of skipping" testHandlerExceptionRedelivery ] @@ -466,40 +474,7 @@ (expected `elem` replayed) testHandlerExceptionRedelivery :: IO ()-testHandlerExceptionRedelivery = withTestEnv $ \env -> do- createTopic env- let payloads = ["n-1", "n-2", "n-3"]- produceMessages env payloads-- thrown <- newIORef False- seen <- newIORef ([] :: [ByteString])- result <- runEff . runError @KafkaError . runTracingNoop $ do- let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore- sub = topics [env.testTopic] <> offsetReset Earliest- runKafkaConsumer props sub $ do- upstream <- kafkaAdapter (testConfig env)- let finiteAdapter = upstream {source = Stream.take 4 upstream.source}- handler Message {envelope} = do- let payload = maybe "" id envelope.payload- liftIO $ modifyIORef' seen (payload :)- hasThrown <- liftIO $ readIORef thrown- if payload == "n-2" && not hasThrown- then liftIO $ do- writeIORef thrown True- throwIO (userError "planned handler exception")- else pure AckOk- appResult <- runApp defaultAppConfig [(ProcessorId "handler-exception", mkProcessor finiteAdapter handler)]- case appResult of- Left appErr -> liftIO $ assertFailure $ "runApp failed: " <> show appErr- Right appHandle -> waitApp appHandle- commitAllOffsets OffsetCommit- case result of- Left (_cs, err) -> assertFailure $ "handler exception scenario failed: " <> show err- Right () -> pure ()-- delivered <- reverse <$> readIORef seen- assertBool ("expected n-2 at least twice, saw " <> show delivered) (countPayload "n-2" delivered >= 2)- assertBool ("expected n-3 after handler exception retry, saw " <> show delivered) ("n-3" `elem` delivered)+testHandlerExceptionRedelivery = liveRetryScenario True True testConfig :: TestEnv -> KafkaAdapterConfig testConfig env =@@ -548,3 +523,101 @@ partition : _ -> pure partition [] -> loop _ -> loop++-- Observe real finalization, not the handler's intended decision. The initial+-- source handshake forces every original successor into the runner's inbox.+testBufferedRetryCommitsSuccessors :: IO ()+testBufferedRetryCommitsSuccessors = liveRetryScenario False False++liveRetryScenario :: Bool -> Bool -> IO ()+liveRetryScenario exactOrder throws = do+ outcome <- timeout 60000000 $ withTestEnv $ \env -> do+ createTopic env+ let payloads = if throws then ["n-1", "n-2", "n-3"] else [BS8.pack ("o-" <> show n) | n <- [0 .. 11 :: Int]]+ retryPayload = if throws then "n-2" else "o-4"+ expected = if throws then ["n-1", "n-2", "n-2", "n-3"] else take 5 payloads <> drop 4 payloads+ end = length payloads+ target = [TopicPartition env.testTopic (PartitionId 0) (PartitionOffset (fromIntegral end))]+ props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore+ sub = topics [env.testTopic] <> offsetReset Earliest+ produceMessagesBounded env payloads+ seen <- newIORef []+ completed <- newIORef []+ retried <- newIORef False+ queued <- newEmptyMVar+ state <- newKafkaAdapterState+ result <- runEff . runError @KafkaError . runTracingNoop $+ runKafkaConsumer props sub $ do+ upstream <- kafkaAdapterWith state (testConfig env)+ let observe ingested =+ pure+ ingested+ { ack =+ AckHandle+ ( \decision -> do+ let AckHandle finalize = ingested.ack+ finalize decision+ case decision of+ AckOk -> liftIO $ atomicModifyIORef' completed (\xs -> (xs <> [maybe "" id ingested.envelope.payload], ()))+ _ -> pure ()+ )+ }+ adapter = upstream {source = afterYields end queued (Stream.mapM observe upstream.source)}+ handler Message {envelope} = do+ let payload = maybe "" id envelope.payload+ liftIO $ modifyIORef' seen (<> [payload])+ first <-+ if payload == retryPayload+ then liftIO $ atomicModifyIORef' retried (\x -> (True, not x))+ else pure False+ if first+ then do+ liftIO $ await "initial live inbox" queued+ if throws then liftIO $ throwIO (userError "planned handler exception") else pure (AckRetry (RetryDelay 0))+ else pure AckOk+ awaitCommit remaining = do+ offsets <- withConsumerLock state $ committed (Timeout 100) [(env.testTopic, PartitionId 0)]+ if offsets == target+ then pure ()+ else+ if remaining <= (0 :: Int)+ then liftIO $ assertFailure ("committed offset did not reach log end: " <> show offsets)+ else liftIO (threadDelay 100000) >> awaitCommit (remaining - 1)+ app <- runApp defaultAppConfig [(ProcessorId "live-buffered-retry", mkProcessor adapter handler)]+ case app of+ Left err -> liftIO $ assertFailure (show err)+ Right handle ->+ Exception.finally+ ( do+ liftIO $ waitUntil "live successful finalizations" ((== payloads) <$> readIORef completed)+ awaitCommit 300+ drained <- stopAppGracefully defaultShutdownConfig handle+ liftIO $ assertEqual "successful drain" True drained+ )+ (stopApp handle)+ case result of+ Left (_, err) -> assertFailure (show err)+ Right () -> pure ()+ delivered <- readIORef seen+ if exactOrder then assertEqual "handler trace" expected delivered else pure ()+ -- A fresh assigned consumer must retain the durable end offset and stay idle.+ restarted <- runEff . runError @KafkaError $ runKafkaConsumer props sub $ do+ offsets <- committed (Timeout 1000) [(env.testTopic, PartitionId 0)]+ liftIO $ assertEqual "committed after close" target offsets+ let assigned remaining = do+ batch <- pollMessageBatch (Timeout 100) (BatchSize 100)+ liftIO $ assertEqual "no restart deliveries" [] [() | Right _ <- batch]+ owned <- assignment+ if Map.lookup env.testTopic owned == Just [PartitionId 0]+ then pure ()+ else+ if remaining <= (0 :: Int)+ then liftIO $ assertFailure "restart assignment timed out"+ else assigned (remaining - 1)+ assigned 100+ batches <- forM [1 .. 10 :: Int] $ \_ -> pollMessageBatch (Timeout 100) (BatchSize 100)+ liftIO $ assertEqual "assigned restart stays idle" [] [() | Right _ <- concat batches]+ case restarted of+ Left (_, err) -> assertFailure (show err)+ Right () -> pure ()+ assertEqual "bounded live session" (Just ()) outcome
+ test/Shibuya/Adapter/Kafka/SourceGateTest.hs view
@@ -0,0 +1,269 @@+module Shibuya.Adapter.Kafka.SourceGateTest (tests, afterYields, await, waitUntil) where++import Control.Concurrent (threadDelay)+import Control.Concurrent.Async qualified as Async+import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar)+import Control.Concurrent.STM (atomically, writeTVar)+import Control.Exception qualified as IOExc+import Control.Monad (unless)+import Data.IORef (atomicModifyIORef', newIORef, readIORef)+import Data.Int (Int64)+import Data.List.NonEmpty qualified as NE+import Effectful (Eff, IOE, liftIO, runEff, (:>))+import Effectful.Error.Static (runErrorNoCallStack)+import Effectful.Exception qualified as Exception+import Kafka.Consumer.Types (ConsumerRecord (..), Offset (..), RebalanceEvent (..))+import Kafka.ScriptedLog+import Kafka.Types (BatchSize (..), KafkaError, PartitionId (..), Timeout (..))+import Shibuya.Adapter (Adapter (..))+import Shibuya.Adapter.Kafka (KafkaAdapterConfig (..), kafkaAdapter, kafkaRebalanceHandler)+import Shibuya.Adapter.Kafka.AckHandleTest hiding (tests)+import Shibuya.Adapter.Kafka.Internal (KafkaAdapterState (..), Polled (..), gateRecords, kafkaSource, mkAckHandle, newKafkaAdapterState)+import Shibuya.App (ProcessorId (..), defaultAppConfig, defaultShutdownConfig, mkBatchProcessor, mkProcessor, runApp, stopApp, stopAppGracefully)+import Shibuya.Batch (BatchConfig (..), BatchInfo (..), BatchTrigger (..), ackAllOk, ackExcept, defaultBatchConfig)+import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..), RetryDelay (..))+import Shibuya.Core.AckHandle (AckHandle (..))+import Shibuya.Core.Ingested (Message (..))+import Shibuya.Core.Types (Cursor (..), Envelope (..))+import Shibuya.Telemetry.Effect (runTracingNoop)+import Streamly.Data.Fold qualified as Fold+import Streamly.Data.Stream (Stream)+import Streamly.Data.Stream qualified as Stream+import System.Timeout (timeout)+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertEqual, assertFailure, testCase)++tests :: TestTree+tests =+ testGroup+ "Source gate"+ [ testCase "a retried record blocks its successors (scripted log)" (scriptedRetry 1),+ testCase "a twice-retried record blocks its successors (scripted log)" (scriptedRetry 2),+ testCase "pre-seek successors are dropped" testStale,+ testCase "replay holds successor until success" (testHold "ok"),+ testCase "retry drops held successor" (testHold "retry"),+ testCase "shutdown wakes held successor" (testHold "shutdown"),+ testCase "halt wakes held successor" (testHold "halt"),+ testCase "revocation wakes held successor" (testHold "revoke"),+ testCase "dead letter releases held successor" (testHold "dead-letter"),+ testCase "cancellation releases held gate worker" (testHold "cancel"),+ testCase "fatal acknowledgement wakes held successor" (testHold "fatal"),+ testCase "missing replay fences old callbacks" testGap,+ testCase "seek affects only its own partition" testPartitionIsolation,+ testCase "requested incomplete seek makes records stale" testPendingSeek,+ testCase "batch timeout releases replay and stores all successors" testBatchLiveness+ ]++-- Signal on the pull after the nth yield, before pulling upstream again.+afterYields :: (IOE :> es) => Int -> MVar () -> Stream (Eff es) a -> Stream (Eff es) a+afterYields count ready = Stream.unfoldrM step . (0,)+ where+ step (n, stream) = do+ if n == count then liftIO (putMVar ready ()) else pure ()+ next <- Stream.uncons stream+ pure $ case next of+ Nothing -> Nothing+ Just (value, rest) -> Just (value, (n + 1, rest))++await :: String -> MVar a -> IO a+await label var = timeout 5000000 (takeMVar var) >>= maybe (assertFailure (label <> " timed out") >> error label) pure++waitUntil :: String -> IO Bool -> IO ()+waitUntil label check = do+ result <- timeout 10000000 loop+ unless (result == Just ()) $ assertFailure (label <> " timed out")+ where+ loop = check >>= \done -> unless done (threadDelay 1000 >> loop)++scriptedRetry :: Int -> IO ()+scriptedRetry retries = do+ scripted <- newScriptedLog 12+ deliveries <- newIORef ([] :: [Int64])+ attempts <- newIORef (0 :: Int)+ queued <- newEmptyMVar+ outcome <- timeout 30000000 $+ runEff . runErrorNoCallStack @KafkaError . runScriptedConsumer scripted . runTracingNoop $ do+ upstream <- kafkaAdapter KafkaAdapterConfig {topics = [scriptedTopic], pollTimeout = Timeout 5, batchSize = BatchSize 100}+ let adapter = upstream {source = afterYields 12 queued upstream.source}+ handler Message {envelope} = do+ let offset = case envelope.cursor of+ Just (CursorInt value) -> fromIntegral value+ _ -> -1+ liftIO $ atomicModifyIORef' deliveries (\xs -> (xs <> [offset], ()))+ if offset == 4+ then do+ attempt <- liftIO $ atomicModifyIORef' attempts (\n -> (n + 1, n))+ if attempt == 0 then liftIO (await "initial inbox" queued) else pure ()+ pure $ if attempt < retries then AckRetry (RetryDelay 0) else AckOk+ else pure AckOk+ app <- runApp defaultAppConfig [(ProcessorId "scripted-retry", mkProcessor adapter handler)]+ case app of+ Left err -> liftIO $ assertFailure (show err)+ Right handle ->+ Exception.finally+ ( do+ liftIO $+ waitUntil "all offsets stored" ((== [0 .. 11]) <$> readIORef scripted.stored)+ `IOExc.onException` ( do+ stores <- readIORef scripted.stored+ trace <- readIORef deliveries+ putStrLn ("stored=" <> show stores <> " deliveries=" <> show trace)+ )+ drained <- stopAppGracefully defaultShutdownConfig handle+ liftIO $ assertEqual "successful drain" True drained+ )+ (stopApp handle)+ observed <- readIORef deliveries+ stores <- readIORef scripted.stored+ seeks <- readIORef scripted.seeks+ assertEqual ("completion; deliveries=" <> show observed <> "; stored=" <> show stores) (Just (Right ())) outcome+ assertEqual "handler order" ([0 .. 4] <> replicate (retries - 1) 4 <> [4 .. 11]) observed+ assertEqual "stored offsets" [0 .. 11] stores+ assertEqual "seeks" (replicate retries 4) seeks++right :: (Show e) => Either e a -> IO a+right = either (\err -> assertFailure (show err) >> error "unexpected Left") pure++gateOffsets :: KafkaAdapterState -> [Polled] -> IO [Int64]+gateOffsets state values = runEff $ Stream.fold Fold.toList $ fmap (either (const (-1)) (unOffset . (.crOffset))) $ gateRecords state $ Stream.fromList (map Right values)++testStale :: IO ()+testStale = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer mock (finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))) >>= assertRight+ gateOffsets state [Polled 0 (recordAt 43), Polled 0 (recordAt 44)] >>= assertEqual "stale records" []++testHold :: String -> IO ()+testHold mode = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer mock (finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))) >>= assertRight+ emitted <- newEmptyMVar+ offered <- newEmptyMVar+ successor <- newEmptyMVar+ let input =+ Stream.mapM+ ( \p -> do+ if p.record.crOffset == Offset 43 then liftIO (putMVar offered ()) else pure ()+ pure (Right p)+ )+ $ Stream.fromList [Polled 1 (recordAt 42), Polled 1 (recordAt 43)]+ worker =+ runEff+ $ Stream.fold Fold.drain+ $ Stream.mapM+ ( \case+ Right cr | cr.crOffset == Offset 42 -> liftIO $ do+ ack <- runFinalizer mock (mkAckHandle state testConfig cr) >>= right+ putMVar emitted ack+ Right _ -> liftIO (putMVar successor ())+ Left _ -> error "unexpected error"+ )+ $ gateRecords state input+ Async.withAsync worker $ \async -> do+ AckHandle finalize <- await "replay emission" emitted+ await "successor offered" offered+ timeout 200000 (takeMVar successor) >>= assertEqual "held successor" Nothing+ case mode of+ "shutdown" -> atomically $ writeTVar state.shutdownVar True+ "revoke" -> kafkaRebalanceHandler state (error "unused consumer") (RebalanceRevoke [(scriptedTopic, PartitionId 0)])+ "cancel" -> Async.cancel async+ "fatal" -> do+ atomicModifyIORef' mock (\m -> (m {storeFailuresRemaining = 99}, ()))+ failed <- IOExc.try @IOExc.SomeException $ runFinalizer mock (finalizeRecord state (recordAt 1) {crPartition = PartitionId 1} AckOk)+ case failed of+ Left _ -> pure ()+ Right _ -> assertFailure "expected exhausted acknowledgement"+ _ ->+ runFinalizer+ mock+ ( finalize+ ( case mode of+ "retry" -> AckRetry (RetryDelay 0)+ "halt" -> AckHalt (HaltFatal "stop")+ "dead-letter" -> AckDeadLetter MaxRetriesExceeded+ _ -> AckOk+ )+ )+ >>= assertRight+ if mode == "cancel"+ then do+ timeout 1000000 (runFinalizer mock (finalize AckOk)) >>= maybe (assertFailure "lock leaked") assertRight+ else do+ timeout 1000000 (Async.wait async) >>= assertEqual "gate worker exits" (Just ())+ if mode `elem` ["ok", "dead-letter"]+ then await "successor released" successor+ else timeout 10000 (takeMVar successor) >>= assertEqual "successor dropped" Nothing+ if mode == "fatal"+ then do+ result <- runFinalizer mock $ Stream.fold Fold.drain $ gateRecords state $ kafkaSource state testConfig+ case result of+ Left _ -> pure ()+ Right _ -> assertFailure "fatal slot must fail before polling"+ else pure ()++testGap :: IO ()+testGap = mapM_ check [AckOk, AckRetry (RetryDelay 0)]+ where+ check decision = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ old <- runFinalizer mock (mkAckHandle state testConfig (recordAt 46)) >>= right+ runFinalizer mock (finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))) >>= assertRight+ gateOffsets state [Polled 1 (recordAt 45), Polled 1 (recordAt 46)] >>= assertEqual "gap resumes" [45, 46]+ let AckHandle finalize = old+ runFinalizer mock (finalize decision >> finalizeRecord state (recordAt 45) AckOk) >>= assertRight+ m <- readIORef mock+ assertEqual "only new delivery stores" [Offset 45] m.storedOffsets+ assertEqual "old callback cannot seek" 1 (length m.seekCalls)++testPartitionIsolation :: IO ()+testPartitionIsolation = do+ state <- newKafkaAdapterState+ mock <- newIORef defaultMockState+ runFinalizer mock (finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0))) >>= assertRight+ gateOffsets state [Polled 0 (recordAt 43) {crPartition = PartitionId 1}] >>= assertEqual "other partition stays current" [43]+ kafkaRebalanceHandler state (error "unused") (RebalanceRevoke [(scriptedTopic, PartitionId 0)])+ gateOffsets state [Polled 1 (recordAt 42)] >>= assertEqual "revoked partition" []++testPendingSeek :: IO ()+testPendingSeek = do+ state <- newKafkaAdapterState+ started <- newEmptyMVar+ release <- newEmptyMVar+ mock <- newIORef defaultMockState {seekBlock = Just (started, release)}+ Async.withAsync (runFinalizer mock (finalizeRecord state (recordAt 42) (AckRetry (RetryDelay 0)))) $ \worker -> do+ await "seek entered" started+ gateOffsets state [Polled 0 (recordAt 42)] >>= assertEqual "pending seek drops record" []+ putMVar release ()+ Async.wait worker >>= assertRight+ gateOffsets state [Polled 1 (recordAt 42)] >>= assertEqual "completed seek admits replay" [42]++testBatchLiveness :: IO ()+testBatchLiveness = do+ scripted <- newScriptedLog 12+ first <- newIORef True+ batches <- newIORef []+ outcome <- timeout 10000000 $+ runEff . runErrorNoCallStack @KafkaError . runScriptedConsumer scripted . runTracingNoop $ do+ adapter <- kafkaAdapter KafkaAdapterConfig {topics = [scriptedTopic], pollTimeout = Timeout 5, batchSize = BatchSize 100}+ let config = defaultBatchConfig {batchSize = 12, batchTimeout = 0.05, tickInterval = Just 0.01}+ handler info messages = do+ isFirst <- liftIO $ atomicModifyIORef' first (\x -> (False, x))+ liftIO $ atomicModifyIORef' batches (\xs -> (xs <> [(info.trigger, info.size)], ()))+ pure $ if isFirst then ackExcept [((NE.head messages).envelope.messageId, AckRetry (RetryDelay 0))] else ackAllOk+ app <- runApp defaultAppConfig [(ProcessorId "batch-replay", mkBatchProcessor adapter handler config)]+ case app of+ Left err -> liftIO $ assertFailure (show err)+ Right handle ->+ Exception.finally+ ( do+ liftIO $ waitUntil "batch stores" ((== [0 .. 11]) <$> readIORef scripted.stored)+ drained <- stopAppGracefully defaultShutdownConfig handle+ liftIO $ assertEqual "drain" True drained+ )+ (stopApp handle)+ assertEqual "completion" (Just (Right ())) outcome+ observed <- readIORef batches+ unless ((TriggerTimeout, 1) `elem` observed) $ assertFailure ("missing replay timeout: " <> show observed)