SubscriptionCheckpointWorker.hs

{-# LANGUAGE NumericUnderscores #-}

module Test.SubscriptionCheckpointWorker (spec) where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar, tryPutMVar)
import Control.Concurrent.STM (atomically, check, modifyTVar', newTVarIO, readTVar, readTVarIO, writeTVar)
import Control.Exception (SomeException, finally, fromException, try)
import Control.Lens ((&), (.~), (^.))
import Control.Monad (forM, void)
import Control.Monad.IO.Class (liftIO)
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.List (sortOn)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Vector qualified as V
import Effectful (runEff)
import Kiroku.Store
import Kiroku.Store.Subscription.Effect qualified as SubEff
import Kiroku.Store.Subscription.Stream (subscriptionAckStream)
import Kiroku.Store.Subscription.Stream qualified as Buffer
import Streamly.Data.Stream qualified as Stream
import Test.Helpers (makeEvent, validConsumerGroup, waitForPublisher, waitWithTimeout, withTestStore, withTestStoreSettings)
import Test.Hspec

spec :: Spec
spec = describe "subscription checkpoint worker policies" $ do
  it "starts a non-group worker FromBeginning and reports the durable seed" $ do
    eventsRef <- newIORef []
    resolutionsRef <- newIORef []
    let name = SubscriptionName "worker-from-beginning"
        observe evt = case evt of
          KirokuEventSubscriptionCheckpointResolved initialization NonGroup ->
            modifyIORef' resolutionsRef (initialization :)
          _ -> pure ()
        tweak settings = settings & #eventHandler .~ Just observe
    withTestStoreSettings tweak $ \store -> do
      _ <- appendBatch store "worker-from-beginning-events" 3
      waitForPublisher store (GlobalPosition 3)
      let handler event = do
            modifyIORef' eventsRef (event ^. #globalPosition :)
            seen <- length <$> readIORef eventsRef
            pure (if seen >= 3 then Stop else Continue)
          config =
            (defaultSubscriptionConfig name AllStreams handler)
              { missingCheckpointPolicy = FromBeginning
              }
      handle <- subscribe store config
      expectClean handle

      reverse <$> readIORef eventsRef
        `shouldReturn` fmap GlobalPosition [1, 2, 3]
      inventoryPositions store
        `shouldReturn` [(SubscriptionCheckpointKey name 0, GlobalPosition 3)]

    readIORef resolutionsRef
      `shouldReturn` [InitializedCheckpoint FromBeginning (SubscriptionCheckpointKey name 0) (GlobalPosition 0)]

  it "forms a clean FromCurrentHead cut while appends race startup" $ do
    initializationReady <- newEmptyMVar
    deliveredRef <- newIORef []
    let name = SubscriptionName "worker-current-head-race"
        key = SubscriptionCheckpointKey name 0
        observe evt = case evt of
          KirokuEventSubscriptionCheckpointResolved initialization NonGroup
            | initializationKey initialization == key ->
                void (tryPutMVar initializationReady initialization)
          _ -> pure ()
        tweak settings = settings & #eventHandler .~ Just observe
    withTestStoreSettings tweak $ \store -> do
      _ <- appendBatch store "worker-current-head-race-events" 10
      waitForPublisher store (GlobalPosition 10)
      gate <- newEmptyMVar
      let handler event = do
            modifyIORef' deliveredRef (event ^. #globalPosition :)
            pure $ if event ^. #eventType == EventType "RaceSentinel" then Stop else Continue
          config =
            (defaultSubscriptionConfig name AllStreams handler)
              { missingCheckpointPolicy = FromCurrentHead,
                batchSize = either (error . show) Prelude.id (mkBatchSize 3)
              }
      subscribeThread <- Async.async (takeMVar gate >> subscribe store config)
      appendThread <- Async.async $ do
        takeMVar gate
        forM [1 .. 20 :: Int] $ \i ->
          appendExisting store "worker-current-head-race-events" ("Racing" <> T.pack (show i))
      putMVar gate ()
      putMVar gate ()
      handle <- Async.wait subscribeThread
      initialization <- waitMVar "checkpoint resolution" initializationReady
      racePositions <- Async.wait appendThread
      raceTail <- case reverse racePositions of
        [] -> expectationFailure "expected racing appends" >> error "unreachable"
        position : _ -> pure position
      let seed = checkpointInitializationPosition initialization
      waitForPublisher store raceTail
      finalPosition <- appendExisting store "worker-current-head-race-events" "RaceSentinel"
      waitForPublisher store finalPosition
      expectClean handle

      initialization `shouldBe` InitializedCheckpoint FromCurrentHead key seed
      delivered <- reverse <$> readIORef deliveredRef
      delivered `shouldBe` positionsAfter seed finalPosition

  it "initializes each consumer-group member independently at the current head" $ do
    resolutionsVar <- newTVarIO []
    handlerCalled <- newTVarIO False
    let name = SubscriptionName "worker-current-head-members"
        observe evt = case evt of
          KirokuEventSubscriptionCheckpointResolved initialization GroupMember {} ->
            atomically (modifyTVar' resolutionsVar (initialization :))
          _ -> pure ()
        tweak settings = settings & #eventHandler .~ Just observe
    withTestStoreSettings tweak $ \store -> do
      _ <- appendBatch store "worker-current-head-members-events" 6
      waitForPublisher store (GlobalPosition 6)
      let config member =
            ( defaultSubscriptionConfig name AllStreams $ \_ -> do
                atomically (writeTVar handlerCalled True)
                pure Continue
            )
              { consumerGroup = Just (validConsumerGroup member 2),
                missingCheckpointPolicy = FromCurrentHead
              }
      handles <- mapM (subscribe store . config) [0, 1]
      atomically $ do
        resolutions <- readTVar resolutionsVar
        check (length resolutions >= 2)
      mapM_ cancel handles
      mapM_ wait handles

      resolutions <- sortOn initializationKey <$> readTVarIO resolutionsVar
      resolutions
        `shouldBe` [ InitializedCheckpoint FromCurrentHead (SubscriptionCheckpointKey name 0) (GlobalPosition 6),
                     InitializedCheckpoint FromCurrentHead (SubscriptionCheckpointKey name 1) (GlobalPosition 6)
                   ]
      readTVarIO handlerCalled `shouldReturn` False
      inventoryPositions store
        `shouldReturn` [ (SubscriptionCheckpointKey name 0, GlobalPosition 6),
                         (SubscriptionCheckpointKey name 1, GlobalPosition 6)
                       ]

  it "preserves FailIfMissing through the bracketed plain-IO entry point" $
    withTestStore $ \store -> do
      let name = SubscriptionName "worker-bracketed-missing"
          key = SubscriptionCheckpointKey name 0
          config =
            (defaultSubscriptionConfig name AllStreams (\_ -> pure Continue))
              { missingCheckpointPolicy = FailIfMissing
              }
      outcome <- withSubscription store config wait
      case outcome of
        Left err
          | Just (SubscriptionCheckpointMissing actual) <- fromException err ->
              actual `shouldBe` key
        other -> expectationFailure ("expected bracketed missing refusal, got: " <> show other)

  it "preserves FromCurrentHead through the higher-order effect entry point" $ do
    initializationReady <- newEmptyMVar
    deliveredRef <- newIORef []
    let name = SubscriptionName "worker-effect-current-head"
        key = SubscriptionCheckpointKey name 0
        observe evt = case evt of
          KirokuEventSubscriptionCheckpointResolved initialization NonGroup
            | initializationKey initialization == key ->
                void (tryPutMVar initializationReady initialization)
          _ -> pure ()
        tweak settings = settings & #eventHandler .~ Just observe
    withTestStoreSettings tweak $ \store -> do
      _ <- appendBatch store "worker-effect-current-head-events" 5
      waitForPublisher store (GlobalPosition 5)
      runEff $ SubEff.runSubscription store $ do
        let config =
              ( defaultSubscriptionConfig name AllStreams $ \event -> do
                  liftIO (modifyIORef' deliveredRef (event ^. #globalPosition :))
                  pure Stop
              )
                { missingCheckpointPolicy = FromCurrentHead
                }
        handle <- SubEff.subscribe config
        initialization <- liftIO (waitMVar "effect checkpoint resolution" initializationReady)
        liftIO $ initialization `shouldBe` InitializedCheckpoint FromCurrentHead key (GlobalPosition 5)
        position <- liftIO $ appendExisting store "worker-effect-current-head-events" "EffectFuture"
        liftIO (waitForPublisher store position)
        liftIO (expectClean handle)
      reverse <$> readIORef deliveredRef `shouldReturn` [GlobalPosition 6]

  it "preserves FailIfMissing through the Streamly bridge" $
    withTestStore $ \store -> do
      let name = SubscriptionName "worker-streamly-missing"
          key = SubscriptionCheckpointKey name 0
          config =
            (defaultSubscriptionConfig name AllStreams (\_ -> pure Continue))
              { missingCheckpointPolicy = FailIfMissing
              }
      (stream, cancelStream) <- subscriptionAckStream store config (either (error . show) Prelude.id (Buffer.mkStreamBufferSize 1))
      pulled <- finally (try (Stream.uncons stream)) cancelStream
      case pulled of
        Left err
          | Just (SubscriptionCheckpointMissing actual) <- fromException (err :: SomeException) ->
              actual `shouldBe` key
        Left err -> expectationFailure ("expected typed Streamly refusal, got: " <> show err)
        Right _ -> expectationFailure "expected Streamly bridge startup to fail"

initializationKey :: CheckpointInitialization -> SubscriptionCheckpointKey
initializationKey = \case
  ExistingCheckpoint key _ -> key
  InitializedCheckpoint _ key _ -> key

positionsAfter :: GlobalPosition -> GlobalPosition -> [GlobalPosition]
positionsAfter (GlobalPosition start) (GlobalPosition end) =
  fmap GlobalPosition [start + 1 .. end]

inventoryPositions :: KirokuStore -> IO [(SubscriptionCheckpointKey, GlobalPosition)]
inventoryPositions store = do
  Right (SubscriptionCheckpointInventory _ rows) <- runStoreIO store subscriptionCheckpointInventory
  pure
    [ (SubscriptionCheckpointKey name member, position)
    | SubscriptionCheckpoint name member position _ <- V.toList rows
    ]

appendBatch :: KirokuStore -> Text -> Int -> IO GlobalPosition
appendBatch store stream count = do
  let events = [makeEvent ("History" <> T.pack (show i)) (Aeson.object []) | i <- [1 .. count]]
  Right result <- runStoreIO store $ appendToStream (StreamName stream) NoStream events
  pure (result ^. #globalPosition)

appendExisting :: KirokuStore -> Text -> Text -> IO GlobalPosition
appendExisting store stream typ = do
  Right result <-
    runStoreIO store $
      appendToStream (StreamName stream) StreamExists [makeEvent typ (Aeson.object [])]
  pure (result ^. #globalPosition)

expectClean :: SubscriptionHandle -> IO ()
expectClean handle = do
  result <- waitWithTimeout 15_000_000 handle
  case result of
    Left message -> expectationFailure message
    Right (Left err) -> expectationFailure ("subscription failed: " <> show err)
    Right (Right ()) -> pure ()

waitMVar :: String -> MVar a -> IO a
waitMVar label var = do
  result <- Async.race (threadDelay 10_000_000) (takeMVar var)
  case result of
    Left () -> expectationFailure ("timed out waiting for " <> label) >> error "unreachable"
    Right value -> pure value