Metrics.hs

-- | Metrics and state tracking for processors.
-- Provides introspection into what's happening in the system.
module Shibuya.Core.Metrics
  ( -- * Processor State
    ProcessorState (..),
    ProcessorId (..),

    -- * In-Flight Tracking
    InFlightInfo (..),
    emptyInFlightInfo,

    -- * Stream Statistics
    StreamStats (..),
    emptyStreamStats,

    -- * Batch Statistics
    BatchStats (..),
    emptyBatchStats,
    incBatchesEmitted,
    addBatchedMessages,
    incPartialFailures,
    incSizeTriggered,
    incTimeoutTriggered,
    incFlushTriggered,

    -- * Combined Metrics
    ProcessorMetrics (..),
    emptyProcessorMetrics,
    HotCounters (..),
    MetricsHandle (..),
    newMetricsHandle,
    sampleMetrics,
    incrementReceived,
    beginProcessing,
    AckDecisionMetric (..),
    finishProcessing,
    finishFinalizationFailure,
    BatchTriggerMetric (..),
    recordBatchOutcomeMetrics,

    -- * Metrics Map
    MetricsMap,

    -- * Metrics Updates
    incReceived,
    incProcessed,
    incFailed,
  )
where

import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVarIO)
import Control.Monad (unless, void, when)
import Data.Aeson (FromJSON (..), FromJSONKey (..), ToJSON (..), ToJSONKey (..), object, withObject, (.:))
import Data.Aeson qualified as Aeson
import Data.Atomics.Counter (AtomicCounter, incrCounter, readCounter)
import Data.Atomics.Counter qualified as Counter
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import Data.Text qualified as Text
import Shibuya.Prelude

-- | Processor identifier.
newtype ProcessorId = ProcessorId {unProcessorId :: Text}
  deriving stock (Eq, Ord, Show, Generic)
  deriving newtype (ToJSON, FromJSON, ToJSONKey, FromJSONKey)

-- | Tracks concurrent in-flight messages.
data InFlightInfo = InFlightInfo
  { -- | Currently processing count
    inFlight :: !Int,
    -- | Configured max concurrency (1 for Serial)
    maxConcurrency :: !Int
  }
  deriving stock (Eq, Show, Generic)

instance ToJSON InFlightInfo where
  toJSON info =
    object
      [ "inFlight" Aeson..= info.inFlight,
        "maxConcurrency" Aeson..= info.maxConcurrency
      ]

instance FromJSON InFlightInfo where
  parseJSON = withObject "InFlightInfo" $ \v ->
    InFlightInfo
      <$> v .: "inFlight"
      <*> v .: "maxConcurrency"

-- | Create empty in-flight info with given max concurrency.
emptyInFlightInfo :: Int -> InFlightInfo
emptyInFlightInfo = InFlightInfo 0

-- | Processor runtime state.
data ProcessorState
  = -- | Waiting for messages
    Idle
  | -- | Currently processing (in-flight info, last activity time)
    Processing !InFlightInfo !UTCTime
  | -- | Failed with error (error message, timestamp)
    Failed !Text !UTCTime
  | -- | Processor has been stopped
    Stopped
  deriving stock (Eq, Show, Generic)

instance ToJSON ProcessorState where
  toJSON Idle = object ["status" Aeson..= ("idle" :: Text)]
  toJSON (Processing info lastActivity) =
    object
      [ "status" Aeson..= ("processing" :: Text),
        "inFlight" Aeson..= info.inFlight,
        "maxConcurrency" Aeson..= info.maxConcurrency,
        "lastActivity" Aeson..= lastActivity
      ]
  toJSON (Failed err timestamp) =
    object
      [ "status" Aeson..= ("failed" :: Text),
        "error" Aeson..= err,
        "timestamp" Aeson..= timestamp
      ]
  toJSON Stopped = object ["status" Aeson..= ("stopped" :: Text)]

instance FromJSON ProcessorState where
  parseJSON = withObject "ProcessorState" $ \v -> do
    status <- v .: "status"
    case status :: Text of
      "idle" -> pure Idle
      "processing" -> do
        inFlightCount <- v .: "inFlight"
        maxConc <- v .: "maxConcurrency"
        lastActivity <- v .: "lastActivity"
        pure $ Processing (InFlightInfo inFlightCount maxConc) lastActivity
      "failed" -> Failed <$> v .: "error" <*> v .: "timestamp"
      "stopped" -> pure Stopped
      other -> fail $ "Unknown processor state: " <> Text.unpack other

-- | Stream statistics.
data StreamStats = StreamStats
  { -- | Messages received from stream
    received :: !Int,
    -- | Messages successfully processed
    processed :: !Int,
    -- | Messages that failed processing
    failed :: !Int
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON, FromJSON)

-- | Empty stream stats.
emptyStreamStats :: StreamStats
emptyStreamStats = StreamStats 0 0 0

-- | Batch-processing statistics, tracked alongside per-message stream stats.
data BatchStats = BatchStats
  { -- | Number of batches emitted and executed.
    batchesEmitted :: !Int,
    -- | Total messages across all emitted batches.
    batchedMessages :: !Int,
    -- | Batches with a genuine partial failure: the handler returned normally
    -- and named at least one message in its decision map with a failing
    -- decision (dead-letter or retry) while acking the rest. Counted per batch,
    -- not per message, so it does not double-count the per-message 'failed'
    -- counter.
    partialFailures :: !Int,
    -- | Batches emitted because they reached the configured size.
    sizeTriggered :: !Int,
    -- | Batches emitted because their timeout elapsed.
    timeoutTriggered :: !Int,
    -- | Batches emitted because the processor was draining (flush).
    flushTriggered :: !Int
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON, FromJSON)

-- | Empty batch stats (all zero).
emptyBatchStats :: BatchStats
emptyBatchStats = BatchStats 0 0 0 0 0 0

-- | Combined processor metrics.
data ProcessorMetrics = ProcessorMetrics
  { -- | Current state
    state :: !ProcessorState,
    -- | Per-message statistics
    stats :: !StreamStats,
    -- | Batch statistics
    batch :: !BatchStats,
    -- | When the processor started
    startedAt :: !UTCTime
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON, FromJSON)

-- | Empty processor metrics.
emptyProcessorMetrics :: UTCTime -> ProcessorMetrics
emptyProcessorMetrics now =
  ProcessorMetrics
    { state = Idle,
      stats = emptyStreamStats,
      batch = emptyBatchStats,
      startedAt = now
    }

-- | Map of processor IDs to their metrics.
type MetricsMap = Map ProcessorId ProcessorMetrics

-- | Hot per-message counters for one processor.
--
-- These counters are updated by fetch-and-add operations on the message hot
-- path. The colder 'ProcessorMetrics' TVar inside 'MetricsHandle' stores state,
-- batch counters, and metadata that change less frequently.
data HotCounters = HotCounters
  { received :: !AtomicCounter,
    processed :: !AtomicCounter,
    failed :: !AtomicCounter,
    inFlight :: !AtomicCounter
  }

-- | Write-side metrics handle for one processor.
data MetricsHandle = MetricsHandle
  { hot :: !HotCounters,
    maxConcurrencyRef :: !(IORef Int),
    burstStartedRef :: !(IORef UTCTime),
    stateActiveRef :: !(IORef Bool),
    cold :: !(TVar ProcessorMetrics)
  }

newMetricsHandle :: UTCTime -> IO MetricsHandle
newMetricsHandle now = do
  received <- Counter.newCounter 0
  processed <- Counter.newCounter 0
  failed <- Counter.newCounter 0
  inFlight <- Counter.newCounter 0
  maxConcurrencyRef <- newIORef 1
  burstStartedRef <- newIORef now
  stateActiveRef <- newIORef False
  cold <- newTVarIO (emptyProcessorMetrics now)
  pure
    MetricsHandle
      { hot =
          HotCounters
            { received = received,
              processed = processed,
              failed = failed,
              inFlight = inFlight
            },
        maxConcurrencyRef = maxConcurrencyRef,
        burstStartedRef = burstStartedRef,
        stateActiveRef = stateActiveRef,
        cold = cold
      }

sampleMetrics :: MetricsHandle -> IO ProcessorMetrics
sampleMetrics handle = do
  coldSnapshot <- readTVarIO handle.cold
  received <- readCounter handle.hot.received
  processed <- readCounter handle.hot.processed
  failed <- readCounter handle.hot.failed
  inFlight <- readCounter handle.hot.inFlight
  maxConcurrency <- readIORef handle.maxConcurrencyRef
  burstStartedAt <- readIORef handle.burstStartedRef
  let sampledStats =
        StreamStats
          { received = received,
            processed = processed,
            failed = failed
          }
      sampledState = case coldSnapshot.state of
        Failed err timestamp -> Failed err timestamp
        Stopped -> Stopped
        _ | inFlight > 0 -> Processing (InFlightInfo inFlight maxConcurrency) burstStartedAt
        _ -> Idle
  pure coldSnapshot {state = sampledState, stats = sampledStats}

incrementReceived :: MetricsHandle -> IO ()
incrementReceived handle =
  void $ incrCounter 1 handle.hot.received

beginProcessing :: MetricsHandle -> Int -> IO Int
beginProcessing handle maxConcurrency = do
  currentInflight <- incrCounter 1 handle.hot.inFlight
  when (currentInflight == 1) $ do
    stateActive <- readIORef handle.stateActiveRef
    unless stateActive $ do
      now <- getCurrentTime
      writeIORef handle.maxConcurrencyRef maxConcurrency
      writeIORef handle.burstStartedRef now
      writeIORef handle.stateActiveRef True
      atomically $
        modifyTVar' handle.cold $ \m ->
          m {state = Processing (InFlightInfo currentInflight maxConcurrency) now}
  pure currentInflight

finishProcessing :: MetricsHandle -> Either Text AckDecisionMetric -> IO ()
finishProcessing handle result = do
  case result of
    Right CountProcessed -> void $ incrCounter 1 handle.hot.processed
    Right CountFailed -> void $ incrCounter 1 handle.hot.failed
    Right CountNeither -> pure ()
    Right (CountHalt _) -> pure ()
    Left _ -> void $ incrCounter 1 handle.hot.failed
  void $ incrCounter (-1) handle.hot.inFlight
  case result of
    Left failureText -> setFailed failureText
    Right (CountHalt reasonText) -> setFailed reasonText
    _ -> pure ()
  where
    setFailed failureText = do
      now <- getCurrentTime
      writeIORef handle.stateActiveRef False
      atomically $
        modifyTVar' handle.cold $ \m ->
          m {state = Failed failureText now}

finishFinalizationFailure :: MetricsHandle -> Text -> IO ()
finishFinalizationFailure handle failureText =
  finishProcessing handle (Left failureText)

data AckDecisionMetric
  = CountProcessed
  | CountFailed
  | CountNeither
  | CountHalt !Text

recordBatchOutcomeMetrics ::
  MetricsHandle ->
  BatchTriggerMetric ->
  Int ->
  Bool ->
  Bool ->
  [AckDecisionMetric] ->
  Maybe Text ->
  IO ()
recordBatchOutcomeMetrics handle trigger size handlerThrew partialInc decisions firstHalt = do
  let (processedDelta, failedDelta) =
        foldl'
          ( \(processedAcc, failedAcc) decision ->
              if handlerThrew
                then (processedAcc, failedAcc + 1)
                else case decision of
                  CountProcessed -> (processedAcc + 1, failedAcc)
                  CountFailed -> (processedAcc, failedAcc + 1)
                  CountNeither -> (processedAcc, failedAcc)
                  CountHalt _ -> (processedAcc, failedAcc)
          )
          (0, 0)
          decisions
  when (processedDelta /= 0) $
    void $
      incrCounter processedDelta handle.hot.processed
  when (failedDelta /= 0) $
    void $
      incrCounter failedDelta handle.hot.failed
  void $ incrCounter (-1) handle.hot.inFlight
  now <- getCurrentTime
  atomically $
    modifyTVar' handle.cold $ \m ->
      let newState = case firstHalt of
            Just reasonText -> Failed reasonText now
            Nothing -> case m.state of
              Failed {} -> m.state
              Stopped -> Stopped
              _ -> m.state
          newBatch =
            incTriggerMetric trigger
              . (if partialInc then incPartialFailures else id)
              . addBatchedMessages size
              . incBatchesEmitted
              $ m.batch
       in m {state = newState, batch = newBatch}
  case firstHalt of
    Just _ -> writeIORef handle.stateActiveRef False
    Nothing -> pure ()

data BatchTriggerMetric
  = CountTriggerSize
  | CountTriggerTimeout
  | CountTriggerFlush

incTriggerMetric :: BatchTriggerMetric -> BatchStats -> BatchStats
incTriggerMetric CountTriggerSize = incSizeTriggered
incTriggerMetric CountTriggerTimeout = incTimeoutTriggered
incTriggerMetric CountTriggerFlush = incFlushTriggered

-- | Increment received count.
incReceived :: StreamStats -> StreamStats
incReceived StreamStats {received, processed, failed} =
  StreamStats {received = received + 1, processed = processed, failed = failed}

-- | Increment processed count.
incProcessed :: StreamStats -> StreamStats
incProcessed StreamStats {received, processed, failed} =
  StreamStats {received = received, processed = processed + 1, failed = failed}

-- | Increment failed count.
incFailed :: StreamStats -> StreamStats
incFailed StreamStats {received, processed, failed} =
  StreamStats {received = received, processed = processed, failed = failed + 1}

-- | Increment the emitted-batch counter.
incBatchesEmitted :: BatchStats -> BatchStats
incBatchesEmitted s = s {batchesEmitted = s.batchesEmitted + 1}

-- | Add to the total batched-messages counter.
addBatchedMessages :: Int -> BatchStats -> BatchStats
addBatchedMessages n s = s {batchedMessages = s.batchedMessages + n}

-- | Increment the partial-failure batch counter.
incPartialFailures :: BatchStats -> BatchStats
incPartialFailures s = s {partialFailures = s.partialFailures + 1}

-- | Increment the size-trigger counter.
incSizeTriggered :: BatchStats -> BatchStats
incSizeTriggered s = s {sizeTriggered = s.sizeTriggered + 1}

-- | Increment the timeout-trigger counter.
incTimeoutTriggered :: BatchStats -> BatchStats
incTimeoutTriggered s = s {timeoutTriggered = s.timeoutTriggered + 1}

-- | Increment the flush-trigger counter.
incFlushTriggered :: BatchStats -> BatchStats
incFlushTriggered s = s {flushTriggered = s.flushTriggered + 1}