SubscriptionPauseResume.hs

{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}

{- | Regression tests for ExecPlan 41 M2 — recoverable backpressure.

Before M2, a slow AllStreams consumer that filled its bounded queue was
__terminal__: the publisher marked it @Overflowed@ and the worker threw
'Kiroku.Store.Subscription.Types.SubscriptionOverflowed' and died (overflow
policy 'Kiroku.Store.Subscription.Types.DropSubscription'). M2 adds the
recoverable 'Kiroku.Store.Subscription.Types.PauseAndResume' policy (now the
default): when the queue fills the publisher /pauses/ delivery rather than
killing the subscriber, and the worker — once its slow handler catches up —
drains the stale queue and re-reads the events it missed directly from the
database from its checkpoint. No event is lost and the checkpoint advances
monotonically.

The first spec proves the recovery: a handler that blocks long enough to fill a
@queueCapacity = 1@ subscriber while five events are appended still delivers all
five in order, stops cleanly (not via 'SubscriptionOverflowed'), advances the
checkpoint to the last position, and emits a
'Kiroku.Store.Observability.KirokuEventSubscriptionPaused' followed by a
'Kiroku.Store.Observability.KirokuEventSubscriptionResumed'.

The second spec proves the fail-fast path is preserved: the identical scenario
under 'DropSubscription' still surfaces 'SubscriptionOverflowed' on @wait@.

To confirm the first spec actually pins the new behavior: temporarily set its
@overflowPolicy@ to 'DropSubscription' and it fails with
@Left (SubscriptionOverflowed ...)@ where it expected @Right ()@.
-}
module Test.SubscriptionPauseResume (spec) where

import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Concurrent.STM (atomically, newTVarIO, readTVar, writeTVar)
import Control.Exception qualified
import Control.Lens ((&), (.~), (^.))
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.Int (Int32)
import Data.Text qualified as T
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Kiroku.Store
import Kiroku.Store.SQL qualified as SQL
import Test.Helpers (makeEvent, waitForPublisher, waitWithTimeout, withTestStoreSettings)
import Test.Hspec

-- Read the durable checkpoint (@kiroku.subscriptions.last_seen@) for a non-group
-- subscription (member 0), or 'Nothing' if no row exists yet.
readCheckpoint :: KirokuStore -> T.Text -> IO (Maybe Int)
readCheckpoint store subName = do
    result <- Pool.use (store ^. #pool) (Session.statement (subName, 0 :: Int32) SQL.getCheckpointMemberStmt)
    case result of
        Left err -> error ("readCheckpoint failed: " <> show err)
        Right mPos -> pure (fmap fromIntegral mPos)

spec :: Spec
spec = describe "subscription FSM — recoverable backpressure (EP-41 M2)" $ do
    it "pauses a slow AllStreams consumer and resumes, delivering all events" $ do
        evtRef <- newIORef ([] :: [KirokuEvent])
        let evtHandler e = modifyIORef' evtRef (e :)
        withTestStoreSettings (& #eventHandler .~ Just evtHandler) $ \store -> do
            firstSeen <- newEmptyMVar
            release <- newEmptyMVar
            delivered <- newIORef ([] :: [GlobalPosition])
            seenCount <- newTVarIO (0 :: Int)
            let handler' evt = do
                    modifyIORef' delivered ((evt ^. #globalPosition) :)
                    n <- atomically $ do
                        c <- readTVar seenCount
                        writeTVar seenCount (c + 1)
                        pure (c + 1)
                    if n == 1
                        then putMVar firstSeen () >> takeMVar release
                        else pure ()
                    -- Stop once all five have been seen, so `wait` resolves Right ().
                    if n >= 5 then pure Stop else pure Continue
            let cfg =
                    SubscriptionConfig
                        { name = SubscriptionName "pause-resume-test"
                        , target = AllStreams
                        , handler = handler'
                        , batchSize = 100
                        , queueCapacity = 1
                        , overflowPolicy = PauseAndResume
                        , consumerGroup = Nothing
                        , consumerGroupGuard = False
                        , missingCheckpointPolicy = FromBeginning
                        , retryPolicy = defaultRetryPolicy
                        , eventTypeFilter = AllEventTypes
                        , selector = Nothing
                        }
            handle <- subscribe store cfg
            -- First append: the worker reads it from the queue and the handler
            -- blocks inside it on `release`, so the worker stops draining.
            Right _ <- runStoreIO store $ appendToStream (StreamName "pr-1") NoStream [makeEvent "E1" (Aeson.object [])]
            takeMVar firstSeen
            -- With the worker stuck, append four more one at a time. Capacity is 1,
            -- so the second of these fills the queue and the publisher pauses the
            -- subscriber for the rest (skipping, not dropping).
            let appendOne i = do
                    let sn = StreamName ("pr-" <> T.pack (show (i :: Int)))
                    Right _ <- runStoreIO store $ appendToStream sn NoStream [makeEvent "Ex" (Aeson.object [])]
                    waitForPublisher store (GlobalPosition (fromIntegral i))
            appendOne 2
            appendOne 3
            appendOne 4
            appendOne 5
            -- Release the slow handler; the worker drains, resumes, re-catches-up
            -- from its checkpoint, and delivers the skipped events.
            putMVar release ()
            result <- waitWithTimeout 15_000_000 handle
            case result of
                Left timeout -> expectationFailure timeout
                Right (Left e) -> expectationFailure ("expected a clean stop, got: " <> show e)
                Right (Right ()) -> pure ()
            ds <- reverse <$> readIORef delivered
            ds `shouldBe` map GlobalPosition [1 .. 5]
            cp <- readCheckpoint store "pause-resume-test"
            cp `shouldBe` Just 5
            evts <- readIORef evtRef
            let isPaused (KirokuEventSubscriptionPaused (SubscriptionName "pause-resume-test") _ _) = True
                isPaused _ = False
                isResumed (KirokuEventSubscriptionResumed (SubscriptionName "pause-resume-test") _ _) = True
                isResumed _ = False
            any isPaused evts `shouldBe` True
            any isResumed evts `shouldBe` True

    it "still fails fast under DropSubscription" $ do
        withTestStoreSettings Prelude.id $ \store -> do
            firstSeen <- newEmptyMVar
            release <- newEmptyMVar
            seenCount <- newTVarIO (0 :: Int)
            let handler' _evt = do
                    n <- atomically $ do
                        c <- readTVar seenCount
                        writeTVar seenCount (c + 1)
                        pure (c + 1)
                    if n == 1
                        then putMVar firstSeen () >> takeMVar release
                        else pure ()
                    pure Continue
            let cfg =
                    SubscriptionConfig
                        { name = SubscriptionName "dropsub-overflow-test"
                        , target = AllStreams
                        , handler = handler'
                        , batchSize = 100
                        , queueCapacity = 1
                        , overflowPolicy = DropSubscription
                        , consumerGroup = Nothing
                        , consumerGroupGuard = False
                        , missingCheckpointPolicy = FromBeginning
                        , retryPolicy = defaultRetryPolicy
                        , eventTypeFilter = AllEventTypes
                        , selector = Nothing
                        }
            handle <- subscribe store cfg
            Right _ <- runStoreIO store $ appendToStream (StreamName "ds-1") NoStream [makeEvent "E1" (Aeson.object [])]
            takeMVar firstSeen
            let appendOne i = do
                    let sn = StreamName ("ds-" <> T.pack (show (i :: Int)))
                    Right _ <- runStoreIO store $ appendToStream sn NoStream [makeEvent "Ex" (Aeson.object [])]
                    waitForPublisher store (GlobalPosition (fromIntegral i))
            appendOne 2
            appendOne 3
            appendOne 4
            appendOne 5
            putMVar release ()
            result <- waitWithTimeout 15_000_000 handle
            case result of
                Left timeout -> expectationFailure timeout
                Right (Right ()) -> expectationFailure "expected SubscriptionOverflowed, got a clean stop"
                Right (Left e) ->
                    case Control.Exception.fromException e of
                        Just (SubscriptionOverflowed sn) -> sn `shouldBe` SubscriptionName "dropsub-overflow-test"
                        Nothing -> expectationFailure ("expected SubscriptionOverflowed, got: " <> show e)