StreamHistoryGuard.hs

{-# LANGUAGE MultilineStrings #-}
{-# LANGUAGE NumericUnderscores #-}

module Test.StreamHistoryGuard (spec) where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async qualified as Async
import Control.Lens ((^.))
import Control.Monad (forM_, when)
import Data.Aeson qualified as Aeson
import Data.Text qualified as Text
import Data.Vector qualified as Vector
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Hasql.Statement (Statement, preparable)
import Hasql.Transaction qualified as Tx
import Kiroku.Store
import Test.Helpers (makeEvent, withTestStore)
import Test.Hspec

spec :: Spec
spec = describe "stream history guard" $ do
  it "returns exact metadata and pages with production cursor semantics" $
    withTestStore $ \store -> do
      let stream = StreamName "guard-metadata"
      Right _ <- runStoreIO store $ appendToStream stream NoStream (events 3)
      Right _ <- runStoreIO store $ setStreamTruncateBefore stream (StreamVersion 2)
      result <- runTx store $ do
        locked <- lockStreamHistoryForReplayTx stream
        page <- readStreamForwardTx stream (StreamVersion 0) 10
        pure (locked, page)
      case result of
        (Right info, page) -> do
          info ^. #version `shouldBe` StreamVersion 3
          info ^. #truncateBefore `shouldBe` StreamVersion 2
          info ^. #deletedAt `shouldBe` Nothing
          fmap (^. #streamVersion) (Vector.toList page)
            `shouldBe` [StreamVersion 2, StreamVersion 3]
        (Left unavailable, _) -> expectationFailure ("guard unexpectedly unavailable: " <> show unavailable)

  it "returns typed unavailable results for missing and reserved streams" $
    withTestStore $ \store -> do
      missing <- runTx store (lockStreamHistoryForReplayTx (StreamName "guard-missing"))
      missing `shouldBe` Left (StreamHistoryNotFound (StreamName "guard-missing"))
      reserved <- runTx store (lockStreamHistoryForReplayTx (StreamName "$all"))
      reserved `shouldBe` Left (StreamHistoryReserved (StreamName "$all"))

  it "returns soft-deleted metadata and blocks undelete until guard completion" $
    withTestStore $ \store -> do
      let stream = StreamName "guard-undelete"
      Right _ <- runStoreIO store $ appendToStream stream NoStream (events 1)
      Right _ <- runStoreIO store $ softDeleteStream stream
      assertBlockedByGuard store stream $ do
        info <- runTx store (lockStreamHistoryForReplayTx stream)
        info `shouldSatisfy` \case Right StreamInfo {deletedAt = Just _} -> True; _ -> False
        runStoreIO store (undeleteStream stream)

  it "blocks append until guard completion" $
    withTestStore $ \store -> do
      let stream = StreamName "guard-append"
      Right _ <- runStoreIO store $ appendToStream stream NoStream (events 1)
      assertBlockedByGuard store stream $
        runStoreIO store (appendToStream stream AnyVersion (events 1))

  it "releases the guard on rollback and lets the blocked append complete" $
    withTestStore $ \store -> do
      let stream = StreamName "guard-rollback"
      Right _ <- runStoreIO store $ appendToStream stream NoStream (events 1)
      assertBlockedByGuardEnding True store stream $
        runStoreIO store (appendToStream stream AnyVersion (events 1))

  it "blocks link into the guarded stream until guard completion" $
    withTestStore $ \store -> do
      let source = StreamName "guard-link-source"
          target = StreamName "guard-link-target"
      Right _ <- runStoreIO store $ appendToStream source NoStream (events 1)
      Right _ <- runStoreIO store $ appendToStream target NoStream (events 1)
      Right sourceRows <- runStoreIO store $ readStreamForward source (StreamVersion 0) 10
      let eventId = Vector.head sourceRows ^. #eventId
      assertBlockedByGuard store target $
        runStoreIO store (linkToStream target [eventId])

  it "blocks soft delete and logical truncate until guard completion" $
    withTestStore $ \store -> do
      let soft = StreamName "guard-soft-delete"
          truncated = StreamName "guard-truncate"
      Right _ <- runStoreIO store $ appendToStream soft NoStream (events 1)
      Right _ <- runStoreIO store $ appendToStream truncated NoStream (events 2)
      assertBlockedByGuard store soft $
        runStoreIO store (softDeleteStream soft)
      assertBlockedByGuard store truncated $
        runStoreIO store (setStreamTruncateBefore truncated (StreamVersion 2))

  it "blocks hard delete of the guarded origin until guard completion" $
    withTestStore $ \store -> do
      let stream = StreamName "guard-hard-delete"
      Right _ <- runStoreIO store $ appendToStream stream NoStream (events 1)
      assertBlockedByGuard store stream $
        runStoreIO store (hardDeleteStream stream)

  it "blocks hard delete of another origin linked into the guarded stream" $
    withTestStore $ \store -> do
      let origin = StreamName "guard-linked-origin"
          guarded = StreamName "guard-linked-target"
      Right _ <- runStoreIO store $ appendToStream origin NoStream (events 1)
      Right _ <- runStoreIO store $ appendToStream guarded NoStream (events 1)
      Right sourceRows <- runStoreIO store $ readStreamForward origin (StreamVersion 0) 10
      Right _ <- runStoreIO store $ linkToStream guarded [Vector.head sourceRows ^. #eventId]
      assertBlockedByGuard store guarded $
        runStoreIO store (hardDeleteStream origin)

  it "keeps opposing hard deletes and multi-stream appends deadlock-free" $
    withTestStore $ \store ->
      forM_ [1 .. 5 :: Int] $ \index -> do
        let suffix = Text.pack (show index)
            first = StreamName ("guard-race-a-" <> suffix)
            second = StreamName ("guard-race-b-" <> suffix)
        Right _ <- runStoreIO store $ appendToStream first NoStream (events 1)
        Right _ <- runStoreIO store $ appendToStream second NoStream (events 1)
        operations <-
          Async.async $
            Async.concurrently
              (runStoreIO store $ hardDeleteStream first)
              ( runStoreIO store $
                  appendMultiStream
                    [ (first, AnyVersion, events 1),
                      (second, AnyVersion, events 1)
                    ]
              )
        _ <- waitWithin "hard-delete/multi-append race" operations
        pure ()

assertBlockedByGuard :: KirokuStore -> StreamName -> IO result -> IO ()
assertBlockedByGuard = assertBlockedByGuardEnding False

assertBlockedByGuardEnding :: Bool -> KirokuStore -> StreamName -> IO result -> IO ()
assertBlockedByGuardEnding rollBack store stream mutation = do
  guard <- Async.async $ runStoreIO store $ runTransaction $ do
    locked <- lockStreamHistoryForReplayTx stream
    _ <- Tx.statement () holdGuardStmt
    when rollBack Tx.condemn
    pure locked
  waitForGuardPhase store 100
  waiter <- Async.async mutation
  threadDelay 50_000
  Async.poll waiter >>= \case
    Nothing -> pure ()
    Just _ -> expectationFailure "mutation completed while the stream-history guard was held"
  guardResult <- waitWithin "guard transaction" guard
  guardResult `shouldSatisfy` \case Right (Right _) -> True; _ -> False
  _ <- waitWithin "blocked mutation" waiter
  pure ()

waitForGuardPhase :: KirokuStore -> Int -> IO ()
waitForGuardPhase _ 0 = expectationFailure "guard never reached its held phase"
waitForGuardPhase store attempts = do
  result <- Pool.use (store ^. #pool) (Session.statement () guardActiveStmt)
  case result of
    Right True -> pure ()
    Right False -> threadDelay 10_000 >> waitForGuardPhase store (attempts - 1)
    Left err -> expectationFailure ("could not observe guard phase: " <> show err)

waitWithin :: String -> Async.Async value -> IO value
waitWithin label action = do
  result <- Async.race (threadDelay 2_000_000) (Async.wait action)
  case result of
    Left () -> do
      Async.cancel action
      expectationFailure (label <> " timed out")
      error "unreachable"
    Right value -> pure value

runTx :: KirokuStore -> Tx.Transaction value -> IO value
runTx store transaction = do
  result <- runStoreIO store (runTransaction transaction)
  case result of
    Left err -> expectationFailure ("guard transaction failed: " <> show err) >> error "unreachable"
    Right value -> pure value

events :: Int -> [EventData]
events count =
  [makeEvent "Guarded" (Aeson.object [("n", Aeson.toJSON n)]) | n <- [1 .. count]]

holdGuardStmt :: Statement () Bool
holdGuardStmt =
  preparable
    "SELECT pg_sleep(0.4) IS NULL"
    E.noParams
    (D.singleRow (D.column (D.nonNullable D.bool)))

guardActiveStmt :: Statement () Bool
guardActiveStmt =
  preparable
    """
    SELECT EXISTS (
      SELECT 1
      FROM pg_stat_activity
      WHERE state = 'active'
        AND query = 'SELECT pg_sleep(0.4) IS NULL'
    )
    """
    E.noParams
    (D.singleRow (D.column (D.nonNullable D.bool)))