BatchProcessorSpec.hs

module Shibuya.Runner.BatchProcessorSpec (spec) where

import Control.Concurrent.NQE.Supervisor (Strategy (..))
import Control.Concurrent.STM (atomically, check, modifyTVar', newTVarIO, readTVar, readTVarIO, writeTVar)
import Control.Monad (forM_)
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef)
import Data.List (sort)
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NE
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Time (UTCTime (..), fromGregorian)
import Effectful (Eff, IOE, liftIO, runEff, (:>))
import Shibuya (ProcessorFailure (..), ProcessorHalt (..))
import Shibuya.Adapter (Adapter (..))
import Shibuya.Adapter.Mock
  ( TrackingAck (..),
    getTrackedDecisions,
    newTrackingAck,
    trackingAckHandle,
  )
import Shibuya.Batch
  ( BatchConfig (..),
    BatchHandler,
    BatchInfo (..),
    BatchKey (..),
    BatchTrigger (..),
    ackAllOk,
    ackExcept,
    defaultBatchConfig,
    defaultBatchKey,
    withFallback,
  )
import Shibuya.Core.Ack
  ( AckDecision (..),
    DeadLetterCode,
    DeadLetterReason (..),
    HaltReason (..),
    RetryDelay (..),
    mkDeadLetterCode,
  )
import Shibuya.Core.AckHandle (AckHandle (..))
import Shibuya.Core.Ingested (DeliveryStatus (..), Ingested (..), Message (..), mkIngested)
import Shibuya.Core.Metrics
  ( BatchStats (..),
    ProcessorId (..),
    ProcessorMetrics (..),
    StreamStats (..),
  )
import Shibuya.Core.Metrics qualified as Metrics
import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope)
import Shibuya.Internal.Runner.BatchProcessor (runBatchesWithMetrics)
import Shibuya.Internal.Runner.Master (startMaster, stopMaster)
import Shibuya.Internal.Runner.Supervised (SupervisedProcessor (..), runSupervisedBatchWith)
import Shibuya.Policy (Concurrency (..))
import Shibuya.Telemetry.Effect (Tracing, runTracingNoop)
import Streamly.Data.Stream qualified as Stream
import Test.Hspec
import UnliftIO (throwIO, try)
import UnliftIO.Concurrent (threadDelay)

spec :: Spec
spec = describe "Shibuya.Internal.Runner.BatchProcessor" $ do
  describe "delivery status" $ do
    forM_ [Metrics.MetricsBasic, Metrics.MetricsDetailed Metrics.defaultDetailOptions] $ \detail -> do
      forM_ [[1, 3, 5], [1 .. 5], []] $ \retained ->
        it ("samples every status once and retains current members " <> show retained <> " / " <> show detail) $ do
          calls <- newIORef []
          handled <- newIORef []
          (tracked, metrics) <- runEff $ runTracingNoop $ do
            tracking <- newTrackingAck
            (info, batch) <- buildBatch tracking 5 TriggerSize
            let members =
                  NE.zipWith
                    ( \n item ->
                        item
                          { envelope = item.envelope {cursor = Just (CursorInt (fromIntegral n))},
                            deliveryStatus = liftIO $ do
                              previous <- atomicModifyIORef' calls (\xs -> (xs <> [n], xs))
                              -- A superseded status flips on a second sample; two reads
                              -- would wrongly re-admit it and fail both assertions.
                              pure $ if n `elem` retained || n `elem` previous then DeliveryCurrent else DeliverySuperseded
                          }
                    )
                    (NE.fromList [1 .. 5 :: Int])
                    batch
                handler actual messages = do
                  liftIO $ modifyIORef' handled (<> [(actual, map (.envelope.messageId) (NE.toList messages))])
                  pure ackAllOk
            m <- runStatusBatches detail handler [(info, members)]
            t <- getTrackedDecisions tracking
            pure (t, m)
          readIORef calls `shouldReturn` [1 .. 5]
          map fst (reverse tracked) `shouldBe` [MessageId ("msg-" <> Text.pack (show n)) | n <- retained]
          observed <- readIORef handled
          case retained of
            [] -> do
              observed `shouldBe` []
              metrics.batch.batchesEmitted `shouldBe` 0
            _ -> case observed of
              [(actual, ids)] -> do
                actual.size `shouldBe` length retained
                actual.batchKey `shouldBe` defaultBatchKey
                actual.trigger `shouldBe` TriggerSize
                actual.partition `shouldBe` Nothing
                ids `shouldBe` map fst (reverse tracked)
              _ -> expectationFailure (show observed)
          metrics.stats.processed `shouldBe` length retained
          metrics.stats.failed `shouldBe` 0
          metrics.oldestInFlight `shouldBe` Nothing
          case detail of
            Metrics.MetricsBasic -> metrics.latency `shouldBe` Nothing
            _ -> do
              fmap (.count) metrics.latency `shouldBe` (if null retained then Nothing else Just 1)
              (metrics.progress >>= (.latest)) `shouldBe` (if null retained then Nothing else Just (CursorInt (fromIntegral (last retained))))

    it "checks a later batch after the previous batch finalizes" $ do
      current <- newIORef True
      handled <- newIORef []
      finalized <- newIORef []
      metrics <- runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        (info, first) <- buildBatch tracking 1 TriggerSize
        let a = (NE.head first) {ack = AckHandle (\_ -> liftIO $ atomicModifyIORef' current (const (False, ())))}
            b =
              (NE.head first)
                { deliveryStatus = liftIO $ (\x -> if x then DeliveryCurrent else DeliverySuperseded) <$> readIORef current,
                  ack = AckHandle (\_ -> liftIO $ modifyIORef' finalized (<> [2 :: Int]))
                }
            handler _ _ = liftIO (modifyIORef' handled (<> [1 :: Int])) >> pure ackAllOk
        runBatchesWithMetrics (ProcessorId "status-later") Serial handler [(info, NE.singleton a), (info, NE.singleton b)]
      readIORef handled `shouldReturn` [1]
      readIORef finalized `shouldReturn` []
      metrics.batch.batchesEmitted `shouldBe` 1

    it "lets throwing status escape as infrastructure failure" $ do
      handlerCalls <- newIORef (0 :: Int)
      finalized <- newIORef (0 :: Int)
      result <- try @IO @IOError $ runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        (info, batch) <- buildBatch tracking 1 TriggerSize
        let item = (NE.head batch) {deliveryStatus = liftIO $ ioError (userError "batch status infrastructure failure"), ack = AckHandle (\_ -> liftIO $ modifyIORef' finalized (+ 1))}
        runBatchesWithMetrics (ProcessorId "status-throw") Serial (\_ _ -> liftIO (modifyIORef' handlerCalls (+ 1)) >> pure ackAllOk) [(info, NE.singleton item)]
      case result of
        Left _ -> pure ()
        Right _ -> expectationFailure "expected infrastructure failure"
      readIORef handlerCalls `shouldReturn` 0
      readIORef finalized `shouldReturn` 0

  describe "decision resolution and finalization (M1)" $ do
    it "finalizes each of 5 messages successfully with per-message decisions" $ do
      (tracked, metrics) <- runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        batch <- buildBatch tracking 5 TriggerSize
        let handler _info _msgs =
              pure $
                ackExcept
                  [ (MessageId "msg-2", AckDeadLetter MaxRetriesExceeded),
                    (MessageId "msg-4", AckDeadLetter (PoisonPill "x"))
                  ]
        m <- runBatchesWithMetrics (ProcessorId "m1") Serial handler [batch]
        t <- getTrackedDecisions tracking
        pure (t, m)

      -- Each id appears exactly once in the successful-finalization list.
      sort (map fst tracked) `shouldBe` expectedIds
      lookup (MessageId "msg-1") tracked `shouldBe` Just AckOk
      lookup (MessageId "msg-2") tracked `shouldBe` Just (AckDeadLetter MaxRetriesExceeded)
      lookup (MessageId "msg-3") tracked `shouldBe` Just AckOk
      lookup (MessageId "msg-4") tracked `shouldBe` Just (AckDeadLetter (PoisonPill "x"))
      lookup (MessageId "msg-5") tracked `shouldBe` Just AckOk

      metrics.batch.batchesEmitted `shouldBe` 1
      metrics.batch.batchedMessages `shouldBe` 5
      metrics.batch.partialFailures `shouldBe` 1
      metrics.batch.sizeTriggered `shouldBe` 1
      metrics.stats.processed `shouldBe` 3
      metrics.stats.failed `shouldBe` 2

    it "retries a transient finalizer failure and records one successful finalization" $ do
      tracked <- runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        counter <- liftIO $ newIORef (0 :: Int)
        -- A single-message batch whose finalizer throws its first two attempts
        -- then records success on the third.
        let mid = MessageId "flaky-1"
            ing =
              mkIngested (mkEnv 1 mid) (flakyAckHandle counter tracking mid)
            info =
              BatchInfo
                { batchKey = defaultBatchKey,
                  size = 1,
                  trigger = TriggerSize,
                  partition = Nothing
                }
        _ <- runBatchesWithMetrics (ProcessorId "m1-flaky") Serial (\_ _ -> pure ackAllOk) [(info, NE.fromList [ing])]
        getTrackedDecisions tracking

      -- Exactly one successful finalization despite two prior throws.
      tracked `shouldBe` [(MessageId "flaky-1", AckOk)]

    it "preserves application failures in explicit and fallback decisions" $ do
      let code = validDeadLetterCode "keiro.router.selection.recipient_overflow"
          fallbackReason = ApplicationFailure code "fallback policy rejection"
          explicitReason = ApplicationFailure code "explicit policy rejection"
      (tracked, metrics) <- runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        batch <- buildBatch tracking 5 TriggerSize
        let handler _info _msgs =
              pure $
                withFallback
                  (AckDeadLetter fallbackReason)
                  [ (MessageId "msg-2", AckDeadLetter explicitReason),
                    (MessageId "msg-4", AckOk)
                  ]
        m <- runBatchesWithMetrics (ProcessorId "m1-application-failure") Serial handler [batch]
        t <- getTrackedDecisions tracking
        pure (t, m)

      sort (map fst tracked) `shouldBe` expectedIds
      lookup (MessageId "msg-1") tracked `shouldBe` Just (AckDeadLetter fallbackReason)
      lookup (MessageId "msg-2") tracked `shouldBe` Just (AckDeadLetter explicitReason)
      lookup (MessageId "msg-3") tracked `shouldBe` Just (AckDeadLetter fallbackReason)
      lookup (MessageId "msg-4") tracked `shouldBe` Just AckOk
      lookup (MessageId "msg-5") tracked `shouldBe` Just (AckDeadLetter fallbackReason)
      metrics.batch.partialFailures `shouldBe` 1
      metrics.stats.processed `shouldBe` 1
      metrics.stats.failed `shouldBe` 4

  describe "exception fallback (M2)" $ do
    it "finalizes all 5 with AckRetry when the handler throws" $ do
      (tracked, metrics) <- runEff $ runTracingNoop $ do
        tracking <- newTrackingAck
        batch <- buildBatch tracking 5 TriggerTimeout
        let handler _info _msgs = error "boom"
        m <- runBatchesWithMetrics (ProcessorId "m2-exc") Serial handler [batch]
        t <- getTrackedDecisions tracking
        pure (t, m)

      length tracked `shouldBe` 5
      sort (map fst tracked) `shouldBe` expectedIds
      all ((== AckRetry (RetryDelay 0)) . snd) tracked `shouldBe` True
      metrics.stats.failed `shouldBe` 5
      metrics.batch.partialFailures `shouldBe` 0
      metrics.batch.timeoutTriggered `shouldBe` 1

  describe "halt (M2)" $ do
    it "finalizes all 5, sets halt, and the driver throws ProcessorHalt" $ do
      -- Build the tracking ack OUTSIDE the aborted action so it survives the throw.
      tracking <- runEff $ runTracingNoop newTrackingAck
      result <- try $ runEff $ runTracingNoop $ do
        batch <- buildBatch tracking 5 TriggerFlush
        let handler _info _msgs =
              pure $
                withFallback
                  AckOk
                  [(MessageId "msg-3", AckHalt (HaltFatal "halt on 3"))]
        _ <- runBatchesWithMetrics (ProcessorId "m2-halt") Serial handler [batch]
        pure ()

      case result of
        Left (ProcessorHalt reason) -> reason `shouldBe` HaltFatal "halt on 3"
        Right () -> expectationFailure "expected ProcessorHalt"

      tracked <- runEff $ runTracingNoop $ getTrackedDecisions tracking
      sort (map fst tracked) `shouldBe` expectedIds
      lookup (MessageId "msg-3") tracked `shouldBe` Just (AckHalt (HaltFatal "halt on 3"))

    it "skips already-ready batches after halt and finalizes them with retry" $ do
      tracking <- runEff $ runTracingNoop newTrackingAck
      observedRef <- newIORef ([] :: [[MessageId]])
      result <- try $ runEff $ runTracingNoop $ do
        first <- buildRangeBatch tracking [1, 2] TriggerSize
        second <- buildRangeBatch tracking [3, 4, 5] TriggerFlush
        let handler _info msgs = do
              let ids = [ing.envelope.messageId | ing <- NE.toList msgs]
              liftIO $ modifyIORef' observedRef (ids :)
              if MessageId "msg-2" `elem` ids
                then pure $ withFallback AckOk [(MessageId "msg-2", AckHalt (HaltFatal "halt on 2"))]
                else pure ackAllOk
        _ <- runBatchesWithMetrics (ProcessorId "m3-halt-skip") Serial handler [first, second]
        pure ()

      case result of
        Left (ProcessorHalt reason) -> reason `shouldBe` HaltFatal "halt on 2"
        Right () -> expectationFailure "expected ProcessorHalt"

      observed <- readIORef observedRef
      observed `shouldBe` [[MessageId "msg-1", MessageId "msg-2"]]

      tracked <- runEff $ runTracingNoop $ getTrackedDecisions tracking
      lookup (MessageId "msg-1") tracked `shouldBe` Just AckOk
      lookup (MessageId "msg-2") tracked `shouldBe` Just (AckHalt (HaltFatal "halt on 2"))
      lookup (MessageId "msg-3") tracked `shouldBe` Just (AckRetry (RetryDelay 0))
      lookup (MessageId "msg-4") tracked `shouldBe` Just (AckRetry (RetryDelay 0))
      lookup (MessageId "msg-5") tracked `shouldBe` Just (AckRetry (RetryDelay 0))
      sort (map fst tracked) `shouldBe` expectedIds

  describe "finalizer failure (M2)" $ do
    it "exhausts retries, names the failed MessageId, and still attempts the rest" $ do
      -- Two-message batch: msg-1 finalizer always throws; msg-2 records normally.
      tracking <- runEff $ runTracingNoop newTrackingAck
      result <- try $ runEff $ runTracingNoop $ do
        let mid1 = MessageId "perm-1"
            mid2 = MessageId "perm-2"
            ing1 =
              mkIngested (mkEnv 1 mid1) alwaysFailAckHandle
            ing2 =
              mkIngested (mkEnv 2 mid2) (trackingAckHandle tracking mid2)
            info =
              BatchInfo
                { batchKey = defaultBatchKey,
                  size = 2,
                  trigger = TriggerSize,
                  partition = Nothing
                }
        _ <- runBatchesWithMetrics (ProcessorId "m2-fin") Serial (\_ _ -> pure ackAllOk) [(info, NE.fromList [ing1, ing2])]
        pure ()

      case result of
        Left (ProcessorFailure msg (Just messageId)) -> do
          msg `shouldSatisfy` Text.isInfixOf "perm-1"
          messageId `shouldBe` MessageId "perm-1"
        Left other -> expectationFailure ("unexpected processor failure: " <> show other)
        Right () -> expectationFailure "expected ProcessorFailure from exhausted finalization"

      -- The other message was still attempted and finalized despite msg-1 failing.
      tracked <- runEff $ runTracingNoop $ getTrackedDecisions tracking
      tracked `shouldBe` [(MessageId "perm-2", AckOk)]

  describe "keyed concurrency (M2)" $ do
    it "does not overlap batches with the same BatchKey under Async" $ do
      violation <- runEff $ runTracingNoop $ do
        activeKeys <- liftIO $ newTVarIO Set.empty
        violated <- liftIO $ newTVarIO False
        -- Handler marks its key active on entry (flagging a violation if it was
        -- already active), holds it briefly, then clears it.
        let handler info _msgs = do
              liftIO $ atomically $ do
                active <- readTVar activeKeys
                if info.batchKey `Set.member` active
                  then writeTVar violated True
                  else pure ()
                modifyTVar' activeKeys (Set.insert info.batchKey)
              liftIO $ threadDelay 20000
              liftIO $ atomically $ modifyTVar' activeKeys (Set.delete info.batchKey)
              pure ackAllOk
        -- Three batches for key "a" (must serialize) plus two for "b".
        batches <- traverse (\(i, k) -> buildKeyedBatch i (BatchKey k)) (zip [1 ..] ["a", "a", "a", "b", "b"])
        _ <- runBatchesWithMetrics (ProcessorId "m2-keyed") (Async 3) handler batches
        liftIO $ readTVarIO violated
      violation `shouldBe` False

expectedIds :: [MessageId]
expectedIds = [MessageId ("msg-" <> tshow i) | i <- [1 .. 5 :: Int]]

-- | Envelope for message @i@ with the given id.
mkEnv :: Int -> MessageId -> Envelope String
mkEnv i mid =
  (mkEnvelope mid ("payload-" <> show i)) {enqueuedAt = Just testTime}

-- | Build a batch of @n@ messages (ids @msg-1@..@msg-n@) whose acks record into
-- the given TrackingAck.
buildBatch ::
  (IOE :> es) =>
  TrackingAck ->
  Int ->
  BatchTrigger ->
  Eff es (BatchInfo, NonEmpty (Ingested es String))
buildBatch tracking n trig = do
  let mk i =
        let mid = MessageId ("msg-" <> tshow i)
         in mkIngested (mkEnv i mid) (trackingAckHandle tracking mid)
      msgs = map mk [1 .. n]
      info =
        BatchInfo
          { batchKey = defaultBatchKey,
            size = n,
            trigger = trig,
            partition = Nothing
          }
  pure (info, NE.fromList msgs)

-- | Build a batch containing the exact message numbers requested.
buildRangeBatch ::
  (IOE :> es) =>
  TrackingAck ->
  [Int] ->
  BatchTrigger ->
  Eff es (BatchInfo, NonEmpty (Ingested es String))
buildRangeBatch tracking ids trig = do
  let mk i =
        let mid = MessageId ("msg-" <> tshow i)
         in mkIngested (mkEnv i mid) (trackingAckHandle tracking mid)
      msgs = map mk ids
      info =
        BatchInfo
          { batchKey = defaultBatchKey,
            size = length ids,
            trigger = trig,
            partition = Nothing
          }
  pure (info, NE.fromList msgs)

-- | A single-message batch with the given batch key (no-op tracking ack); used
-- to exercise the keyed scheduler.
buildKeyedBatch ::
  (IOE :> es) =>
  Int ->
  BatchKey ->
  Eff es (BatchInfo, NonEmpty (Ingested es String))
buildKeyedBatch i key = do
  let mid = MessageId ("k-" <> tshow i)
      ing =
        mkIngested (mkEnv i mid) (AckHandle (\_ -> pure ()))
      info =
        BatchInfo
          { batchKey = key,
            size = 1,
            trigger = TriggerSize,
            partition = Nothing
          }
  pure (info, NE.fromList [ing])

-- | A finalizer that throws its first two attempts, then records success.
flakyAckHandle :: (IOE :> es) => IORef Int -> TrackingAck -> MessageId -> AckHandle es
flakyAckHandle counter tracking mid =
  AckHandle $ \decision -> do
    n <- liftIO $ atomicModifyIORef' counter (\c -> (c + 1, c))
    if n < 2
      then liftIO $ throwIO (userError "flaky finalize")
      else liftIO $ modifyIORef' tracking.trackedDecisions ((mid, decision) :)

-- | A finalizer that always throws (permanent adapter failure).
alwaysFailAckHandle :: (IOE :> es) => AckHandle es
alwaysFailAckHandle = AckHandle $ \_ -> liftIO $ throwIO (userError "permanent finalize failure")

testTime :: UTCTime
testTime = UTCTime (fromGregorian 2026 1 1) 0

tshow :: (Show a) => a -> Text
tshow = Text.pack . show

validDeadLetterCode :: Text -> DeadLetterCode
validDeadLetterCode code =
  case mkDeadLetterCode code of
    Left err -> error $ "invalid test fixture: " <> show err
    Right valid -> valid

-- Exercise the actual supervised batch runner in both metrics modes.
runStatusBatches :: (IOE :> es, Tracing :> es) => Metrics.MetricsDetail -> BatchHandler es msg -> [(BatchInfo, NonEmpty (Ingested es msg))] -> Eff es ProcessorMetrics
runStatusBatches detail handler batches = do
  master <- startMaster IgnoreAll
  let values = concatMap (NE.toList . snd) batches
      adapter = Adapter "status-batch" (Stream.fromList values) (pure ())
      config = defaultBatchConfig {batchSize = length values}
  proc <- runSupervisedBatchWith detail master 10 (ProcessorId "status-batch") Serial config adapter handler
  liftIO $ atomically $ readTVar proc.done >>= check
  metrics <- liftIO $ Metrics.sampleMetrics proc.metrics
  liftIO $ Metrics.slotMisses proc.metrics `shouldReturn` 0
  stopMaster master
  pure metrics