StartupFailureSurfacing.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.StartupFailureSurfacing (spec) where

import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM (atomically, newTVarIO, readTVarIO, writeTVar)
import Control.Exception (SomeException, fromException)
import Control.Lens ((&), (.~), (^.))
import Control.Monad (replicateM_)
import Data.Generics.Labels ()
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.IntMap.Strict qualified as IntMap
import Data.Map.Strict qualified as Map
import Hasql.Pool qualified as Pool
import Kiroku.Store
import Kiroku.Store.Subscription.EventPublisher qualified as Pub
import Kiroku.Store.Subscription.Worker (withLoadCheckpointHookForTest)
import Test.Helpers (waitWithTimeout, withTestStore, withTestStoreSettings)
import Test.Hspec

spec :: Spec
spec = describe "startup failure surfacing" $ do
  it "fails subscription startup loudly when checkpoint load fails" $ do
    handlerCalled <- newTVarIO False
    observedRef <- newIORef ([] :: [KirokuEvent])
    let subName = SubscriptionName "load-checkpoint-fails-loudly"
        cfg =
          defaultSubscriptionConfig subName AllStreams $ \_ -> do
            atomically (writeTVar handlerCalled True)
            pure Continue
        observe evt = modifyIORef' observedRef (evt :)
        injectLoadFailure _ = pure (Just (Left Pool.AcquisitionTimeoutUsageError))
        tweak settings = settings & #eventHandler .~ Just observe

    withTestStoreSettings tweak $ \store ->
      withLoadCheckpointHookForTest injectLoadFailure $ do
        handle <- subscribe store cfg
        result <- waitWithTimeout 5_000_000 handle
        case result of
          Right (Left e)
            | Just Pool.AcquisitionTimeoutUsageError <- fromException e -> pure ()
          Left timeout -> expectationFailure timeout
          Right other -> expectationFailure ("expected checkpoint load failure, got: " <> show other)

    wasCalled <- readTVarIO handlerCalled
    wasCalled `shouldBe` False
    observed <- reverse <$> readIORef observedRef
    any (isLoadCheckpointError subName) observed `shouldBe` True
    any (isCrashedStop subName) observed `shouldBe` True

  it "refuses a missing checkpoint before Started or handler delivery" $ do
    handlerCalled <- newTVarIO False
    observedRef <- newIORef ([] :: [KirokuEvent])
    let subName = SubscriptionName "fail-if-checkpoint-missing"
        key = SubscriptionCheckpointKey subName 0
        cfg =
          ( defaultSubscriptionConfig subName AllStreams $ \_ -> do
              atomically (writeTVar handlerCalled True)
              pure Continue
          )
            { missingCheckpointPolicy = FailIfMissing
            }
        observe evt = modifyIORef' observedRef (evt :)
        tweak settings = settings & #eventHandler .~ Just observe

    withTestStoreSettings tweak $ \store -> do
      handle <- subscribe store cfg
      result <- waitWithTimeout 5_000_000 handle
      case result of
        Right (Left e)
          | Just (SubscriptionCheckpointMissing actualKey) <- fromException e ->
              actualKey `shouldBe` key
        Left timeout -> expectationFailure timeout
        Right other -> expectationFailure ("expected missing-checkpoint refusal, got: " <> show other)

      runStoreIO store subscriptionCheckpointInventory
        `shouldReturn` Right (SubscriptionCheckpointInventory (GlobalPosition 0) mempty)

    readTVarIO handlerCalled `shouldReturn` False
    observed <- reverse <$> readIORef observedRef
    observed `shouldSatisfy` any (isMissingCheckpoint key)
    observed `shouldSatisfy` all (not . isStarted subName)

  it "leaves no publisher or subscription registry entries after a subscribe/cancel storm" $
    withTestStore $ \store -> do
      let cfg = defaultSubscriptionConfig (SubscriptionName "subscribe-cancel-storm") AllStreams (\_ -> pure Continue)
      -- 200 iterations gives async exceptions many chances to land in the
      -- pre-fork window without making the test expensive on normal runs.
      replicateM_ 200 $ do
        pendingAsync <- Async.async (subscribe store cfg)
        Async.cancel pendingAsync
        outcome <- Async.waitCatch pendingAsync
        case outcome of
          Right handle -> cancel handle >> (() <$ wait handle)
          Left (_ :: SomeException) -> pure ()

      subs <- readTVarIO (Pub.subscribers (store ^. #publisher))
      reg <- readTVarIO (store ^. #subscriptionRegistry)
      IntMap.null subs `shouldBe` True
      Map.null reg `shouldBe` True

isLoadCheckpointError :: SubscriptionName -> KirokuEvent -> Bool
isLoadCheckpointError expected = \case
  KirokuEventSubscriptionDbError actual LoadCheckpoint Pool.AcquisitionTimeoutUsageError _
    | actual == expected -> True
  _ -> False

isCrashedStop :: SubscriptionName -> KirokuEvent -> Bool
isCrashedStop expected = \case
  KirokuEventSubscriptionStopped actual _ (StopWorkerCrashed e) _
    | actual == expected,
      Just Pool.AcquisitionTimeoutUsageError <- fromException e ->
        True
  _ -> False

isMissingCheckpoint :: SubscriptionCheckpointKey -> KirokuEvent -> Bool
isMissingCheckpoint expected = \case
  KirokuEventSubscriptionCheckpointMissing (SubscriptionCheckpointMissing actual) _ -> actual == expected
  _ -> False

isStarted :: SubscriptionName -> KirokuEvent -> Bool
isStarted expected = \case
  KirokuEventSubscriptionStarted actual _ _ -> actual == expected
  _ -> False