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 (DeliveryStatus (..), 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.Correlation (BatchLinks (..), collectBatchLinksBy, consumerSpanArgsLinkedTo)
import Shibuya.Telemetry.Effect
  ( Tracing,
    addAttribute,
    addAttributes,
    addEvent,
    isTracingEnabled,
    recordException,
    setStatus,
    withRootContext,
    withSpan',
  )
import Shibuya.Telemetry.Semantic
  ( attrMessagingBatchMessageCount,
    attrMessagingDestinationName,
    attrMessagingOperation,
    attrMessagingOperationName,
    attrMessagingSystem,
    attrShibuyaBatchKey,
    attrShibuyaBatchLinksDropped,
    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 original@(info, batch) =
  scan 0 (NE.toList batch)
  where
    -- All-current batches retain the original pair without constructing a list.
    scan _ [] = run original
    scan count (member : rest) = do
      status <- member.deliveryStatus
      case status of
        DeliveryCurrent -> scan (count + 1) rest
        DeliverySuperseded -> do
          suffix <- retain rest
          case NE.nonEmpty (take count (NE.toList batch) <> suffix) of
            Nothing -> pure ()
            Just current -> run (info {size = NE.length current}, current)
    retain [] = pure []
    retain (member : rest) = do
      status <- member.deliveryStatus
      suffix <- retain rest
      pure $ case status of
        DeliveryCurrent -> member : suffix
        DeliverySuperseded -> suffix
    run = processCurrentBatchWithDetail detailed metricsHandle procId maxConc stopSignal exitPublisher handler
{-# INLINE processOneBatchWithDetail #-}

processCurrentBatchWithDetail ::
  (IOE :> es, Tracing :> es) =>
  Bool ->
  MetricsHandle ->
  ProcessorId ->
  Int ->
  ProcessorSignal ->
  ProcessorExitPublisher ->
  BatchHandler es msg ->
  (BatchInfo, NonEmpty (Ingested es msg)) ->
  Eff es ()
processCurrentBatchWithDetail detailed metricsHandle procId maxConc stopSignal exitPublisher handler (info, batch) = do
  enabled <- isTracingEnabled
  let ProcessorId pidText = procId
      spanName = processSpanName pidText
      -- A batch has many producers and a span has one parent, so the batch
      -- span is always a root and links each member's creation context
      -- instead (ADR 0014). Disabled tracing reads no member headers.
      openBatchSpan body
        | not enabled = withSpan' spanName consumerSpanArgs body
        | otherwise =
            let !batchLinks = collectBatchLinksBy (.envelope.messageId) (.envelope.traceContext) batch
                !args = case batchLinks.links of
                  [] -> consumerSpanArgs
                  links -> consumerSpanArgsLinkedTo links
             in withRootContext $
                  withSpan' spanName args $ \traceSpan -> do
                    when (batchLinks.dropped > 0) $
                      addAttribute traceSpan attrShibuyaBatchLinksDropped batchLinks.dropped
                    body traceSpan

  openBatchSpan $ \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)),
              (attrMessagingOperationName, toAttribute ("process" :: Text)),
              (attrShibuyaBatchKey, toAttribute keyText),
              (attrShibuyaBatchSize, toAttribute info.size),
              (attrMessagingBatchMessageCount, 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 (NE.head batch).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 processCurrentBatchWithDetail #-}

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