kiroku-store-0.10.0.0: test/Test/ConsumerGroupResize.hs
{-# LANGUAGE MultilineStrings #-}
{-# LANGUAGE NumericUnderscores #-}
module Test.ConsumerGroupResize (spec) where
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM
import Control.Exception (SomeException, fromException)
import Control.Lens ((^.))
import Control.Monad (forM_, void)
import Data.Aeson qualified as Aeson
import Data.Generics.Labels ()
import Data.IORef
import Data.Int (Int32, Int64)
import Data.Set qualified as Set
import Data.Text (Text)
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 (preparable)
import Hasql.Transaction qualified as Tx
import Hasql.Transaction.Sessions qualified as TxSessions
import Kiroku.Store
import Kiroku.Store.SQL qualified as SQL
import System.Timeout (timeout)
import Test.Helpers (makeEvent, validConsumerGroup, waitWithTimeout, withTestStore, withTestStoreSettings)
import Test.Hspec
spec :: Spec
spec = describe "consumer-group resize" $ do
it "rejects invalid sizes and member indices at construction" $ do
mkConsumerGroupSize 0 `shouldBe` Left (InvalidConsumerGroup 0 0)
mkConsumerGroupSize (-1) `shouldBe` Left (InvalidConsumerGroup 0 (-1))
let size2 = groupSize 2
mkConsumerGroup (-1) size2 `shouldBe` Left (InvalidConsumerGroup (-1) 2)
mkConsumerGroup 2 size2 `shouldBe` Left (InvalidConsumerGroup 2 2)
mkConsumerGroup 1 size2 `shouldBe` Right (validConsumerGroup 1 2)
it "refuses a configured size that disagrees with stored topology before creating a member" $ do
refusals <- newIORef []
let observe = \case
KirokuEventSubscriptionGroupSizeMismatch refusal group -> modifyIORef' refusals ((refusal, group) :)
_ -> pure ()
withTestStoreSettings (\settings -> settings{eventHandler = Just observe}) $ \store -> do
seedCheckpoint store "mismatch" 0 2 5
called <- newIORef False
handle <- subscribe store $ config "mismatch" 2 3 $ \_ -> writeIORef called True >> pure Continue
expectMismatch handle "mismatch" 3 [2]
readIORef called `shouldReturn` False
rows store "mismatch" `shouldReturn` [(0, 2, 5)]
readIORef refusals `shouldReturn` [(ConsumerGroupSizeMismatch (SubscriptionName "mismatch") 3 (Vector.singleton 2), GroupMember 2 3)]
it "also refuses an ordinary subscription over an existing larger group" $
withTestStore $ \store -> do
seedCheckpoint store "ordinary-mismatch" 0 2 0
handle <- subscribe store (defaultSubscriptionConfig (SubscriptionName "ordinary-mismatch") AllStreams (\_ -> pure Continue))
expectMismatch handle "ordinary-mismatch" 1 [2]
it "refuses an underestimated derived topology until resized" $
withTestStore $ \store -> do
-- The upgrade test proves a lone legacy member derives size 1.
seedCheckpoint store "underestimated" 0 1 8
refused <- subscribe store (config "underestimated" 0 2 (\_ -> pure Continue))
expectMismatch refused "underestimated" 2 [1]
void $ resize store "underestimated" 2
accepted <- subscribe store (config "underestimated" 0 2 (\_ -> pure Continue))
awaitLive accepted
cancel accepted
rows store "underestimated" `shouldReturn` [(0, 2, 8), (1, 2, 8)]
it "refuses mixed stored sizes even when its own row agrees" $
withTestStore $ \store -> do
seedCheckpoint store "mixed" 0 2 0
seedCheckpoint store "mixed" 1 3 0
handle <- subscribe store (config "mixed" 0 2 (\_ -> pure Continue))
expectMismatch handle "mixed" 2 [2, 3]
it "serializes concurrent initializers with competing group sizes" $
withTestStore $ \store -> do
[firstHandle, secondHandle] <- Async.mapConcurrently (\n -> subscribe store (config "concurrent" 0 n (\_ -> pure Continue))) [2, 3]
-- The winner reaches Live; the loser exits with a typed refusal.
result <- timeout 5_000_000 $ Async.race (wait firstHandle) (wait secondHandle)
case result of
Just (Left outcome) -> assertMismatch outcome
Just (Right outcome) -> assertMismatch outcome
Nothing -> expectationFailure "competing topology startup timed out"
mapM_ cancel [firstHandle, secondHandle]
stored <- rows store "concurrent"
length stored `shouldBe` 1
stored `shouldSatisfy` \case [(0, n, 0)] -> n == 2 || n == 3; _ -> False
it "persists configured topology on initialization, ordinary saves and dead-letter saves" $
withTestStore $ \store -> do
void $ runStoreIO store $ appendToStream (StreamName "persist-1") NoStream [makeEvent "E" (Aeson.object [])]
-- A size-1 group exercises all rows deterministically. Size >1 is
-- checked below via explicit ordinary and dead-letter SQL paths.
handle <- subscribe store (config "saved" 0 1 (\_ -> pure Stop))
expectClean handle
rows store "saved" `shouldReturn` [(0, 1, 1)]
live <- subscribe store ((config "initialized" 2 3 (\_ -> pure Continue)){missingCheckpointPolicy = FromCurrentHead})
awaitLive live
cancel live
rows store "initialized" `shouldReturn` [(2, 3, 1)]
seedCheckpoint store "saved-size" 1 4 20
seedCheckpoint store "saved-size" 1 4 10
rows store "saved-size" `shouldReturn` [(1, 4, 20)]
Right events <- runStoreIO store (readAllForward (GlobalPosition 0) 10)
let event = Vector.head events
EventId eid = event ^. #eventId
params = SQL.DeadLetterParams "dead-letter-size" 1 "unbound" Nothing 4 1 eid (Aeson.object []) "test" 1
runSession store (Session.statement params SQL.insertDeadLetterAndCheckpointStmt)
rows store "dead-letter-size" `shouldReturn` [(1, 4, 1)]
let StreamId sid = event ^. #originalStreamId
ownerStmt =
preparable
"SELECT (((hashtextextended($1::bigint::text, 0) % 4) + 4) % 4)::int4"
(E.param (E.nonNullable E.int8))
(D.singleRow (D.column (D.nonNullable D.int4)))
owner <- runSession store (Session.statement sid ownerStmt)
saved <- subscribe store (config "worker-save-size" owner 4 (\_ -> pure Stop))
expectClean saved
rows store "worker-save-size" `shouldReturn` [(owner, 4, 1)]
deadLettered <- subscribe store (config "worker-dead-letter-size" owner 4 (\_ -> pure (DeadLetter (DeadLetterPoison "explicit consumer decision"))))
awaitLive deadLettered
cancel deadLettered
rows store "worker-dead-letter-size" `shouldReturn` [(owner, 4, 1)]
it "delivers every seeded event after equalizing size 2 to size 3" $
withTestStore $ \store -> do
forM_ [1 .. 40 :: Int] $ \i -> do
result <- runStoreIO store $ appendToStream (StreamName ("resize-" <> Text.pack (show i))) NoStream [makeEvent "E" (Aeson.object [])]
result `shouldSatisfy` either (const False) (const True)
Right events <- runStoreIO store (readAllForward (GlobalPosition 0) 100)
let expected = Set.fromList [event ^. #eventId | event <- Vector.toList events]
-- Member 0 has seen nothing, member 1 has crossed the full log.
-- Remember its actual delivered set; a naive size change loses
-- streams that move from old member 0 to new member 1.
oldFast <- runSession store (Session.statement (0, 1, 2, 100) SQL.readAllForwardConsumerGroupStmt)
naive <- mapM (\(m, p) -> runSession store (Session.statement (p, m, 3, 100) SQL.readAllForwardConsumerGroupStmt)) [(0, 0), (1, 40), (2, 0)]
let alreadySeen = Set.fromList [e ^. #eventId | e <- Vector.toList oldFast]
naiveSeen = Set.fromList [e ^. #eventId | batch <- naive, e <- Vector.toList batch]
Set.union alreadySeen naiveSeen `shouldNotBe` expected
seedCheckpoint store "skewed" 0 2 0
seedCheckpoint store "skewed" 1 2 40
refused <- subscribe store (config "skewed" 1 3 (\_ -> pure Continue))
expectMismatch refused "skewed" 3 [2]
report <- resize store "skewed" 3
report `shouldBe` ConsumerGroupResizeReport (Vector.singleton 2) 2 (groupSize 3) (GlobalPosition 0)
delivered <- newTVarIO alreadySeen
handles <- mapM (\m -> subscribe store (config "skewed" m 3 (\e -> atomically (modifyTVar' delivered (Set.insert (e ^. #eventId))) >> pure Continue))) [0, 1, 2]
complete <- timeout 15_000_000 (atomically (readTVar delivered >>= check . (== expected)))
mapM_ cancel handles
complete `shouldBe` Just ()
readTVarIO delivered `shouldReturn` expected
it "is idempotent when repeated at the same size and retains surviving member identities" $
withTestStore $ \store -> do
seedCheckpoint store "repeat" 0 2 5
seedCheckpoint store "repeat" 1 2 20
oldIds <- memberIds store "repeat"
void $ resize store "repeat" 3
originalRows <- rows store "repeat"
resizedIds <- memberIds store "repeat"
take 2 resizedIds `shouldBe` oldIds
second <- resize store "repeat" 3
rows store "repeat" `shouldReturn` originalRows
memberIds store "repeat" `shouldReturn` resizedIds
originalRows `shouldBe` [(0, 3, 5), (1, 3, 5), (2, 3, 5)]
resumePosition second `shouldBe` GlobalPosition 5
previousMemberCount second `shouldBe` 3
it "starts a missing group at zero and removes obsolete members on shrink" $
withTestStore $ \store -> do
report <- resize store "missing" 3
previousMemberCount report `shouldBe` 0
resumePosition report `shouldBe` GlobalPosition 0
rows store "missing" `shouldReturn` [(0, 3, 0), (1, 3, 0), (2, 3, 0)]
void $ resize store "missing" 1
rows store "missing" `shouldReturn` [(0, 1, 0)]
it "rolls back the complete resize with the caller transaction" $
withTestStore $ \store -> do
seedCheckpoint store "rollback" 0 2 5
seedCheckpoint store "rollback" 1 2 20
void $ runSession store $ TxSessions.transaction TxSessions.ReadCommitted TxSessions.Write $ do
report <- resizeConsumerGroupTx (SubscriptionName "rollback") (groupSize 3)
Tx.condemn
pure report
rows store "rollback" `shouldReturn` [(0, 2, 5), (1, 2, 20)]
groupSize :: Int32 -> ConsumerGroupSize
groupSize = either (error . show) Prelude.id . mkConsumerGroupSize
config :: Text -> Int32 -> Int32 -> EventHandler -> SubscriptionConfig
config name m n handler =
(defaultSubscriptionConfig (SubscriptionName name) AllStreams handler)
{ consumerGroup = Just (validConsumerGroup m n)
}
seedCheckpoint :: KirokuStore -> Text -> Int32 -> Int32 -> Int64 -> IO ()
seedCheckpoint store name m n p = runSession store (Session.statement (name, m, p, n, "unbound", Nothing) SQL.saveCheckpointMemberStmt)
resize :: KirokuStore -> Text -> Int32 -> IO ConsumerGroupResizeReport
resize store name n = runSession store $ TxSessions.transaction TxSessions.ReadCommitted TxSessions.Write (resizeConsumerGroupTx (SubscriptionName name) (groupSize n))
rows :: KirokuStore -> Text -> IO [(Int32, Int32, Int64)]
rows store name = Vector.toList <$> runSession store (Session.statement name stmt)
where
stmt =
preparable
"SELECT consumer_group_member, consumer_group_size, last_seen FROM subscriptions WHERE subscription_name = $1 ORDER BY consumer_group_member"
(E.param (E.nonNullable E.text))
(D.rowVector ((,,) <$> D.column (D.nonNullable D.int4) <*> D.column (D.nonNullable D.int4) <*> D.column (D.nonNullable D.int8)))
runSession :: KirokuStore -> Session.Session a -> IO a
runSession store session = Pool.use (store ^. #pool) session >>= either (error . show) pure
expectMismatch :: SubscriptionHandle -> Text -> Int32 -> [Int32] -> IO ()
expectMismatch handle name n observed = do
result <- waitWithTimeout 5_000_000 handle
case result of
Right (Left exception) -> fromException exception `shouldBe` Just (ConsumerGroupSizeMismatch (SubscriptionName name) n (Vector.fromList observed))
other -> expectationFailure ("expected topology mismatch, got " <> show other)
assertMismatch :: Either SomeException () -> IO ()
assertMismatch = \case
Left exception -> (fromException exception :: Maybe ConsumerGroupSizeMismatch) `shouldSatisfy` maybe False (const True)
Right () -> expectationFailure "expected competing-size startup refusal"
expectClean :: SubscriptionHandle -> IO ()
expectClean handle =
waitWithTimeout 5_000_000 handle >>= \case
Right (Right ()) -> pure ()
other -> expectationFailure (show other)
awaitLive :: SubscriptionHandle -> IO ()
awaitLive handle = do
result <- timeout 5_000_000 loop
result `shouldBe` Just ()
where
loop =
currentState handle >>= \case
Just state | stateName state == "live" -> pure ()
_ -> threadDelay 1_000 >> loop
memberIds :: KirokuStore -> Text -> IO [Int64]
memberIds store name =
runSession store $
Session.statement name $
preparable
"SELECT subscription_id FROM subscriptions WHERE subscription_name = $1 ORDER BY consumer_group_member"
(E.param (E.nonNullable E.text))
(D.rowList (D.column (D.nonNullable D.int8)))