PublisherRestartNoRebroadcast.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.PublisherRestartNoRebroadcast (spec) where

import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Concurrent.STM (atomically, newTVarIO, readTVar, writeTVar)
import Control.Lens ((&), (.~), (^.))
import Control.Monad (when)
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.Text qualified as T
import Kiroku.Store
import Kiroku.Test.Postgres (withMigratedTestDatabase)
import Test.Helpers (caughtUpEventHandler, makeEvent, waitForPublisher, waitForSubscriptionLive, waitWithTimeout)
import Test.Hspec

spec :: Spec
spec =
  describe "publisher restart tail initialization" $ do
    it "does not rebroadcast restart history into a blocked all-stream subscriber" $
      withMigratedTestDatabase $ \connStr -> do
        let settings = defaultConnectionSettings connStr
        withStore settings $ \store ->
          appendBatch store "tail-init-overflow-seed" "TailInitSeed" 1001

        firstSeen <- newEmptyMVar
        releaseFirst <- newEmptyMVar
        caughtUp <- newEmptyMVar
        blockedFirst <- newTVarIO True
        ref <- newIORef ([] :: [GlobalPosition])

        let subName = SubscriptionName "tail-init-no-rebroadcast"
            stopType = EventType "TailInitStop"
            handler' evt = do
              modifyIORef' ref ((evt ^. #globalPosition) :)
              shouldBlock <- atomically $ do
                block <- readTVar blockedFirst
                when block (writeTVar blockedFirst False)
                pure block
              when shouldBlock $ do
                putMVar firstSeen ()
                takeMVar releaseFirst
              if evt ^. #eventType == stopType
                then pure Stop
                else pure Continue
            cfg =
              SubscriptionConfig
                { name = subName,
                  target = AllStreams,
                  handler = handler',
                  batchSize = either (error . show) Prelude.id (mkBatchSize 10),
                  queueCapacity = 1,
                  overflowPolicy = DropSubscription,
                  consumerGroup = Nothing,
                  consumerGroupGuard = False,
                  missingCheckpointPolicy = FromBeginning,
                  targetBindingPolicy = AdoptUnbound,
                  undecodableHandler = Nothing,
                  handlerStallWarnAfter = Nothing,
                  retryPolicy = defaultRetryPolicy,
                  eventTypeFilter = AllEventTypes,
                  selector = Nothing
                }
            liveSettings =
              settings
                & #eventHandler .~ Just (caughtUpEventHandler subName caughtUp Nothing)

        withStore liveSettings $ \store -> do
          handle <- subscribe store cfg
          takeMVar firstSeen

          Right _ <-
            runStoreIO store $
              appendToStream
                (StreamName "tail-init-overflow-live")
                NoStream
                [makeEvent "TailInitDuringCatchUp" (Aeson.object [])]
          waitForPublisher store (GlobalPosition 1002)

          putMVar releaseFirst ()
          waitForSubscriptionLive caughtUp

          Right _ <-
            runStoreIO store $
              appendToStream
                (StreamName "tail-init-overflow-stop")
                NoStream
                [makeEvent "TailInitStop" (Aeson.object [])]

          result <- waitWithTimeout 10_000_000 handle
          case result of
            Left timeout -> expectationFailure timeout
            Right (Left err) -> expectationFailure ("Subscription failed: " <> show err)
            Right (Right ()) -> pure ()

        collected <- reverse <$> readIORef ref
        collected `shouldBe` map GlobalPosition [1 .. 1003]

    it "delivers a live event appended after a restarted subscriber catches up" $
      withMigratedTestDatabase $ \connStr -> do
        let settings = defaultConnectionSettings connStr
            subName = SubscriptionName "tail-init-live-after-catchup"
        withStore settings $ \store ->
          appendBatch store "tail-init-live-seed" "TailInitLiveSeed" 3

        caughtUp <- newEmptyMVar
        ref <- newIORef ([] :: [GlobalPosition])
        let cfg =
              SubscriptionConfig
                { name = subName,
                  target = AllStreams,
                  handler = \evt -> do
                    modifyIORef' ref ((evt ^. #globalPosition) :)
                    if evt ^. #globalPosition == GlobalPosition 4
                      then pure Stop
                      else pure Continue,
                  batchSize = either (error . show) Prelude.id (mkBatchSize 10),
                  queueCapacity = 1,
                  overflowPolicy = DropSubscription,
                  consumerGroup = Nothing,
                  consumerGroupGuard = False,
                  missingCheckpointPolicy = FromBeginning,
                  targetBindingPolicy = AdoptUnbound,
                  undecodableHandler = Nothing,
                  handlerStallWarnAfter = Nothing,
                  retryPolicy = defaultRetryPolicy,
                  eventTypeFilter = AllEventTypes,
                  selector = Nothing
                }
            liveSettings =
              settings
                & #eventHandler .~ Just (caughtUpEventHandler subName caughtUp Nothing)

        withStore liveSettings $ \store -> do
          handle <- subscribe store cfg
          waitForSubscriptionLive caughtUp
          Right _ <-
            runStoreIO store $
              appendToStream
                (StreamName "tail-init-live-after-catchup")
                NoStream
                [makeEvent "TailInitLiveAfterCatchUp" (Aeson.object [])]
          result <- waitWithTimeout 10_000_000 handle
          case result of
            Left timeout -> expectationFailure timeout
            Right (Left err) -> expectationFailure ("Subscription failed: " <> show err)
            Right (Right ()) -> pure ()

        collected <- reverse <$> readIORef ref
        collected `shouldBe` map GlobalPosition [1 .. 4]

    it "keeps per-subscription checkpoints independent from the publisher tail" $
      withMigratedTestDatabase $ \connStr -> do
        let settings = defaultConnectionSettings connStr
        withStore settings $ \store ->
          appendBatch store "tail-init-checkpoint-seed" "TailInitCheckpointSeed" 10

        withStore settings $ \store -> do
          runUntil store (SubscriptionName "tail-init-sub-a") 3
            `shouldReturn` map GlobalPosition [1 .. 3]
          runUntil store (SubscriptionName "tail-init-sub-b") 7
            `shouldReturn` map GlobalPosition [1 .. 7]

        withStore settings $ \store -> do
          runUntil store (SubscriptionName "tail-init-sub-a") 7
            `shouldReturn` map GlobalPosition [4 .. 10]
          runUntil store (SubscriptionName "tail-init-sub-b") 3
            `shouldReturn` map GlobalPosition [8 .. 10]

appendBatch :: KirokuStore -> T.Text -> T.Text -> Int -> IO ()
appendBatch store streamName typePrefix count = do
  let events =
        [ makeEvent (typePrefix <> T.pack (show i)) (Aeson.object [])
        | i <- [1 .. count]
        ]
  result <- runStoreIO store $ appendToStream (StreamName streamName) NoStream events
  case result of
    Left err -> expectationFailure ("append failed: " <> show err)
    Right _ -> pure ()

runUntil :: KirokuStore -> SubscriptionName -> Int -> IO [GlobalPosition]
runUntil store subName targetCount = do
  ref <- newIORef ([] :: [GlobalPosition])
  countRef <- newTVarIO (0 :: Int)
  let handler' evt = do
        modifyIORef' ref ((evt ^. #globalPosition) :)
        n <- atomically $ do
          count <- readTVar countRef
          let count' = count + 1
          writeTVar countRef count'
          pure count'
        if n >= targetCount
          then pure Stop
          else pure Continue
      cfg =
        SubscriptionConfig
          { name = subName,
            target = AllStreams,
            handler = handler',
            batchSize = either (error . show) Prelude.id (mkBatchSize 10),
            queueCapacity = 1,
            overflowPolicy = DropSubscription,
            consumerGroup = Nothing,
            consumerGroupGuard = False,
            missingCheckpointPolicy = FromBeginning,
            targetBindingPolicy = AdoptUnbound,
            undecodableHandler = Nothing,
            handlerStallWarnAfter = Nothing,
            retryPolicy = defaultRetryPolicy,
            eventTypeFilter = AllEventTypes,
            selector = Nothing
          }
  handle <- subscribe store cfg
  result <- waitWithTimeout 10_000_000 handle
  case result of
    Left timeout -> expectationFailure timeout
    Right (Left err) -> expectationFailure ("Subscription failed: " <> show err)
    Right (Right ()) -> pure ()
  reverse <$> readIORef ref