App.hs

-- | Application entry point for running Shibuya queue processors.
module Shibuya.App
  ( -- * Running Processors
    runApp,
    AppConfig (..),
    defaultAppConfig,
    QueueProcessor (..),
    mkProcessor,
    mkBatchProcessor,
    AppHandle,

    -- * AppHandle Operations
    getAppMetrics,
    getAppMaster,
    stopApp,
    stopAppGracefully,
    waitApp,
    pauseProcessor,
    resumeProcessor,
    isProcessorPaused,
    ControlOutcome (..),
    PauseState (..),

    -- * Shutdown Configuration
    ShutdownConfig (..),
    defaultShutdownConfig,

    -- * Supervision Strategy
    SupervisionStrategy (..),

    -- * Errors
    AppError (..),
    Master,
    getAllMetrics,
    getAllMetricsIO,
    getProcessorMetrics,
    getProcessorMetricsIO,

    -- * Batch API (re-exported from "Shibuya.Batch")
    module Shibuya.Batch,

    -- * Re-exports
    ProcessorId (..),
    ProcessorMetrics (..),
  )
where

import Control.Concurrent.Async qualified as Async
import Control.Concurrent.NQE.Supervisor qualified as NQE
import Control.Concurrent.STM
  ( STM,
    atomically,
    check,
    newEmptyTMVarIO,
    newTVarIO,
    orElse,
    putTMVar,
    readTMVar,
    readTVar,
    registerDelay,
    retry,
    throwSTM,
    writeTVar,
  )
import Control.Exception qualified as IOException
import Control.Monad (forM, forM_, void)
import Data.Bifunctor (first)
import Data.Foldable (traverse_)
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Text qualified as Text
import Data.Time.Clock (NominalDiffTime)
import Effectful (Eff, IOE, Limit (..), Persistence (..), UnliftStrategy (..), liftIO, withEffToIO, (:>))
import Effectful.Exception qualified as Exception
import GHC.Generics (Generic)
import Numeric.Natural (Natural)
import OpenTelemetry.Context qualified as TraceContext
import OpenTelemetry.Context.ThreadLocal qualified as TraceContext
import OpenTelemetry.Trace.Core qualified as OTel
import Shibuya.Adapter (Adapter (..))
import Shibuya.Batch
import Shibuya.Core.Error (ConfigError (..), HandlerError (..), PolicyError (..), RuntimeError (..))
import Shibuya.Core.Metrics
  ( MetricsMap,
    ProcessorId (..),
    ProcessorMetrics (..),
  )
import Shibuya.Core.Metrics qualified as Metrics
import Shibuya.Internal.App
  ( AppHandle (..),
    OwnershipFailure (..),
    QueueProcessor (..),
    acquireOwned,
    mkBatchProcessor,
    mkProcessor,
  )
import Shibuya.Internal.Runner.Master
  ( ControlOutcome (..),
    Master,
    PauseState (..),
    drainAndResume,
    getAllMetrics,
    getAllMetricsIO,
    getProcessorMetrics,
    getProcessorMetricsIO,
    markMasterDraining,
    markMasterRunning,
    startMaster,
    stopMaster,
  )
import Shibuya.Internal.Runner.Master qualified as Master
import Shibuya.Internal.Runner.Supervised
  ( SupervisedProcessor (..),
    runSupervisedBatchWith,
    runSupervisedWith,
  )
import Shibuya.Policy (Concurrency (..), OrderingPolicy (..), validatePolicy)
import Shibuya.Telemetry.Effect (Tracing, getTracer, isTracingEnabled)
import Shibuya.Telemetry.Semantic (attrShibuyaProcessorId, internalSpanArgs)
import UnliftIO (SomeException, displayException)
import UnliftIO qualified as UIO

--------------------------------------------------------------------------------
-- Supervision Strategy
--------------------------------------------------------------------------------

-- | Supervision strategy for processor failures.
--
-- This is Shibuya's own type that maps to NQE's supervision strategies,
-- decoupling users from the NQE library.
data SupervisionStrategy
  = -- | Ignore all child exits, keep running.
    -- Failed processors are marked as Failed in metrics but don't affect others.
    IgnoreFailures
  | -- | Stop all processors if any fails.
    -- A single processor failure triggers shutdown of all processors.
    -- Graceful exits, including finite streams completing and handlers returning
    -- @AckHalt@, do not stop sibling processors.
    StopAllOnFailure
  deriving stock (Eq, Show, Generic)

-- | Convert Shibuya's strategy type to NQE's internal type.
toNQEStrategy :: SupervisionStrategy -> NQE.Strategy
toNQEStrategy = \case
  IgnoreFailures -> NQE.IgnoreAll
  StopAllOnFailure -> NQE.IgnoreGraceful

--------------------------------------------------------------------------------
-- Shutdown Configuration
--------------------------------------------------------------------------------

-- | Configuration for graceful shutdown behavior.
data ShutdownConfig = ShutdownConfig
  { -- | Maximum time to wait for in-flight messages to drain.
    -- After this timeout, remaining processors are forcefully stopped.
    -- Default: 30 seconds.
    drainTimeout :: !NominalDiffTime,
    -- | Maximum time for adapter shutdown plus graceful draining before
    -- forced supervisor stop begins.
    -- Default: 60 seconds.
    totalShutdownTimeout :: !NominalDiffTime
  }
  deriving stock (Eq, Show, Generic)

-- | Default shutdown configuration with a 30 second drain timeout and a
-- 60 second bound on the graceful shutdown phase.
defaultShutdownConfig :: ShutdownConfig
defaultShutdownConfig = ShutdownConfig {drainTimeout = 30, totalShutdownTimeout = 60}

--------------------------------------------------------------------------------
-- Errors
--------------------------------------------------------------------------------

-- | Application errors.
-- Uses structured error types from Shibuya.Core.Error.
data AppError
  = -- | Invalid policy configuration
    AppPolicyError !PolicyError
  | -- | Handler execution error
    AppHandlerError !HandlerError
  | -- | Runtime error
    AppRuntimeError !RuntimeError
  | -- | Invalid batch configuration for a 'BatchingProcessor'
    AppBatchConfigError !BatchConfigError
  | -- | Invalid application configuration
    AppConfigInvalid !ConfigError
  deriving stock (Eq, Show)

-- | Configuration for @runApp@.
data AppConfig = AppConfig
  { -- | How processor failures affect siblings.
    strategy :: !SupervisionStrategy,
    -- | Bounded-inbox capacity per processor (backpressure). Must be >= 1.
    inboxSize :: !Int,
    -- | Optional latency and in-flight detail; defaults to basic accounting.
    metricsDetail :: !Metrics.MetricsDetail
  }
  deriving stock (Eq, Show, Generic)

-- | 'IgnoreFailures' with an inbox of 100.
defaultAppConfig :: AppConfig
defaultAppConfig = AppConfig {strategy = IgnoreFailures, inboxSize = 100, metricsDetail = Metrics.defaultMetricsDetail}

-- | Run queue processors concurrently under NQE supervision.
--
-- Each processor runs independently. Returns immediately with a handle
-- for introspection and control.
--
-- Example:
--
-- @
-- result <- runApp defaultAppConfig
--   [ ("orders", QueueProcessor ordersAdapter ordersHandler)
--   , ("events", QueueProcessor eventsAdapter eventsHandler)
--   ]
-- @
runApp ::
  (IOE :> es, Tracing :> es) =>
  -- | Application configuration
  AppConfig ->
  -- | Named processors
  [(ProcessorId, QueueProcessor es)] ->
  Eff es (Either AppError (AppHandle es))
runApp config namedProcessors =
  -- Validate all policies (and batch configs) first
  case validateAppConfig config *> validateUniqueProcessorIds namedProcessors *> validateAllPolicies namedProcessors of
    Left err -> pure $ Left err
    Right () -> do
      let nqeStrategy = toNQEStrategy config.strategy
      startupResult <-
        acquireOwned
          (startMaster nqeStrategy)
          stopMaster
          ( \master -> do
              processors <- spawnProcessors config.metricsDetail master (fromIntegral config.inboxSize) namedProcessors
              markMasterRunning master
              enabled <- isTracingEnabled
              shutdownTracer <- if enabled then Just <$> getTracer else pure Nothing
              shutdownStarted <- liftIO $ newTVarIO False
              shutdownResult <- liftIO newEmptyTMVarIO
              pure
                AppHandle
                  { master = master,
                    processors = Map.fromList processors,
                    shutdownStarted,
                    shutdownResult,
                    shutdownTracer
                  }
          )
      case startupResult of
        Right appHandle -> pure $ Right appHandle
        Left (OwnerAcquisitionFailed failure) -> startupFailure failure
        Left (OwnedActionFailed failure cleanupFailure) ->
          startupFailureWithCleanup failure cleanupFailure
  where
    startupFailure failure
      | Exception.isAsyncException failure = Exception.throwIO failure
      | otherwise =
          pure $ Left $ AppRuntimeError $ SupervisorFailed $ Text.pack $ displayException failure

    startupFailureWithCleanup failure cleanupFailure
      | Exception.isAsyncException failure = Exception.throwIO failure
      | otherwise =
          let cleanupSuffix = case cleanupFailure of
                Nothing -> ""
                Just cleanupException ->
                  "; supervisor cleanup also failed: " <> Text.pack (displayException cleanupException)
           in pure $
                Left $
                  AppRuntimeError $
                    SupervisorFailed $
                      Text.pack (displayException failure) <> cleanupSuffix

-- | Validate app configuration before starting any processor.
validateAppConfig :: AppConfig -> Either AppError ()
validateAppConfig config
  | config.inboxSize < 1 = Left $ AppConfigInvalid $ InvalidInboxSize config.inboxSize
  | Metrics.MetricsDetailed options <- config.metricsDetail,
    options.maxTrackedPartitions < 0 =
      Left $ AppConfigInvalid $ InvalidMaxTrackedPartitions options.maxTrackedPartitions
  | otherwise = Right ()

-- | Reject duplicate processor identifiers before the master or any adapter is
-- acquired. Keeping this separate from handle construction prevents a live
-- processor from being silently discarded by 'Map.fromList'.
validateUniqueProcessorIds :: [(ProcessorId, QueueProcessor es)] -> Either AppError ()
validateUniqueProcessorIds = go Set.empty
  where
    go _ [] = Right ()
    go seen ((pid, _) : rest)
      | pid `Set.member` seen = Left $ AppConfigInvalid $ DuplicateProcessorId pid
      | otherwise = go (Set.insert pid seen) rest

-- | Validate all processor policies (and batch configs) before starting.
validateAllPolicies :: [(ProcessorId, QueueProcessor es)] -> Either AppError ()
validateAllPolicies = traverse_ validateOne
  where
    validateOne (_, qp) = case qp of
      QueueProcessor {ordering, concurrency} ->
        first AppPolicyError (validatePolicy ordering concurrency)
      BatchingProcessor {ordering, concurrency, batchConfig} -> do
        first AppPolicyError (validatePolicy ordering concurrency)
        validateBatchOrdering ordering concurrency
        first AppBatchConfigError (validateBatchConfig batchConfig)

    validateBatchOrdering PartitionedInOrder (Ahead _) =
      Left $
        AppPolicyError $
          InvalidPolicyCombo
            "PartitionedInOrder with Ahead/Async is supported only for QueueProcessor: batching processors schedule by BatchKey, not by Envelope.partition"
    validateBatchOrdering PartitionedInOrder (Async _) =
      Left $
        AppPolicyError $
          InvalidPolicyCombo
            "PartitionedInOrder with Ahead/Async is supported only for QueueProcessor: batching processors schedule by BatchKey, not by Envelope.partition"
    validateBatchOrdering _ _ = Right ()

-- | Spawn all processors under supervision.
spawnProcessors ::
  (IOE :> es, Tracing :> es) =>
  Metrics.MetricsDetail ->
  Master ->
  Natural ->
  [(ProcessorId, QueueProcessor es)] ->
  Eff es [(ProcessorId, (SupervisedProcessor, QueueProcessor es))]
spawnProcessors detail master inboxSize = traverse spawnOne
  where
    spawnOne (procId, qp) = case qp of
      QueueProcessor {adapter, handler, ordering, concurrency} -> do
        sp <- runSupervisedWith detail master inboxSize procId ordering concurrency adapter handler
        pure (procId, (sp, qp))
      BatchingProcessor {adapter, batchHandler, batchConfig, concurrency} -> do
        sp <-
          runSupervisedBatchWith
            detail
            master
            inboxSize
            procId
            concurrency
            batchConfig
            adapter
            batchHandler
        pure (procId, (sp, qp))

--------------------------------------------------------------------------------
-- AppHandle Operations
--------------------------------------------------------------------------------

-- | Get metrics for all processors.
getAppMetrics :: (IOE :> es) => AppHandle es -> Eff es MetricsMap
getAppMetrics appHandle = getAllMetrics appHandle.master

-- | Get the master handle for direct access.
-- This is useful for integrating with the metrics server.
getAppMaster :: AppHandle es -> Master
getAppMaster appHandle = appHandle.master

-- | Gracefully stop all processors with default configuration.
-- Uses 'defaultShutdownConfig' (30 second drain timeout, 60 second graceful-phase bound).
-- For custom timeout, use 'stopAppGracefully'.
stopApp :: (IOE :> es) => AppHandle es -> Eff es ()
stopApp = void . stopAppGracefully defaultShutdownConfig

-- | Gracefully stop all processors with configurable drain and total timeout.
--
-- Shutdown sequence:
-- 1. Signal all adapters to stop producing (close source streams)
-- 2. Wait for processors to drain in-flight messages (with timeout)
-- 3. Force stop any remaining processors after timeout
-- 4. Stop the master coordinator
--
-- Returns whether all processors drained cleanly (True) or were forced (False).
stopAppGracefully :: (IOE :> es) => ShutdownConfig -> AppHandle es -> Eff es Bool
stopAppGracefully config appHandle =
  Exception.mask $ \restore -> do
    isLeader <-
      liftIO $
        atomically $ do
          started <- readTVar appHandle.shutdownStarted
          if started
            then pure False
            else writeTVar appHandle.shutdownStarted True >> pure True
    if isLeader
      then do
        result <- Exception.try @SomeException (restore performShutdown)
        liftIO $ atomically $ putTMVar appHandle.shutdownResult result
        either Exception.throwIO pure result
      else do
        result <- restore $ liftIO $ atomically $ readTMVar appHandle.shutdownResult
        either Exception.throwIO pure result
  where
    performShutdown = Exception.mask $ \restore -> do
      let totalTimeoutMicros = nominalToMicros config.totalShutdownTimeout
      outcome <-
        Exception.try $
          restore $
            withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO ->
              UIO.timeout totalTimeoutMicros (runInIO shutdownAndDrain)

      -- Master cleanup is unconditional: adapter failures, external cancellation,
      -- drain cancellation, and the total deadline all converge here.
      stopOutcome <- Exception.try @SomeException (stopMaster appHandle.master)
      case outcome of
        Left (primaryFailure :: SomeException) -> Exception.throwIO primaryFailure
        Right Nothing -> finishStopOutcome stopOutcome False
        Right (Just drained) -> finishStopOutcome stopOutcome drained

    shutdownAndDrain = do
      markMasterDraining appHandle.master
      drainAndResume appHandle.master

      -- Catch only synchronous adapter failures. External cancellation and the
      -- total timeout must abort this phase so the master is force-stopped.
      shutdownResults <-
        forM (Map.toList appHandle.processors) $ \(procId, processor) ->
          Exception.trySync (traceShutdown procId (shutdownAdapter processor))

      case [failure | Left failure <- shutdownResults] of
        firstFailure : _ -> Exception.throwIO firstFailure
        [] -> do
          let drainTimeoutMicros = nominalToMicros config.drainTimeout
          liftIO $ waitForDrainWithTimeout drainTimeoutMicros (Map.elems appHandle.processors)

    traceShutdown procId action = case appHandle.shutdownTracer of
      Nothing -> action
      Just tracer -> withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO ->
        UIO.mask $ \restore -> do
          context <- TraceContext.getContext
          ready <- newEmptyTMVarIO
          finished <- newEmptyTMVarIO
          -- SDK hooks catch SomeException, including cancellation. Keep them
          -- on a separate owned thread; only the caller runs adapter effects.
          -- Even if span creation swallows cancellation, it can only observe
          -- the published outcome, never start adapter cleanup afterwards.
          let trace = UIO.bracket (TraceContext.attachContext context) TraceContext.detachContext $ \_ ->
                OTel.inSpan' tracer (procId.unProcessorId <> " shutdown") internalSpanArgs $ \sp -> do
                  OTel.addAttribute sp attrShibuyaProcessorId procId.unProcessorId
                  atomically $ putTMVar ready sp
                  result <- atomically $ readTMVar finished
                  either IOException.throwIO pure result
          UIO.withAsync trace $ \worker -> do
            result <- IOException.try @SomeException $ restore $ do
              sp <-
                atomically $
                  readTMVar ready `orElse` (Async.waitCatchSTM worker >>= either throwSTM (const retry))
              UIO.bracket (TraceContext.attachContext $ TraceContext.insertSpan sp context) TraceContext.detachContext $ \_ ->
                runInIO action
            atomically $ putTMVar finished result
            case result of
              -- Do not wait for export after the caller has been cancelled.
              -- withAsync still cancels and joins the tracing worker.
              Left failure | Exception.isAsyncException failure -> IOException.throwIO failure
              _ -> do
                traced <- restore $ UIO.waitCatch worker
                either IOException.throwIO (const $ either IOException.throwIO pure traced) result

    shutdownAdapter (_, qp) = case qp of
      QueueProcessor {adapter} -> adapter.shutdown
      BatchingProcessor {adapter} -> adapter.shutdown

    finishStopOutcome (Left stopFailure) _ = Exception.throwIO stopFailure
    finishStopOutcome (Right ()) drained = pure drained

    nominalToMicros timeout = max 0 (floor (timeout * 1_000_000))

-- | Wait for all processors to be done, with timeout.
-- Returns True if all drained cleanly, False if timeout occurred.
-- Note: Requires -threaded RTS for registerDelay to work properly.
waitForDrainWithTimeout :: Int -> [(SupervisedProcessor, a)] -> IO Bool
waitForDrainWithTimeout timeoutMicros processors = do
  -- Create a timeout TVar that becomes True after the deadline
  timeoutVar <- registerDelay timeoutMicros

  -- Wait for either all done or timeout
  atomically $
    (allDone processors >> pure True)
      `orElse` (readTVar timeoutVar >>= check >> pure False)
  where
    allDone :: [(SupervisedProcessor, a)] -> STM ()
    allDone procs = forM_ procs $ \(sp, _) -> readTVar sp.done >>= check

-- | Wait for all processors to complete.
-- For infinite streams, this will block forever.
-- Use 'stopApp' to gracefully terminate.
--
-- Uses STM to block efficiently until all processors are done,
-- rather than polling.
waitApp :: (IOE :> es) => AppHandle es -> Eff es ()
waitApp appHandle = liftIO $ atomically $ do
  -- Block until all done TVars are True
  forM_ (Map.elems appHandle.processors) $ \(sp, _) ->
    readTVar sp.done >>= check

-- | Pause source admissions; already admitted polls and inbox work may drain.
pauseProcessor :: (IOE :> es) => AppHandle es -> ProcessorId -> Eff es ControlOutcome
pauseProcessor app = Master.pauseProcessor app.master

-- | Resume intake. Successful state is captured by the control transaction.
resumeProcessor :: (IOE :> es) => AppHandle es -> ProcessorId -> Eff es ControlOutcome
resumeProcessor app = Master.resumeProcessor app.master

-- | Nothing for unknown, terminal, draining or uncontrollable processors.
isProcessorPaused :: (IOE :> es) => AppHandle es -> ProcessorId -> Eff es (Maybe Bool)
isProcessorPaused app = Master.isProcessorPaused app.master