BatchProcessor.hs

-- | __Internal module.__ Exposed for the test suite and benchmarks only.
-- No PVP guarantees: anything here may change or disappear in any release.
-- Application authors should import "Shibuya" instead.
--
-- Batch execution stage: run a batch handler over an emitted batch, resolve
-- one acknowledgement decision per retained message, and finalize resiliently.
--
-- This is the reliability heart of batch processing. For each ready batch it:
--
--   1. Opens an OpenTelemetry span scoped to the whole batch.
--   2. Runs the user's 'BatchHandler' under exception isolation.
--   3. On success, uses the returned batch ack; on exception, substitutes the
--      framework default @ackAll (AckRetry (RetryDelay 0))@ (redeliver the whole
--      batch, no data loss).
--   4. Resolves EVERY message in its OWN retained 'NonEmpty' list to one
--      decision, looking each decision up by message id with a fallback.
--   5. Calls each message's idempotent finalizer with bounded retry. If retry
--      is exhausted, records the message id and fails the processor loudly.
--   6. On @AckHalt@, sets a shared halt flag (does not throw); the caller drains
--      then throws a processor halt.
--   7. Records batch metrics.
--
-- The decision loop iterates the framework's retained list, never the handler's
-- output, so handler bugs cannot skip or misassign retained messages.
module Shibuya.Internal.Runner.BatchProcessor
  ( -- * Batch execution
    processOneBatch,
    processBatchesUntilDrained,

    -- * Standalone driver (for tests / finite batch lists)
    runBatchesWithMetrics,
  )
where

import Control.Applicative ((<|>))
import Control.Monad (when)
import Data.Foldable (for_, traverse_)
import Data.HashMap.Strict qualified as HashMap
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe, isJust, listToMaybe)
import Data.Text qualified as Text
import Effectful (Eff, IOE, Limit (..), Persistence (..), UnliftStrategy (..), liftIO, withEffToIO, (:>))
import Effectful.Exception qualified as Exception
import GHC.IO (unsafeUnmask)
import OpenTelemetry.Attributes (toAttribute)
import OpenTelemetry.Trace.Core qualified as OTel
import Shibuya.Batch
  ( BatchAck (..),
    BatchHandler,
    BatchInfo (..),
    BatchKey (..),
    BatchTrigger (..),
    ackAll,
  )
import Shibuya.Core.Ack
  ( AckDecision (..),
    HaltReason (..),
    RetryDelay (..),
  )
import Shibuya.Core.Ingested (Ingested (..), toMessage)
import Shibuya.Core.Metrics
  ( AckDecisionMetric (..),
    BatchTriggerMetric (..),
    MetricsHandle (..),
    ProcessorId (..),
    ProcessorMetrics (..),
    beginProcessing,
    newMetricsHandle,
    recordBatchOutcomeMetrics,
    sampleMetrics,
  )
import Shibuya.Core.Metrics qualified as Metrics
import Shibuya.Core.Types (Envelope (..))
import Shibuya.Internal.Runner.Finalize (finalizeWithRetry)
import Shibuya.Internal.Runner.Halt
  ( ProcessorExit (..),
    ProcessorExitPublisher,
    ProcessorSignal,
    newProcessorExitPublisher,
    newProcessorSignal,
    readProcessorExit,
    requestProcessorExit,
    throwProcessorExit,
  )
import Shibuya.Internal.Runner.KeyedScheduler (runKeyedScheduler)
import Shibuya.Policy (Concurrency (..))
import Shibuya.Prelude
import Shibuya.Telemetry.Effect
  ( Tracing,
    addAttribute,
    addAttributes,
    addEvent,
    recordException,
    setStatus,
    withExtractedContext,
    withSpan',
  )
import Shibuya.Telemetry.Propagation (extractTraceContext)
import Shibuya.Telemetry.Semantic
  ( attrMessagingDestinationName,
    attrMessagingOperation,
    attrMessagingSystem,
    attrShibuyaBatchKey,
    attrShibuyaBatchSize,
    attrShibuyaBatchTrigger,
    attrShibuyaInflightCount,
    attrShibuyaInflightMax,
    consumerSpanArgs,
    eventBatchCompleted,
    eventBatchStarted,
    mkEvent,
    processSpanName,
  )
import Streamly.Data.Fold qualified as Fold
import Streamly.Data.Stream (Stream)
import Streamly.Data.Stream qualified as Stream
import UnliftIO (catchAny)

-- | Execute one emitted batch and finalize every retained message resiliently.
--
-- @maxConc@ is the batch-concurrency limit (reported on the span). @haltRef@ is
-- the shared halt flag: on 'AckHalt' this sets it via 'atomicWriteIORef' and
-- returns normally, letting the stream drain.
processOneBatch ::
  (IOE :> es, Tracing :> es) =>
  MetricsHandle ->
  ProcessorId ->
  Int ->
  ProcessorSignal ->
  ProcessorExitPublisher ->
  BatchHandler es msg ->
  (BatchInfo, NonEmpty (Ingested es msg)) ->
  Eff es ()
processOneBatch metricsHandle = case Metrics.metricsDetailLevel metricsHandle of
  Metrics.MetricsBasic -> processOneBatchWithDetail False metricsHandle
  Metrics.MetricsDetailed _ -> processOneBatchWithDetail True metricsHandle

processOneBatchWithDetail ::
  (IOE :> es, Tracing :> es) =>
  Bool ->
  MetricsHandle ->
  ProcessorId ->
  Int ->
  ProcessorSignal ->
  ProcessorExitPublisher ->
  BatchHandler es msg ->
  (BatchInfo, NonEmpty (Ingested es msg)) ->
  Eff es ()
processOneBatchWithDetail detailed metricsHandle procId maxConc stopSignal exitPublisher handler (info, batch) = do
  -- Use the first message's trace context as the batch span's parent. A batch
  -- may span several traces; picking the first is a pragmatic single parent
  -- (full fan-in links are a later refinement).
  let firstMsg = NE.head batch
      parentCtx = firstMsg.envelope.traceContext >>= extractTraceContext
      ProcessorId pidText = procId

  withExtractedContext parentCtx $
    withSpan' (processSpanName pidText) consumerSpanArgs $ \traceSpan -> do
      -- Framework messaging.* attributes plus batch-scoped attributes.
      let BatchKey keyText = info.batchKey
          frameworkAttrs =
            HashMap.fromList
              [ (attrMessagingSystem, toAttribute ("shibuya" :: Text)),
                (attrMessagingDestinationName, toAttribute pidText),
                (attrMessagingOperation, toAttribute ("process" :: Text)),
                (attrShibuyaBatchKey, toAttribute keyText),
                (attrShibuyaBatchSize, toAttribute info.size),
                (attrShibuyaBatchTrigger, toAttribute (triggerText info.trigger))
              ]
      addAttributes traceSpan frameworkAttrs

      let process currentInflight = do
            addAttribute traceSpan attrShibuyaInflightCount currentInflight
            addAttribute traceSpan attrShibuyaInflightMax maxConc

            addEvent traceSpan (mkEvent eventBatchStarted [])

            terminal <- liftIO $ readProcessorExit stopSignal

            -- Run the handler under exception isolation. On any exception, record it
            -- on the span and substitute the whole-batch retry default.
            (handlerResult, skippedAfterHalt) <-
              case terminal of
                Just _ -> pure (Left (), True)
                Nothing -> do
                  result <-
                    catchAny
                      (Right <$> handler info (toMessage <$> batch))
                      ( \ex -> do
                          recordException traceSpan ex
                          pure (Left ())
                      )
                  pure (result, False)
            let (resolvedAck, handlerThrew) = case handlerResult of
                  Right a -> (a, False)
                  Left () -> (ackAll (AckRetry (RetryDelay 0)), True)

            -- RELIABLE FINALIZATION: iterate OUR OWN retained list, never the
            -- handler's output. For each retained message, choose its decision once
            -- via findWithDefault, then call the idempotent adapter finalizer with
            -- bounded retry. Do not let one adapter failure prevent attempts for the
            -- rest of the batch.
            results <-
              mapM
                ( \ingested -> do
                    let explicitDecision = Map.lookup ingested.envelope.messageId resolvedAck.decisions
                        d = fromMaybe resolvedAck.fallback explicitDecision
                    finalResult <- finalizeWithRetry traceSpan ingested d
                    let recordCursor = liftIO $ Metrics.recordAcknowledgedCursor metricsHandle ingested.envelope.partition ingested.envelope.cursor
                    when detailed $ case (finalResult, d) of
                      (Right (), AckOk) -> recordCursor
                      (Right (), AckDeadLetter _) -> recordCursor
                      _ -> pure ()
                    pure (ingested.envelope.messageId, isJust explicitDecision, d, finalResult)
                )
                (NE.toList batch)
            let decisions = [d | (_, _, d, _) <- results]
                finalizeFailures = [(mid, ex) | (mid, _, _, Left ex) <- results]

            -- Compute halt and partial-failure signals from the resolved decisions.
            let finalizationFailure =
                  case finalizeFailures of
                    [] -> Nothing
                    failed@((firstMessageId, _) : _) ->
                      Just $
                        ProcessorFailed
                          ( "batch finalization failed for message ids: "
                              <> Text.intercalate ", " [tshow mid | (mid, _) <- failed]
                          )
                          (Just firstMessageId)
                firstHalt = listToMaybe [r | AckHalt r <- decisions]
                requestedExit = finalizationFailure <|> (ProcessorHalted <$> firstHalt)
                overrideFailures =
                  [ ()
                  | (_, explicitlyNamed, d, _) <- results,
                    explicitlyNamed,
                    isFailing d
                  ]
                partialInc = not handlerThrew && not (null overrideFailures)

            -- Span status: error on halt or exception, otherwise Ok.
            addEvent traceSpan (mkEvent eventBatchCompleted [])
            case requestedExit of
              Just processorExit -> setStatus traceSpan (OTel.Error (processorExitText processorExit))
              Nothing ->
                if skippedAfterHalt
                  then setStatus traceSpan (OTel.Error "batch skipped after halt")
                  else
                    if handlerThrew
                      then setStatus traceSpan (OTel.Error "batch handler exception")
                      else setStatus traceSpan OTel.Ok

            traverse_ (recordException traceSpan . snd) finalizeFailures

            -- Record metrics: decrement in-flight, add per-message counter deltas,
            -- advance batch counters, set Failed state on halt or exhausted
            -- finalization retry.
            liftIO $
              recordBatchOutcomeMetrics
                metricsHandle
                (triggerMetric info.trigger)
                info.size
                handlerThrew
                partialInc
                (decisionMetric <$> decisions)
                (processorExitText <$> requestedExit)

            -- Halt: set the shared flag; do NOT throw (let the stream drain).
            for_ requestedExit $ \processorExit ->
              liftIO $ requestProcessorExit exitPublisher processorExit
      if not detailed
        then do
          currentInflight <- liftIO $ beginProcessing metricsHandle maxConc
          process currentInflight
        else Exception.mask $ \restore -> do
          currentInflight <- liftIO $ beginProcessing metricsHandle maxConc
          slot <- liftIO $ Metrics.claimInFlightSlot metricsHandle firstMsg.envelope.messageId
          restore (process currentInflight)
            `Exception.onException` liftIO (Metrics.abandonInFlightSlot metricsHandle slot)
          liftIO (Metrics.releaseInFlightSlot metricsHandle slot)
            `Exception.onException` liftIO (Metrics.abandonInFlightSlot metricsHandle slot)
  where
    isFailing :: AckDecision -> Bool
    isFailing (AckDeadLetter _) = True
    isFailing (AckRetry _) = True
    isFailing _ = False
{-# INLINE processOneBatchWithDetail #-}

decisionMetric :: AckDecision -> AckDecisionMetric
decisionMetric AckOk = CountProcessed
decisionMetric (AckRetry _) = CountProcessed
decisionMetric (AckDeadLetter _) = CountFailed
decisionMetric (AckHalt reason) = CountHalt (haltReasonText reason)

triggerMetric :: BatchTrigger -> BatchTriggerMetric
triggerMetric TriggerSize = CountTriggerSize
triggerMetric TriggerTimeout = CountTriggerTimeout
triggerMetric TriggerFlush = CountTriggerFlush

triggerText :: BatchTrigger -> Text
triggerText TriggerSize = "size"
triggerText TriggerTimeout = "timeout"
triggerText TriggerFlush = "flush"

haltReasonText :: HaltReason -> Text
haltReasonText (HaltOrderedStream t) = t
haltReasonText (HaltFatal t) = t

processorExitText :: ProcessorExit -> Text
processorExitText (ProcessorHalted reason) = haltReasonText reason
processorExitText (ProcessorFailed failure _) = failure

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

-- | Fold the ready-batch stream, running each batch under the batch-concurrency
-- policy. Batches with the same batch key are always serialized in emission
-- order; different keys may run concurrently up to the configured bound. This
-- does NOT throw on halt; after draining, the caller inspects @haltRef@ and
-- throws a processor halt (see 'runBatchesWithMetrics').
processBatchesUntilDrained ::
  (IOE :> es, Tracing :> es) =>
  MetricsHandle ->
  ProcessorId ->
  Concurrency ->
  BatchHandler es msg ->
  Stream IO (BatchInfo, NonEmpty (Ingested es msg)) ->
  ProcessorSignal ->
  ProcessorExitPublisher ->
  Eff es ()
processBatchesUntilDrained metricsHandle procId concurrency handler batchesStream stopSignal exitPublisher =
  -- Specialize the complete loop once; keep mode selection and ownership
  -- closures out of the basic batch hot path.
  case Metrics.metricsDetailLevel metricsHandle of
    Metrics.MetricsBasic -> processBatchesUntilDrainedWithDetail False metricsHandle procId concurrency handler batchesStream stopSignal exitPublisher
    Metrics.MetricsDetailed _ -> processBatchesUntilDrainedWithDetail True metricsHandle procId concurrency handler batchesStream stopSignal exitPublisher

processBatchesUntilDrainedWithDetail ::
  (IOE :> es, Tracing :> es) =>
  Bool ->
  MetricsHandle ->
  ProcessorId ->
  Concurrency ->
  BatchHandler es msg ->
  Stream IO (BatchInfo, NonEmpty (Ingested es msg)) ->
  ProcessorSignal ->
  ProcessorExitPublisher ->
  Eff es ()
processBatchesUntilDrainedWithDetail detailed metricsHandle procId concurrency handler batchesStream stopSignal exitPublisher = do
  let maxConc = case concurrency of
        Serial -> 1
        Ahead n -> n
        Async n -> n

  withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do
    -- The supervisor keeps framework coordination masked; only the owned batch
    -- action is unmasked so user code and finalizers remain cancellable.
    let runBatchAction batch =
          runInIO $ processOneBatchWithDetail detailed metricsHandle procId maxConc stopSignal exitPublisher handler batch
        batchAction = unsafeUnmask . runBatchAction
        pendingLimit = max 2 (2 * max 1 maxConc)
    case concurrency of
      Serial ->
        unsafeUnmask $
          Stream.fold Fold.drain $
            Stream.mapM runBatchAction batchesStream
      Ahead n ->
        runKeyedScheduler (max 1 n) pendingLimit (Just . fstBatchKey) batchAction batchesStream
      Async n ->
        runKeyedScheduler (max 1 n) pendingLimit (Just . fstBatchKey) batchAction batchesStream
{-# INLINE processBatchesUntilDrainedWithDetail #-}

fstBatchKey :: (BatchInfo, NonEmpty (Ingested es msg)) -> BatchKey
fstBatchKey (info, _) = info.batchKey

-- | Self-contained driver for finite batch lists (tests / simple setups).
-- Mirrors 'Shibuya.Internal.Runner.Supervised.runWithMetrics': creates its own metrics
-- TVar and halt flag, runs execution to completion, and — after draining —
-- throws a processor halt if a batch requested a halt. Returns the final metrics.
runBatchesWithMetrics ::
  (IOE :> es, Tracing :> es) =>
  ProcessorId ->
  Concurrency ->
  BatchHandler es msg ->
  [(BatchInfo, NonEmpty (Ingested es msg))] ->
  Eff es ProcessorMetrics
runBatchesWithMetrics procId concurrency handler batches = do
  now <- liftIO getCurrentTime
  metricsHandle <- liftIO $ newMetricsHandle now
  stopSignal <- liftIO newProcessorSignal
  let exitPublisher = newProcessorExitPublisher stopSignal

  let batchesStream = Stream.fromList batches
  processBatchesUntilDrained metricsHandle procId concurrency handler batchesStream stopSignal exitPublisher

  maybeExit <- liftIO $ readProcessorExit stopSignal
  case maybeExit of
    Just processorExit -> liftIO $ throwProcessorExit processorExit
    Nothing -> liftIO $ sampleMetrics metricsHandle