PublisherDropCounter.hs

{-# LANGUAGE OverloadedRecordDot #-}

module Test.PublisherDropCounter (spec) where

import Control.Concurrent.STM
import Control.Exception (bracket)
import Data.Aeson (object)
import Data.IntMap.Strict qualified as IntMap
import Data.Text (Text)
import Data.Vector qualified as V
import Kiroku.Store
import Kiroku.Store.Subscription.EventPublisher qualified as Pub
import System.Timeout (timeout)
import Test.Helpers (makeEvent, waitForPublisher, withTestStore)
import Test.Hspec

spec :: Spec
spec = describe "publisher drop counter" $ do
  it "counts drops, keeps the newest batch and preserves the idempotent wrapper" $ withTestStore $ \store -> do
    baseline <- IntMap.size <$> readTVarIO (Pub.subscribers store.publisher)
    bracket (atomically (Pub.subscribePublisherWith store.publisher 1 DropOldest)) Pub.unsubscribe $ \sub -> do
      appendAndWait store "drop-a" 1
      readTVarIO sub.subscriptionDropped `shouldReturn` 0
      appendAndWait store "drop-b" 2
      readTVarIO sub.subscriptionDropped `shouldReturn` 1
      readTVarIO sub.subscriptionStatus `shouldReturn` Pub.Active
      newest sub.subscriptionQueue `shouldReturn` [GlobalPosition 2]
      bracket (atomically (Pub.subscribePublisher store.publisher 1 DropOldest)) (\(_, _, stop) -> stop) $ \(queue, _, stop) -> do
        appendAndWait store "drop-c" 3
        appendAndWait store "drop-d" 4
        newest queue `shouldReturn` [GlobalPosition 4]
        stop >> stop
      Pub.unsubscribe sub >> Pub.unsubscribe sub
    (IntMap.size <$> readTVarIO (Pub.subscribers store.publisher)) `shouldReturn` baseline
  it "does not count PauseAndResume or DropSubscription overflow" $ withTestStore $ \store ->
    mapM_
      ( \policy -> bracket (atomically (Pub.subscribePublisherWith store.publisher 1 policy)) Pub.unsubscribe $ \sub -> do
          Right pos <- runStoreIO store (appendToStream (StreamName "policies-a") AnyVersion [makeEvent "E" (object [])])
          bounded (waitForPublisher store pos.globalPosition)
          Right pos2 <- runStoreIO store (appendToStream (StreamName "policies-a") AnyVersion [makeEvent "E" (object [])])
          bounded (waitForPublisher store pos2.globalPosition)
          readTVarIO sub.subscriptionDropped `shouldReturn` 0
          readTVarIO sub.subscriptionStatus `shouldReturn` (if policy == PauseAndResume then Pub.Paused else Pub.Overflowed)
      )
      [PauseAndResume, DropSubscription]

appendAndWait :: KirokuStore -> Text -> Int -> IO ()
appendAndWait store name n = do
  Right _ <- runStoreIO store (appendToStream (StreamName name) NoStream [makeEvent "E" (object [])])
  bounded (waitForPublisher store (GlobalPosition (fromIntegral n)))

bounded :: IO a -> IO a
bounded action = timeout 5_000_000 action >>= maybe (fail "publisher timeout") pure

newest :: TBQueue DecodedBatch -> IO [GlobalPosition]
newest queue = do
  Just (UnchangedBatch events) <- atomically (tryReadTBQueue queue)
  atomically (isEmptyTBQueue queue) `shouldReturn` True
  pure (map (.globalPosition) (V.toList events))