SubscriptionRegistry.hs

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

-- | Tests for the central subscription-state registry (ExecPlan 45 / EP-1 of
-- MasterPlan 7).
--
-- These prove the registry's observable behaviour against a real migrated
-- PostgreSQL database:
--
--   * Several live subscriptions (including consumer-group members) appear in the
--     'Kiroku.Store.Subscription.subscriptionStates' snapshot with sensible state
--     labels, matching keys, and a valid FSM cursor, and a live handle's
--     'currentState' agrees with its snapshot entry (the single-source-of-truth
--     reshape).
--   * A subscription's key disappears from the snapshot when it stops cleanly, is
--     cancelled, or crashes — and its held handle's 'currentState' then returns
--     'Nothing'.
--   * A stale duplicate-key worker's cleanup cannot delete a newer worker's
--     replacement entry (the per-worker token defends the active entry).
module Test.SubscriptionRegistry (spec) where

import Control.Concurrent (threadDelay)
import Control.Concurrent.STM (readTVarIO)
import Control.Exception (throwIO)
import Control.Lens ((^.))
import Data.Aeson qualified as Aeson
import Data.Foldable (for_)
import Data.Generics.Labels ()
import Data.Int (Int32)
import Data.IntMap.Strict qualified as IntMap
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust, isNothing)
import Data.Text (Text)
import Kiroku.Store
import Kiroku.Store.Subscription.EventPublisher qualified as Pub
import Test.Helpers (makeEvent, validConsumerGroup, waitForPublisher, withTestStore)
import Test.Hspec

-- | A plain @$all@ subscription config whose handler never stops.
plainCont :: Text -> SubscriptionConfig
plainCont nm = defaultSubscriptionConfig (SubscriptionName nm) AllStreams (\_ -> pure Continue)

-- | A size-@n@ category-group config for member @m@ whose handler never stops.
groupCont :: Text -> Text -> Int32 -> Int32 -> SubscriptionConfig
groupCont nm cat m n =
  (defaultSubscriptionConfig (SubscriptionName nm) (Category (CategoryName cat)) (\_ -> pure Continue))
    { consumerGroup = Just (validConsumerGroup m n)
    }

-- | A size-@n@ @$all@ group config for member @m@ whose handler never stops.
groupAllCont :: Text -> Int32 -> Int32 -> SubscriptionConfig
groupAllCont nm m n =
  (defaultSubscriptionConfig (SubscriptionName nm) AllStreams (\_ -> pure Continue))
    { consumerGroup = Just (validConsumerGroup m n)
    }

publisherSubscriberCount :: KirokuStore -> IO Int
publisherSubscriberCount store =
  IntMap.size <$> readTVarIO (Pub.subscribers (store ^. #publisher))

-- | Poll the snapshot until the keyed entry reaches the given 'statePhase' label.
waitUntilPhase :: Int -> KirokuStore -> (SubscriptionName, Int32) -> Text -> IO Bool
waitUntilPhase budget store key phase
  | budget <= 0 = pure False
  | otherwise = do
      m <- subscriptionStates store
      case Map.lookup key m of
        Just v | (v ^. #statePhase) == phase -> pure True
        _ -> threadDelay 20_000 >> waitUntilPhase (budget - 20_000) store key phase

-- | Poll the snapshot until the keyed entry is present.
waitUntilPresent :: Int -> KirokuStore -> (SubscriptionName, Int32) -> IO Bool
waitUntilPresent budget store key
  | budget <= 0 = pure False
  | otherwise = do
      m <- subscriptionStates store
      if Map.member key m
        then pure True
        else threadDelay 20_000 >> waitUntilPresent (budget - 20_000) store key

-- | Poll the snapshot until the keyed entry is absent.
waitUntilAbsent :: Int -> KirokuStore -> (SubscriptionName, Int32) -> IO Bool
waitUntilAbsent budget store key
  | budget <= 0 = pure False
  | otherwise = do
      m <- subscriptionStates store
      if Map.member key m
        then threadDelay 20_000 >> waitUntilAbsent (budget - 20_000) store key
        else pure True

spec :: Spec
spec = describe "subscription registry (EP-1 / plan 45)" $ do
  it "registers every live subscription with a sensible state and position" $
    withTestStore $ \store -> do
      let plainName = SubscriptionName "reg-plain"
          groupName = SubscriptionName "reg-group"
      Right _ <- runStoreIO store $ appendToStream (StreamName "reg-1") NoStream [makeEvent "A" (Aeson.object [])]
      Right _ <- runStoreIO store $ appendToStream (StreamName "reg-2") NoStream [makeEvent "B" (Aeson.object [])]
      waitForPublisher store (GlobalPosition 2)
      hPlain <- subscribe store (plainCont "reg-plain")
      hG0 <- subscribe store (groupCont "reg-group" "reg" 0 2)
      hG1 <- subscribe store (groupCont "reg-group" "reg" 1 2)
      okP <- waitUntilPhase 5_000_000 store (plainName, 0) "live"
      ok0 <- waitUntilPhase 5_000_000 store (groupName, 0) "live"
      ok1 <- waitUntilPhase 5_000_000 store (groupName, 1) "live"
      (okP, ok0, ok1) `shouldBe` (True, True, True)
      snap <- subscriptionStates store
      Map.keys snap `shouldMatchList` [(plainName, 0), (groupName, 0), (groupName, 1)]
      for_ (Map.toList snap) $ \((nm, mbr), v) -> do
        (v ^. #subscriptionName) `shouldBe` nm
        (v ^. #member) `shouldBe` mbr
        (v ^. #statePhase) `shouldSatisfy` (`elem` (["catching_up", "live"] :: [Text]))
        let GlobalPosition p = v ^. #cursor
        p `shouldSatisfy` (>= 0)
      -- currentState on a live handle agrees with its snapshot entry.
      Just s <- currentState hPlain
      case Map.lookup (plainName, 0) snap of
        Just v -> stateName s `shouldBe` (v ^. #statePhase)
        Nothing -> expectationFailure "plain subscription missing from snapshot"
      mapM_ cancel [hPlain, hG0, hG1]

  it "removes a subscription's entry when it stops cleanly" $
    withTestStore $ \store -> do
      let nm = SubscriptionName "reg-stop"
      Right _ <- runStoreIO store $ appendToStream (StreamName "stop-1") NoStream [makeEvent "A" (Aeson.object [])]
      waitForPublisher store (GlobalPosition 1)
      handle <- subscribe store (defaultSubscriptionConfig nm AllStreams (\_ -> pure Stop))
      _ <- wait handle
      gone <- waitUntilAbsent 5_000_000 store (nm, 0)
      gone `shouldBe` True
      cs <- currentState handle
      cs `shouldSatisfy` isNothing

  it "removes a subscription's entry when it is cancelled" $
    withTestStore $ \store -> do
      let nm = SubscriptionName "reg-cancel"
      Right _ <- runStoreIO store $ appendToStream (StreamName "cancel-1") NoStream [makeEvent "A" (Aeson.object [])]
      waitForPublisher store (GlobalPosition 1)
      handle <- subscribe store (plainCont "reg-cancel")
      live <- waitUntilPhase 5_000_000 store (nm, 0) "live"
      live `shouldBe` True
      cancel handle
      _ <- wait handle
      gone <- waitUntilAbsent 5_000_000 store (nm, 0)
      gone `shouldBe` True
      cs <- currentState handle
      cs `shouldSatisfy` isNothing

  it "removes a subscription's entry when its handler crashes" $
    withTestStore $ \store -> do
      let nm = SubscriptionName "reg-crash"
      Right _ <- runStoreIO store $ appendToStream (StreamName "crash-1") NoStream [makeEvent "A" (Aeson.object [])]
      waitForPublisher store (GlobalPosition 1)
      handle <- subscribe store (defaultSubscriptionConfig nm AllStreams (\_ -> throwIO (userError "boom")))
      res <- wait handle
      case res of
        Left _ -> pure ()
        Right () -> expectationFailure "expected the crashing handler to fail the worker"
      gone <- waitUntilAbsent 5_000_000 store (nm, 0)
      gone `shouldBe` True
      cs <- currentState handle
      cs `shouldSatisfy` isNothing

  it "stale duplicate-key cleanup does not remove the replacement entry" $
    withTestStore $ \store -> do
      let nm = SubscriptionName "reg-dup"
      h1 <- subscribe store (plainCont "reg-dup")
      present1 <- waitUntilPresent 5_000_000 store (nm, 0)
      present1 `shouldBe` True
      -- A second subscribe with the same (name, member) supersedes h1's
      -- entry in the registry (insert is synchronous, before the fork).
      h2 <- subscribe store (plainCont "reg-dup")
      -- Stopping the first worker must not delete the second's live entry.
      cancel h1
      _ <- wait h1
      snap <- subscriptionStates store
      Map.member (nm, 0) snap `shouldBe` True
      cs1 <- currentState h1
      cs1 `shouldSatisfy` isNothing
      cs2 <- currentState h2
      cs2 `shouldSatisfy` isJust
      cancel h2

  describe "publisher queue registration" $ do
    it "registers publisher queues only for non-group AllStreams subscriptions" $
      withTestStore $ \store -> do
        initial <- publisherSubscriberCount store
        initial `shouldBe` 0

        hCategory <- subscribe store (defaultSubscriptionConfig (SubscriptionName "reg-cat-no-queue") (Category (CategoryName "regcat")) (\_ -> pure Continue))
        publisherSubscriberCount store `shouldReturn` 0

        hGroupCategory <- subscribe store (groupCont "reg-group-cat-no-queue" "regcat" 0 2)
        publisherSubscriberCount store `shouldReturn` 0

        hGroupAll <- subscribe store (groupAllCont "reg-group-all-no-queue" 0 2)
        publisherSubscriberCount store `shouldReturn` 0

        hAll <- subscribe store (plainCont "reg-all-has-queue")
        publisherSubscriberCount store `shouldReturn` 1

        cancel hAll
        _ <- wait hAll
        publisherSubscriberCount store `shouldReturn` 0

        cancel hCategory
        _ <- wait hCategory
        publisherSubscriberCount store `shouldReturn` 0

        cancel hGroupCategory
        _ <- wait hGroupCategory
        publisherSubscriberCount store `shouldReturn` 0

        cancel hGroupAll
        _ <- wait hGroupAll
        publisherSubscriberCount store `shouldReturn` 0