diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,106 @@ # Changelog +## 0.8.0.0 — 2026-07-04 + +### Breaking Changes + +- New application code should import the `Shibuya` umbrella module. It + re-exports the stable app, handler, batch, retry, policy, metrics, and + tracing surface. `Shibuya.Core` remains as a deprecated compatibility + re-export for this release. +- Runner internals moved under `Shibuya.Internal.Runner.*` and carry no PVP + stability guarantee. The public moves are: + `Shibuya.Runner.Metrics` -> `Shibuya.Core.Metrics`; + `Shibuya.Runner.{Master,Supervised,Batcher,BatchProcessor}` -> + `Shibuya.Internal.Runner.*`; `Shibuya.Runner.{Halt,Ingester}` -> + `Shibuya.Internal.Runner.*` as hidden implementation modules. + `Shibuya.Runner.Serial` and `Shibuya.Runner.Processor` were deleted. + `Shibuya.Prelude` is no longer exposed. +- `AppHandle` and `Master` are now opaque in public modules. Use + `getAppMaster`, `getAppMetrics`, `waitApp`, `stopApp`, and + `stopAppGracefully` instead of record fields. Metrics servers should import + `Master`, `getAllMetricsIO`, and `getProcessorMetricsIO` from + `Shibuya.App`. +- `runApp` now takes an `AppConfig` record and validates `inboxSize`. + Invalid sizes return `Left (AppConfigInvalid (InvalidInboxSize n))` + instead of crashing or stalling. + + ```haskell + -- before + runApp IgnoreFailures 100 processors + + -- after + runApp defaultAppConfig processors + + -- custom strategy/size + runApp defaultAppConfig {strategy = IgnoreFailures, inboxSize = 100} processors + ``` + +- Handlers now receive `Message es msg` instead of `Ingested es msg`. + Handlers that only use `.envelope` or `.lease` need a signature update; + handlers can no longer call the adapter finalizer directly. The framework + owns finalization and calls `AckHandle.finalize` after the handler returns + or throws. Adapter finalizers should remain idempotent or phase-tracked. +- `BatchHandler` likewise receives `NonEmpty (Message es msg)` instead of + `NonEmpty (Ingested es msg)`. +- Removed dead or misleading surface: `HandlerError.HandlerTimeout`, + `RuntimeError.InboxOverflow`, `StreamStats.dropped`, `incDropped`, + the always-zero `shibuya_messages_dropped_total` Prometheus metric, and + `Shibuya.Telemetry.Config`. +- Renamed `Ordering` to `OrderingPolicy`, eliminating the need for + `import Prelude hiding (Ordering)`. `runSupervised` also takes the ordering + policy before `Concurrency`. +- `BatchingProcessor` now rejects `PartitionedInOrder` combined with `Ahead` + or `Async`, because batches are scheduled by `BatchKey`, not + `Envelope.partition`. +- Renamed `Shibuya.Stream.batchStream` to `chunksOf`. + +### New Features + +- Added `mkEnvelope :: MessageId -> msg -> Envelope msg`. It fills optional + metadata with defaults; set optional fields with record update syntax. This + is the recommended construction path because future optional `Envelope` + fields can then be added without forcing adapter code to change, avoiding + the kind of major bumps required by 0.5.0.0 and 0.7.0.0. +- Added `mkIngested :: Envelope msg -> AckHandle es -> Ingested es msg`. +- Added `Message es msg`, the read-only handler view containing envelope and + optional lease. +- `PartitionedInOrder` with `Ahead` or `Async` is now enforced for + single-message processors. Messages sharing an `Envelope.partition` key are + processed and acknowledged in arrival order, distinct partitions run + concurrently up to the configured bound, and messages with no partition key + are unconstrained. + +### Bug Fixes + +- Handler exceptions on the single-message runner path now finalize the message + with `AckRetry (RetryDelay 0)` using the same bounded finalizer retry helper + as batch processing, so adapters observe a disposition instead of silently + losing or stranding the delivery. Exhausted finalizer retry now halts the + processor loudly with the failed message id. +- Batch accumulation failures, including exceptions from user-provided + `batchKey` functions, now fail the processor loudly instead of letting the + batcher consumer die while the processor reports clean completion. +- Batch processing now isolates `AckHalt`: batches already emitted after a halt + no longer re-enter the user batch handler, and their messages are finalized + with `AckRetry (RetryDelay 0)` instead of being left unacknowledged. +- The keyed batch scheduler now keeps its pending queue bounded and owns its + reader and worker threads with structured cleanup, preventing unbounded pulls + under slow handlers and preventing finalization after forced shutdown returns. + +### Other Changes + +- Corrected `Ahead` documentation. It preserves stream-yield order only; + handler execution and acknowledgement may complete in any order. +- `TraceHeaders` remains a type alias with the same representation as + `Headers`; its docs now clarify that it is only the parsed W3C trace-context + projection, while `Headers` is the complete broker header set. +- `shibuya-pgmq-adapter` and `shibuya-kafka-adapter` should migrate envelope + construction to `mkEnvelope` in their next releases. +- Removed unused dependencies from `shibuya-core`: `effectful-core`, + `generic-lens`, `lens`, `uuid`, and `vector`. Internal unlift imports now use + the stable top-level `Effectful` exports. + ## 0.7.1.0 — 2026-06-15 ### Bug Fixes diff --git a/shibuya-core.cabal b/shibuya-core.cabal --- a/shibuya-core.cabal +++ b/shibuya-core.cabal @@ -1,6 +1,6 @@ cabal-version: 3.12 name: shibuya-core -version: 0.7.1.0 +version: 0.8.0.0 synopsis: Supervised queue processing framework for Haskell description: A supervised queue processing framework inspired by Broadway (Elixir). @@ -20,35 +20,39 @@ library import: warnings exposed-modules: + Shibuya Shibuya.Adapter Shibuya.Adapter.Mock Shibuya.App + Shibuya.Batch Shibuya.Core Shibuya.Core.Ack Shibuya.Core.AckHandle Shibuya.Core.Error Shibuya.Core.Ingested Shibuya.Core.Lease + Shibuya.Core.Metrics Shibuya.Core.Retry Shibuya.Core.Types Shibuya.Handler + Shibuya.Internal.App + Shibuya.Internal.Runner.BatchProcessor + Shibuya.Internal.Runner.Batcher + Shibuya.Internal.Runner.Finalize + Shibuya.Internal.Runner.KeyedScheduler + Shibuya.Internal.Runner.Master + Shibuya.Internal.Runner.Supervised Shibuya.Policy - Shibuya.Prelude - Shibuya.Runner.Master - Shibuya.Runner.Metrics - Shibuya.Runner.Supervised Shibuya.Stream Shibuya.Telemetry - Shibuya.Telemetry.Config Shibuya.Telemetry.Effect Shibuya.Telemetry.Propagation Shibuya.Telemetry.Semantic other-modules: - Shibuya.Runner.Halt - Shibuya.Runner.Ingester - Shibuya.Runner.Processor - Shibuya.Runner.Serial + Shibuya.Internal.Runner.Halt + Shibuya.Internal.Runner.Ingester + Shibuya.Prelude default-extensions: DeriveAnyClass @@ -63,17 +67,15 @@ build-depends: aeson ^>=2.2, + atomic-primops ^>=0.8.8, base ^>=4.21.0.0, bytestring ^>=0.12.2.0, containers ^>=0.7, deepseq ^>=1.5, effectful ^>=2.6.1.0, - effectful-core ^>=2.6.1.0, - generic-lens ^>=2.3.0.0, hs-opentelemetry-api ^>=1.0, hs-opentelemetry-propagator-w3c ^>=1.0, hs-opentelemetry-semantic-conventions ^>=1.40, - lens ^>=5.3.5, nqe ^>=0.6, random ^>=1.2, stm ^>=2.5, @@ -83,8 +85,6 @@ time ^>=1.14, unliftio ^>=0.2, unordered-containers ^>=0.2, - uuid ^>=1.3, - vector ^>=0.13, hs-source-dirs: src default-language: GHC2024 @@ -96,6 +96,7 @@ hs-source-dirs: test main-is: Main.hs default-extensions: + DataKinds DeriveAnyClass DerivingStrategies DuplicateRecordFields @@ -112,10 +113,18 @@ -with-rtsopts=-N other-modules: + Shibuya.App.BatchSpec + Shibuya.App.LifecycleSpec + Shibuya.Batch.ReliabilitySpec + Shibuya.Batch.TestHarness + Shibuya.BatchSpec Shibuya.Core.AckSpec Shibuya.Core.RetrySpec Shibuya.Core.TypesSpec Shibuya.PolicySpec + Shibuya.Runner.BatchProcessorSpec + Shibuya.Runner.BatcherSpec + Shibuya.Runner.PartitionOrderingSpec Shibuya.Runner.SupervisedSpec Shibuya.RunnerSpec Shibuya.Telemetry.EffectSpec @@ -126,6 +135,7 @@ QuickCheck ^>=2.15, base ^>=4.21.0.0, bytestring, + containers, effectful, hs-opentelemetry-api, hs-opentelemetry-exporter-in-memory ^>=1.0, @@ -139,4 +149,3 @@ time, unliftio, unordered-containers, - vector, diff --git a/src/Shibuya.hs b/src/Shibuya.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya.hs @@ -0,0 +1,136 @@ +-- | The Shibuya framework: supervised queue processing with explicit acks. +-- This is the single import an application author needs. +module Shibuya + ( -- * Running an application + runApp, + AppConfig (..), + defaultAppConfig, + AppError (..), + QueueProcessor (..), + mkProcessor, + mkBatchProcessor, + AppHandle, + getAppMetrics, + getAppMaster, + waitApp, + stopApp, + stopAppGracefully, + ShutdownConfig (..), + defaultShutdownConfig, + SupervisionStrategy (..), + + -- * Messages and envelopes + MessageId (..), + Cursor (..), + Attempt (..), + Envelope (..), + mkEnvelope, + Headers, + TraceHeaders, + Message (..), + + -- * Handlers and acks + Handler, + AckDecision (..), + RetryDelay (..), + DeadLetterReason (..), + HaltReason (..), + ProcessorHalt (..), + + -- * Batch processing + BatchHandler, + BatchConfig (..), + defaultBatchConfig, + BatchKey (..), + defaultBatchKey, + BatchInfo (..), + BatchTrigger (..), + BatchAck (..), + ackAllOk, + ackAll, + ackExcept, + withFallback, + failMessages, + BatchConfigError (..), + validateBatchConfig, + + -- * Retry helpers + module Shibuya.Core.Retry, + + -- * Policies + OrderingPolicy (..), + Concurrency (..), + validatePolicy, + + -- * Adapter authoring + Adapter (..), + AckHandle (..), + Lease (..), + Ingested (..), + mkIngested, + toMessage, + + -- * Errors + PolicyError (..), + HandlerError (..), + RuntimeError (..), + ConfigError (..), + + -- * Metrics and introspection + Master, + ProcessorId (..), + ProcessorState (..), + ProcessorMetrics (..), + StreamStats (..), + BatchStats (..), + InFlightInfo (..), + MetricsMap, + + -- * Tracing + Tracing, + runTracing, + runTracingNoop, + ) +where + +import Shibuya.Adapter (Adapter (..)) +import Shibuya.App + ( AppConfig (..), + AppError (..), + AppHandle, + Master, + QueueProcessor (..), + ShutdownConfig (..), + SupervisionStrategy (..), + defaultAppConfig, + defaultShutdownConfig, + getAppMaster, + getAppMetrics, + mkBatchProcessor, + mkProcessor, + runApp, + stopApp, + stopAppGracefully, + waitApp, + ) +import Shibuya.Batch +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..), RetryDelay (..)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Error (ConfigError (..), HandlerError (..), PolicyError (..), RuntimeError (..)) +import Shibuya.Core.Ingested (Ingested (..), Message (..), mkIngested, toMessage) +import Shibuya.Core.Lease (Lease (..)) +import Shibuya.Core.Metrics + ( BatchStats (..), + InFlightInfo (..), + MetricsMap, + ProcessorId (..), + ProcessorMetrics (..), + ProcessorState (..), + StreamStats (..), + ) +import Shibuya.Core.Retry +import Shibuya.Core.Types (Attempt (..), Cursor (..), Envelope (..), Headers, MessageId (..), TraceHeaders, mkEnvelope) +import Shibuya.Handler (Handler) +import Shibuya.Internal.Runner.Halt (ProcessorHalt (..)) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..), validatePolicy) +import Shibuya.Telemetry.Effect (Tracing, runTracing, runTracingNoop) diff --git a/src/Shibuya/Adapter/Mock.hs b/src/Shibuya/Adapter/Mock.hs --- a/src/Shibuya/Adapter/Mock.hs +++ b/src/Shibuya/Adapter/Mock.hs @@ -9,16 +9,18 @@ newTrackingAck, trackingAckHandle, getTrackedDecisions, + mkTrackedIngested, + trackedListAdapter, ) where -import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef) import Effectful (Eff, IOE, liftIO, (:>)) import Shibuya.Adapter (Adapter (..)) import Shibuya.Core.Ack (AckDecision) import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Ingested (Ingested) -import Shibuya.Core.Types (MessageId (..)) +import Shibuya.Core.Ingested (Ingested, mkIngested) +import Shibuya.Core.Types (Envelope (..), MessageId (..)) import Streamly.Data.Stream qualified as Stream -- | Create an adapter from a list of ingested messages. @@ -45,7 +47,7 @@ AckHandle es trackingAckHandle tracking msgId = AckHandle $ \decision -> - liftIO $ modifyIORef' tracking.trackedDecisions ((msgId, decision) :) + liftIO $ atomicModifyIORef' tracking.trackedDecisions (\xs -> ((msgId, decision) : xs, ())) -- | Create a new TrackingAck. newTrackingAck :: (IOE :> es) => Eff es TrackingAck @@ -54,3 +56,19 @@ -- | Get all tracked decisions. getTrackedDecisions :: (IOE :> es) => TrackingAck -> Eff es [(MessageId, AckDecision)] getTrackedDecisions tracking = liftIO $ readIORef tracking.trackedDecisions + +-- | Wrap an envelope into an ingested message whose acknowledgement is recorded by the +-- given tracker, keyed by the envelope's own message id. The lease is +-- @Nothing@. Every call to the resulting handle's @finalize@ appends one +-- @(messageId, decision)@ pair to the tracking list, so duplicate finalizes are +-- observable. +mkTrackedIngested :: (IOE :> es) => TrackingAck -> Envelope msg -> Ingested es msg +mkTrackedIngested tracking env = + mkIngested env (trackingAckHandle tracking env.messageId) + +-- | Build an adapter from a list of envelopes where every message's acknowledgement +-- is recorded into one shared tracker. Combine with 'getTrackedDecisions' to +-- assert one successful finalization per message across a normal run. +trackedListAdapter :: (IOE :> es) => TrackingAck -> [Envelope msg] -> Adapter es msg +trackedListAdapter tracking envs = + listAdapter (map (mkTrackedIngested tracking) envs) diff --git a/src/Shibuya/App.hs b/src/Shibuya/App.hs --- a/src/Shibuya/App.hs +++ b/src/Shibuya/App.hs @@ -2,9 +2,12 @@ module Shibuya.App ( -- * Running Processors runApp, + AppConfig (..), + defaultAppConfig, QueueProcessor (..), mkProcessor, - AppHandle (..), + mkBatchProcessor, + AppHandle, -- * AppHandle Operations getAppMetrics, @@ -22,7 +25,15 @@ -- * Errors AppError (..), + Master, + getAllMetrics, + getAllMetricsIO, + getProcessorMetrics, + getProcessorMetricsIO, + -- * Batch API (re-exported from "Shibuya.Batch") + module Shibuya.Batch, + -- * Re-exports ProcessorId (..), ProcessorMetrics (..), @@ -32,8 +43,8 @@ import Control.Concurrent.NQE.Supervisor qualified as NQE import Control.Concurrent.STM (STM, atomically, check, orElse, readTVar, registerDelay) import Control.Monad (forM_, void) +import Data.Bifunctor (first) import Data.Foldable (traverse_) -import Data.Map.Strict (Map) import Data.Map.Strict qualified as Map import Data.Text qualified as Text import Data.Time.Clock (NominalDiffTime) @@ -41,27 +52,31 @@ import GHC.Generics (Generic) import Numeric.Natural (Natural) import Shibuya.Adapter (Adapter (..)) -import Shibuya.Core.Error (HandlerError (..), PolicyError (..), RuntimeError (..)) -import Shibuya.Handler (Handler) -import Shibuya.Policy (Concurrency (..), Ordering (..), validatePolicy) -import Shibuya.Runner.Master +import Shibuya.Batch +import Shibuya.Core.Error (ConfigError (..), HandlerError (..), PolicyError (..), RuntimeError (..)) +import Shibuya.Core.Metrics + ( MetricsMap, + ProcessorId (..), + ProcessorMetrics (..), + ) +import Shibuya.Internal.App (AppHandle (..), QueueProcessor (..), mkBatchProcessor, mkProcessor) +import Shibuya.Internal.Runner.Master ( Master, getAllMetrics, + getAllMetricsIO, + getProcessorMetrics, + getProcessorMetricsIO, startMaster, stopMaster, ) -import Shibuya.Runner.Metrics - ( MetricsMap, - ProcessorId (..), - ProcessorMetrics (..), - ) -import Shibuya.Runner.Supervised +import Shibuya.Internal.Runner.Supervised ( SupervisedProcessor (..), runSupervised, + runSupervisedBatch, ) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..), validatePolicy) import Shibuya.Telemetry.Effect (Tracing) -import UnliftIO (SomeException, catch, displayException) -import Prelude hiding (Ordering) +import UnliftIO (SomeException, catch, displayException, try) -------------------------------------------------------------------------------- -- Supervision Strategy @@ -77,6 +92,8 @@ 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) @@ -84,7 +101,7 @@ toNQEStrategy :: SupervisionStrategy -> NQE.Strategy toNQEStrategy = \case IgnoreFailures -> NQE.IgnoreAll - StopAllOnFailure -> NQE.KillAll + StopAllOnFailure -> NQE.IgnoreGraceful -------------------------------------------------------------------------------- -- Shutdown Configuration @@ -116,33 +133,25 @@ AppHandlerError !HandlerError | -- | Runtime error AppRuntimeError !RuntimeError + | -- | Invalid batch configuration for a 'BatchingProcessor' + AppBatchConfigError !BatchConfigError + | -- | Invalid application configuration + AppConfigInvalid !ConfigError deriving stock (Eq, Show) --- | A queue processor pairs an adapter with its handler. --- The message type is existentially hidden, allowing heterogeneous queues. -data QueueProcessor es where - QueueProcessor :: - { adapter :: Adapter es msg, - handler :: Handler es msg, - ordering :: Ordering, - concurrency :: Concurrency - } -> - QueueProcessor es - --- | Convenience constructor with default policies (Unordered + Serial). --- Provides backward compatibility with existing code. -mkProcessor :: Adapter es msg -> Handler es msg -> QueueProcessor es -mkProcessor adapter handler = QueueProcessor adapter handler Unordered Serial - --- | Handle for a running multi-queue application. --- Provides introspection and control over all processors. -data AppHandle es = AppHandle - { -- | The master coordinator - master :: !Master, - -- | Map of processor IDs to their handles - processors :: !(Map ProcessorId (SupervisedProcessor, QueueProcessor es)) +-- | Configuration for @runApp@. +data AppConfig = AppConfig + { -- | How processor failures affect siblings. + strategy :: !SupervisionStrategy, + -- | Bounded-inbox capacity per processor (backpressure). Must be >= 1. + inboxSize :: !Int } + deriving stock (Eq, Show, Generic) +-- | 'IgnoreFailures' with an inbox of 100. +defaultAppConfig :: AppConfig +defaultAppConfig = AppConfig {strategy = IgnoreFailures, inboxSize = 100} + -- | Run queue processors concurrently under NQE supervision. -- -- Each processor runs independently. Returns immediately with a handle @@ -151,51 +160,74 @@ -- Example: -- -- @ --- result <- runApp IgnoreFailures 100 +-- result <- runApp defaultAppConfig -- [ ("orders", QueueProcessor ordersAdapter ordersHandler) -- , ("events", QueueProcessor eventsAdapter eventsHandler) -- ] -- @ runApp :: (IOE :> es, Tracing :> es) => - -- | Supervision strategy - SupervisionStrategy -> - -- | Inbox size for backpressure - Int -> + -- | Application configuration + AppConfig -> -- | Named processors [(ProcessorId, QueueProcessor es)] -> Eff es (Either AppError (AppHandle es)) -runApp strategy inboxSize namedProcessors = - -- Validate all policies first - case validateAllPolicies namedProcessors of - Left err -> pure $ Left $ AppPolicyError err +runApp config namedProcessors = + -- Validate all policies (and batch configs) first + case validateAppConfig config *> validateAllPolicies namedProcessors of + Left err -> pure $ Left err Right () -> do - let nqeStrategy = toNQEStrategy strategy + let nqeStrategy = toNQEStrategy config.strategy catch ( do - -- Start the master coordinator master <- startMaster nqeStrategy - - -- Spawn each processor under supervision - processors <- spawnProcessors master (fromIntegral inboxSize) namedProcessors - - pure $ - Right - AppHandle - { master = master, - processors = Map.fromList processors - } + spawnResult <- try $ spawnProcessors master (fromIntegral config.inboxSize) namedProcessors + case spawnResult of + Left (e :: SomeException) -> do + stopMaster master + pure $ Left $ AppRuntimeError $ SupervisorFailed $ Text.pack $ displayException e + Right processors -> + pure $ + Right + AppHandle + { master = master, + processors = Map.fromList processors + } ) ( \(e :: SomeException) -> pure $ Left $ AppRuntimeError $ SupervisorFailed $ Text.pack $ displayException e ) --- | Validate all processor policies before starting. -validateAllPolicies :: [(ProcessorId, QueueProcessor es)] -> Either PolicyError () +-- | Validate app configuration before starting any processor. +validateAppConfig :: AppConfig -> Either AppError () +validateAppConfig config + | config.inboxSize < 1 = Left $ AppConfigInvalid $ InvalidInboxSize config.inboxSize + | otherwise = Right () + +-- | Validate all processor policies (and batch configs) before starting. +validateAllPolicies :: [(ProcessorId, QueueProcessor es)] -> Either AppError () validateAllPolicies = traverse_ validateOne where - validateOne (_, QueueProcessor _ _ ord conc) = validatePolicy ord conc + 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) => @@ -205,9 +237,21 @@ Eff es [(ProcessorId, (SupervisedProcessor, QueueProcessor es))] spawnProcessors master inboxSize = traverse spawnOne where - spawnOne (procId, qp@(QueueProcessor adapter handler _ordering concurrency)) = do - sp <- runSupervised master inboxSize procId concurrency adapter handler - pure (procId, (sp, qp)) + spawnOne (procId, qp) = case qp of + QueueProcessor {adapter, handler, ordering, concurrency} -> do + sp <- runSupervised master inboxSize procId ordering concurrency adapter handler + pure (procId, (sp, qp)) + BatchingProcessor {adapter, batchHandler, batchConfig, concurrency} -> do + sp <- + runSupervisedBatch + master + inboxSize + procId + concurrency + batchConfig + adapter + batchHandler + pure (procId, (sp, qp)) -------------------------------------------------------------------------------- -- AppHandle Operations @@ -254,7 +298,9 @@ pure drained where - shutdownAdapter (_, QueueProcessor adapter _ _ _) = adapter.shutdown + shutdownAdapter (_, qp) = case qp of + QueueProcessor {adapter} -> adapter.shutdown + BatchingProcessor {adapter} -> adapter.shutdown -- | Wait for all processors to be done, with timeout. -- Returns True if all drained cleanly, False if timeout occurred. diff --git a/src/Shibuya/Batch.hs b/src/Shibuya/Batch.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Batch.hs @@ -0,0 +1,187 @@ +-- | Public vocabulary for batch processing. +-- +-- This module defines the types a user needs to opt a processor into batching: +-- the batch handler, its configuration, the grouping key, and the batch +-- acknowledgement result. It adds no runtime behavior; the accumulation engine +-- and execution stage live in the internal @Shibuya.Internal.Runner.*@ modules. +-- +-- == Acknowledgement decision contract +-- +-- Given an emitted batch and the @BatchAck@ a @BatchHandler@ returns, the +-- framework resolves one @AckDecision@ for /every/ message in its own +-- retained batch list. For each retained message it looks the message's +-- @MessageId@ up in @decisions@; if the id is absent it uses @fallback@. The +-- handler's return value only /supplies decisions/ — it never drives which +-- messages are acked. Consequently decision resolution is complete and +-- deterministic regardless of what the handler returns (wrong length, +-- reordered, missing ids all degrade gracefully to the fallback). This +-- requires @MessageId@s to be unique within a batch, which holds for every +-- real adapter and the mock adapter. The runtime execution stage applies each +-- resolved decision through the message's idempotent @AckHandle.finalize@ with +-- bounded retries. +module Shibuya.Batch + ( -- * Grouping key + BatchKey (..), + defaultBatchKey, + + -- * Emission trigger + BatchTrigger (..), + + -- * Batch metadata + BatchInfo (..), + + -- * Configuration + BatchConfig (..), + defaultBatchConfig, + BatchConfigError (..), + validateBatchConfig, + + -- * Handler + BatchHandler, + + -- * Acknowledgement result + BatchAck (..), + ackAllOk, + ackAll, + ackExcept, + withFallback, + failMessages, + ) +where + +import Control.DeepSeq (NFData) +import Data.List.NonEmpty (NonEmpty) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.String (IsString) +import Effectful (Eff) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason) +import Shibuya.Core.Ingested (Message) +import Shibuya.Core.Types (Envelope, MessageId) +import Shibuya.Prelude + +-- | Groups messages into independent sub-batches within one processor. +-- Messages sharing a key accumulate together; each key has its own size +-- counter and timeout. Compute one per message via the batch config's @batchKey@. +newtype BatchKey = BatchKey {unBatchKey :: Text} + deriving stock (Eq, Ord, Show, Generic) + deriving newtype (IsString) + deriving anyclass (NFData) + +-- | The key used when a configuration does not distinguish sub-batches. +defaultBatchKey :: BatchKey +defaultBatchKey = BatchKey "default" + +-- | Why the framework emitted a batch. +data BatchTrigger + = -- | Reached the configured @batchSize@. + TriggerSize + | -- | @batchTimeout@ elapsed since the batch's first message arrived. + TriggerTimeout + | -- | The processor is draining/shutting down; a partial batch was flushed. + TriggerFlush + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Metadata about an emitted batch, passed to the batch handler alongside +-- the messages. +data BatchInfo = BatchInfo + { -- | The key all messages in this batch share. + batchKey :: !BatchKey, + -- | How many messages are in this batch (always >= 1). + size :: !Int, + -- | Why this batch was emitted. + trigger :: !BatchTrigger, + -- | Partition of the batch's first message, if the envelope had one. + partition :: !(Maybe Text) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (NFData) + +-- | Configuration for a batching processor. +-- +-- @es@ is the effect stack and @msg@ the payload type, matching the +-- batch handler this config is paired with. +data BatchConfig es msg = BatchConfig + { -- | Emit a batch once it holds this many messages. Must be >= 1. + batchSize :: !Int, + -- | Emit a batch this long after its first message arrives, even if it + -- has not reached @batchSize@. Must be > 0. + batchTimeout :: !NominalDiffTime, + -- | Compute a message's sub-batch key from its envelope. Use + -- @const defaultBatchKey@ for a single undivided batch. + batchKey :: !(Envelope msg -> BatchKey), + -- | How often the timeout ticker scans for timed-out batches. @Nothing@ + -- means "use @batchTimeout@". Flush latency is bounded by this interval. + -- Must be > 0 when @Just@. + tickInterval :: !(Maybe NominalDiffTime) + } + +-- | Batch of at most 100, flushed after 1 second, one undivided sub-batch, +-- ticker granularity equal to the timeout. Matches Broadway's defaults +-- (@batch_size: 100@, @batch_timeout: 1000@ ms). +defaultBatchConfig :: BatchConfig es msg +defaultBatchConfig = + BatchConfig + { batchSize = 100, + batchTimeout = 1, + batchKey = const defaultBatchKey, + tickInterval = Nothing + } + +-- | A batch handler processes a whole batch at once and returns how to +-- acknowledge it. Unlike 'Shibuya.Handler.Handler' (one message -> one +-- decision), a batch handler receives every message in the batch plus its +-- @BatchInfo@, and returns a single @BatchAck@ describing per-message +-- outcomes. See the module haddock for the acknowledgement decision contract. +type BatchHandler es msg = BatchInfo -> NonEmpty (Message es msg) -> Eff es BatchAck + +-- | How to acknowledge every message in a batch. The framework resolves each +-- retained message's decision by looking its @MessageId@ up in @decisions@, or +-- by using @fallback@ if absent. The execution stage applies those decisions to +-- the messages' idempotent finalizers. See the module haddock. +data BatchAck = BatchAck + { -- | Per-message overrides, keyed by @MessageId@. + decisions :: !(Map MessageId AckDecision), + -- | Decision for any message not present in 'decisions'. + fallback :: !AckDecision + } + deriving stock (Show, Generic) + +-- | Acknowledge every message as successfully processed. The common case. +ackAllOk :: BatchAck +ackAllOk = BatchAck Map.empty AckOk + +-- | Apply one decision to every message in the batch. +ackAll :: AckDecision -> BatchAck +ackAll = BatchAck Map.empty + +-- | Acknowledge everything @AckOk@ except the listed messages, which get their +-- given decisions. Use for partial failure within an otherwise-successful batch. +ackExcept :: [(MessageId, AckDecision)] -> BatchAck +ackExcept overrides = BatchAck (Map.fromList overrides) AckOk + +-- | Give the listed messages their decisions and everything else @fb@. +withFallback :: AckDecision -> [(MessageId, AckDecision)] -> BatchAck +withFallback fb overrides = BatchAck (Map.fromList overrides) fb + +-- | Dead-letter the listed messages (with reasons) and acknowledge the rest OK. +failMessages :: [(MessageId, DeadLetterReason)] -> BatchAck +failMessages fs = + BatchAck (Map.fromList [(mid, AckDeadLetter r) | (mid, r) <- fs]) AckOk + +-- | Why a batch config is invalid. +data BatchConfigError + = BatchSizeNotPositive !Int + | BatchTimeoutNotPositive !NominalDiffTime + | TickIntervalNotPositive !NominalDiffTime + deriving stock (Eq, Show, Generic) + +-- | Validate a batch configuration. @batchSize@ must be >= 1, @batchTimeout@ +-- must be > 0, and @tickInterval@ (when set) must be > 0. +validateBatchConfig :: BatchConfig es msg -> Either BatchConfigError () +validateBatchConfig cfg + | cfg.batchSize < 1 = Left (BatchSizeNotPositive cfg.batchSize) + | cfg.batchTimeout <= 0 = Left (BatchTimeoutNotPositive cfg.batchTimeout) + | Just t <- cfg.tickInterval, t <= 0 = Left (TickIntervalNotPositive t) + | otherwise = Right () diff --git a/src/Shibuya/Core.hs b/src/Shibuya/Core.hs --- a/src/Shibuya/Core.hs +++ b/src/Shibuya/Core.hs @@ -1,96 +1,5 @@ --- | Shibuya Core - Public API --- --- This module re-exports all public types and functions from the Shibuya framework. --- Import this module for application development. --- --- Example: --- --- @ --- import Shibuya.Core --- --- main = runEff $ do --- let processor = QueueProcessor myAdapter myHandler --- result <- runApp IgnoreFailures 100 [(ProcessorId "main", processor)] --- case result of --- Right handle -> waitApp handle --- Left err -> print err --- @ -module Shibuya.Core - ( -- * Message Types - MessageId (..), - Cursor (..), - Attempt (..), - Envelope (..), - Headers, - - -- * Ack Semantics - AckDecision (..), - RetryDelay (..), - DeadLetterReason (..), - HaltReason (..), - - -- * Handler Types - AckHandle (..), - Lease (..), - Ingested (..), - - -- * Handler - Handler, - - -- * Adapter - Adapter (..), - - -- * Policy - Ordering (..), - Concurrency (..), - validatePolicy, - - -- * App - runApp, - AppError (..), - QueueProcessor (..), - mkProcessor, - AppHandle (..), - waitApp, - stopApp, - stopAppGracefully, - getAppMetrics, - - -- * Shutdown Configuration - ShutdownConfig (..), - defaultShutdownConfig, - - -- * Supervision Strategy - SupervisionStrategy (..), - - -- * Errors - PolicyError (..), - HandlerError (..), - RuntimeError (..), - - -- * Metrics - ProcessorId (..), - ProcessorState (..), - ProcessorMetrics (..), - StreamStats (..), - InFlightInfo (..), - MetricsMap, - - -- * Halt - ProcessorHalt (..), - ) -where +-- | Compatibility re-export for the pre-1.0 public API. +-- Prefer importing "Shibuya" in new code. +module Shibuya.Core {-# DEPRECATED "Use Shibuya instead." #-} (module Shibuya) where -import Shibuya.Adapter (Adapter (..)) -import Shibuya.App (AppError (..), AppHandle (..), QueueProcessor (..), ShutdownConfig (..), SupervisionStrategy (..), defaultShutdownConfig, getAppMetrics, mkProcessor, runApp, stopApp, stopAppGracefully, waitApp) -import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..), RetryDelay (..)) -import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Error (HandlerError (..), PolicyError (..), RuntimeError (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Core.Lease (Lease (..)) -import Shibuya.Core.Types (Attempt (..), Cursor (..), Envelope (..), Headers, MessageId (..)) -import Shibuya.Handler (Handler) -import Shibuya.Policy (Concurrency (..), Ordering (..), validatePolicy) -import Shibuya.Runner.Halt (ProcessorHalt (..)) -import Shibuya.Runner.Metrics (InFlightInfo (..), MetricsMap, ProcessorId (..), ProcessorMetrics (..), ProcessorState (..), StreamStats (..)) -import Prelude hiding (Ordering) +import Shibuya diff --git a/src/Shibuya/Core/AckHandle.hs b/src/Shibuya/Core/AckHandle.hs --- a/src/Shibuya/Core/AckHandle.hs +++ b/src/Shibuya/Core/AckHandle.hs @@ -1,6 +1,12 @@ -- | AckHandle for mechanical acknowledgment. -- This is the Broadway.Acknowledger equivalent, but typed. --- Must be called exactly once; adapter enforces idempotency. +-- +-- The framework calls 'finalize' at most once per delivery with a resolved +-- decision, except when retrying the same decision after a transient finalizer +-- failure. Adapters must therefore make finalization idempotent or +-- phase-tracked: re-running a partially completed finalization with the same +-- decision must be safe. The framework never intentionally finalizes one +-- delivery with two different decisions. module Shibuya.Core.AckHandle ( AckHandle (..), ) @@ -13,6 +19,9 @@ -- The handler's AckDecision is passed to finalize, which performs -- the actual commit/retry/dead-letter operation. newtype AckHandle es = AckHandle - { -- | Finalize the message with the given decision + { -- | Finalize the message with the given decision. + -- + -- This action may be retried with the same decision under the framework's + -- bounded retry schedule when it throws. finalize :: AckDecision -> Eff es () } diff --git a/src/Shibuya/Core/Error.hs b/src/Shibuya/Core/Error.hs --- a/src/Shibuya/Core/Error.hs +++ b/src/Shibuya/Core/Error.hs @@ -12,9 +12,14 @@ -- * Runtime Errors RuntimeError (..), runtimeErrorToText, + + -- * Configuration Errors + ConfigError (..), + configErrorToText, ) where +import Data.Text qualified as Text import Shibuya.Prelude -- | Policy validation errors. @@ -31,24 +36,29 @@ data HandlerError = -- | Handler threw an exception HandlerException !Text - | -- | Handler timed out - HandlerTimeout deriving stock (Eq, Show, Generic) -- | Convert handler error to text for display. handlerErrorToText :: HandlerError -> Text handlerErrorToText (HandlerException msg) = msg -handlerErrorToText HandlerTimeout = "Handler timed out" -- | Runtime errors during processing. data RuntimeError = -- | Supervisor failed SupervisorFailed !Text - | -- | Inbox overflow (backpressure failure) - InboxOverflow deriving stock (Eq, Show, Generic) -- | Convert runtime error to text for display. runtimeErrorToText :: RuntimeError -> Text runtimeErrorToText (SupervisorFailed msg) = msg -runtimeErrorToText InboxOverflow = "Inbox overflow" + +-- | Application configuration errors, detected before any processor starts. +data ConfigError + = -- | inboxSize must be >= 1; 0 stalls ingestion, negatives are nonsense. + InvalidInboxSize !Int + deriving stock (Eq, Show, Generic) + +-- | Convert configuration error to text for display. +configErrorToText :: ConfigError -> Text +configErrorToText (InvalidInboxSize n) = + "inboxSize must be >= 1, got " <> Text.pack (show n) diff --git a/src/Shibuya/Core/Ingested.hs b/src/Shibuya/Core/Ingested.hs --- a/src/Shibuya/Core/Ingested.hs +++ b/src/Shibuya/Core/Ingested.hs @@ -3,6 +3,9 @@ -- Exactly one thing flows through the system. module Shibuya.Core.Ingested ( Ingested (..), + Message (..), + toMessage, + mkIngested, ) where @@ -10,8 +13,8 @@ import Shibuya.Core.Lease (Lease) import Shibuya.Core.Types (Envelope) --- | What handlers receive for processing. --- Contains the message envelope, ack handle, and optional lease. +-- | Framework-side message with the adapter-provided ack finalizer. +-- Adapters construct this; application handlers receive @Message@ instead. data Ingested es msg = Ingested { -- | Message metadata and payload envelope :: !(Envelope msg), @@ -20,3 +23,29 @@ -- | Optional lease for visibility timeout extension lease :: !(Maybe (Lease es)) } + +-- | The read-only view a handler receives: envelope plus optional lease, +-- deliberately without an ack handle. The framework owns finalization. +data Message es msg = Message + { -- | Message metadata and payload + envelope :: !(Envelope msg), + -- | Optional lease for visibility timeout extension + lease :: !(Maybe (Lease es)) + } + +-- | Project the framework-side value to the handler-facing view. +toMessage :: Ingested es msg -> Message es msg +toMessage ingested = + Message + { envelope = ingested.envelope, + lease = ingested.lease + } + +-- | Construct an ingested message with no lease. +mkIngested :: Envelope msg -> AckHandle es -> Ingested es msg +mkIngested envelope ack = + Ingested + { envelope = envelope, + ack = ack, + lease = Nothing + } diff --git a/src/Shibuya/Core/Metrics.hs b/src/Shibuya/Core/Metrics.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Core/Metrics.hs @@ -0,0 +1,420 @@ +-- | 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} diff --git a/src/Shibuya/Core/Retry.hs b/src/Shibuya/Core/Retry.hs --- a/src/Shibuya/Core/Retry.hs +++ b/src/Shibuya/Core/Retry.hs @@ -1,7 +1,7 @@ -- | Exponential backoff policy and pure evaluator. -- --- Handlers compute a 'RetryDelay' from a 'BackoffPolicy' and the current --- delivery 'Attempt'. The pure evaluator takes a jitter sample in @[0,1)@ +-- Handlers compute a retry delay from a backoff policy and the current +-- delivery attempt. The pure evaluator takes a jitter sample in @[0,1)@ -- and is suitable for property tests. module Shibuya.Core.Retry ( -- * Policy @@ -122,7 +122,7 @@ -- | One-line helper for the common case: read 'envelope.attempt', compute a -- backoff delay, and return 'AckRetry'. -- --- Treats 'Nothing' attempt as @'Attempt' 0@ (first delivery), so handlers +-- Treats @Nothing@ attempt as @Attempt 0@ (first delivery), so handlers -- consuming envelopes from adapters that do not track redeliveries still -- get a sensible base-delay retry. retryWithBackoff :: diff --git a/src/Shibuya/Core/Types.hs b/src/Shibuya/Core/Types.hs --- a/src/Shibuya/Core/Types.hs +++ b/src/Shibuya/Core/Types.hs @@ -13,6 +13,7 @@ -- * Message Envelope Envelope (..), + mkEnvelope, -- * Message Headers Headers, @@ -63,7 +64,11 @@ type Headers = [(ByteString, ByteString)] -- | W3C Trace Context headers for distributed tracing. --- Contains traceparent and optionally tracestate headers. +-- +-- This is a type alias, not a newtype: it deliberately shares the same +-- wire representation as 'Headers'. Use it only for the narrow parsed +-- trace-context projection (@traceparent@ and optional @tracestate@); +-- use 'Headers' for the full, non-lossy broker header set. type TraceHeaders = [(ByteString, ByteString)] -- | Normalized message envelope (Broadway.Message equivalent). @@ -82,12 +87,12 @@ -- | All message headers as delivered by the source broker, in -- order and including duplicates. -- - -- 'Nothing' means the adapter does not surface headers at all; - -- 'Just []' means the adapter surfaces headers and this message - -- carried none. The W3C trace headers ('traceparent' / - -- 'tracestate') appear here verbatim /in addition to/ their - -- parsed form in 'traceContext'; this field is the faithful, - -- non-lossy view and 'traceContext' is the narrow projection the + -- @Nothing@ means the adapter does not surface headers at all; + -- @Just []@ means the adapter surfaces headers and this message + -- carried none. The W3C trace headers (@traceparent@ / + -- @tracestate@) appear here verbatim /in addition to/ their + -- parsed form in @traceContext@; this field is the faithful, + -- non-lossy view and @traceContext@ is the narrow projection the -- framework uses to re-establish a parent span. headers :: !(Maybe Headers), -- | Optional zero-indexed delivery counter. @@ -115,11 +120,27 @@ } deriving stock (Eq, Show, Functor, Generic) --- | Manual 'NFData' so the @attributes@ field's 'Attribute' values do +-- | Construct an envelope from required fields, defaulting optional metadata. +-- Set optional fields with record-update syntax on the result. +mkEnvelope :: MessageId -> msg -> Envelope msg +mkEnvelope messageId payload = + Envelope + { messageId = messageId, + cursor = Nothing, + partition = Nothing, + enqueuedAt = Nothing, + traceContext = Nothing, + headers = Nothing, + attempt = Nothing, + attributes = mempty, + payload = payload + } + +-- | Manual 'NFData' for envelopes so the @attributes@ field's 'Attribute' values do -- not require an upstream NFData instance (which -- @hs-opentelemetry-api@ does not currently ship). Forces every other -- field deeply and reduces 'attributes' to WHNF — every 'Attribute' --- leaf is a small primitive ('Text', 'Bool', 'Double', 'Int64'), so +-- leaf is a small primitive (@Text@, @Bool@, @Double@, @Int64@), so -- WHNF is enough to evaluate the contained values when the HashMap -- is itself in WHNF. instance (NFData msg) => NFData (Envelope msg) where diff --git a/src/Shibuya/Handler.hs b/src/Shibuya/Handler.hs --- a/src/Shibuya/Handler.hs +++ b/src/Shibuya/Handler.hs @@ -8,8 +8,13 @@ import Effectful (Eff) import Shibuya.Core.Ack (AckDecision) -import Shibuya.Core.Ingested (Ingested) +import Shibuya.Core.Ingested (Message) -- | Handler function type. --- Takes an ingested message and returns an ack decision. -type Handler es msg = Ingested es msg -> Eff es AckDecision +-- Takes a read-only message view and returns an ack decision. +-- +-- If a handler throws, the framework records the failure and finalizes the +-- message with @AckRetry (RetryDelay 0)@ so the message is not lost. A handler +-- that needs a different disposition should catch its own exception and return +-- the desired ack decision. +type Handler es msg = Message es msg -> Eff es AckDecision diff --git a/src/Shibuya/Internal/App.hs b/src/Shibuya/Internal/App.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/App.hs @@ -0,0 +1,62 @@ +-- | __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. +module Shibuya.Internal.App + ( QueueProcessor (..), + mkProcessor, + mkBatchProcessor, + AppHandle (..), + ) +where + +import Data.Map.Strict (Map) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Batch (BatchConfig, BatchHandler) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Handler (Handler) +import Shibuya.Internal.Runner.Master (Master) +import Shibuya.Internal.Runner.Supervised (SupervisedProcessor) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) + +-- | A queue processor pairs an adapter with a handler. The message type is +-- existentially hidden, allowing heterogeneous queues in one @runApp@ call. +-- +-- @QueueProcessor@ processes one message at a time; @BatchingProcessor@ groups +-- messages into batches (see "Shibuya.Batch") and runs a batch handler over each. +data QueueProcessor es where + QueueProcessor :: + { adapter :: Adapter es msg, + handler :: Handler es msg, + ordering :: OrderingPolicy, + concurrency :: Concurrency + } -> + QueueProcessor es + BatchingProcessor :: + { adapter :: Adapter es msg, + batchHandler :: BatchHandler es msg, + batchConfig :: BatchConfig es msg, + ordering :: OrderingPolicy, + concurrency :: Concurrency + } -> + QueueProcessor es + +-- | Convenience constructor with default policies (Unordered + Serial). +-- Provides backward compatibility with existing code. +mkProcessor :: Adapter es msg -> Handler es msg -> QueueProcessor es +mkProcessor adapter handler = QueueProcessor adapter handler Unordered Serial + +-- | Convenience constructor for a batching processor with safe default policies +-- (Unordered ordering + Serial concurrency, i.e. one batch at a time). +mkBatchProcessor :: + Adapter es msg -> BatchHandler es msg -> BatchConfig es msg -> QueueProcessor es +mkBatchProcessor adapter batchHandler batchConfig = + BatchingProcessor adapter batchHandler batchConfig Unordered Serial + +-- | Handle for a running multi-queue application. +-- Provides introspection and control over all processors. +data AppHandle es = AppHandle + { -- | The master coordinator + master :: !Master, + -- | Map of processor IDs to their handles + processors :: !(Map ProcessorId (SupervisedProcessor, QueueProcessor es)) + } diff --git a/src/Shibuya/Internal/Runner/BatchProcessor.hs b/src/Shibuya/Internal/Runner/BatchProcessor.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/BatchProcessor.hs @@ -0,0 +1,324 @@ +-- | __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 Data.Foldable (for_, traverse_) +import Data.HashMap.Strict qualified as HashMap +import Data.IORef (IORef, atomicWriteIORef, newIORef, readIORef) +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 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.Types (Envelope (..)) +import Shibuya.Internal.Runner.Finalize (finalizeWithRetry) +import Shibuya.Internal.Runner.Halt (ProcessorHalt (..)) +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, throwIO) + +-- | 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 -> + IORef (Maybe HaltReason) -> + BatchHandler es msg -> + (BatchInfo, NonEmpty (Ingested es msg)) -> + Eff es () +processOneBatch metricsHandle procId maxConc haltRef 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 + + -- Increment in-flight (a batch counts as one in-flight unit) and report it. + currentInflight <- liftIO $ beginProcessing metricsHandle maxConc + addAttribute traceSpan attrShibuyaInflightCount currentInflight + addAttribute traceSpan attrShibuyaInflightMax maxConc + + addEvent traceSpan (mkEvent eventBatchStarted []) + + alreadyHalted <- liftIO $ readIORef haltRef + + -- Run the handler under exception isolation. On any exception, record it + -- on the span and substitute the whole-batch retry default. + (handlerResult, skippedAfterHalt) <- + case alreadyHalted 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 + 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 finalizationHalt = + case finalizeFailures of + [] -> Nothing + failed -> + Just $ + HaltFatal $ + "batch finalization failed for message ids: " + <> Text.intercalate ", " [tshow mid | (mid, _) <- failed] + firstHalt = finalizationHalt <|> listToMaybe [r | AckHalt r <- decisions] + 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 firstHalt of + Just reason -> setStatus traceSpan (OTel.Error (haltReasonText reason)) + 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) + (haltReasonText <$> firstHalt) + + -- Halt: set the shared flag; do NOT throw (let the stream drain). + for_ firstHalt $ \reason -> + liftIO $ atomicWriteIORef haltRef (Just reason) + where + isFailing :: AckDecision -> Bool + isFailing (AckDeadLetter _) = True + isFailing (AckRetry _) = True + isFailing _ = False + +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 + +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)) -> + IORef (Maybe HaltReason) -> + Eff es () +processBatchesUntilDrained metricsHandle procId concurrency handler batchesStream haltRef = do + let maxConc = case concurrency of + Serial -> 1 + Ahead n -> n + Async n -> n + + withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do + let batchAction = runInIO . processOneBatch metricsHandle procId maxConc haltRef handler + pendingLimit = max 2 (2 * max 1 maxConc) + case concurrency of + Serial -> + Stream.fold Fold.drain $ + Stream.mapM batchAction batchesStream + Ahead n -> + runKeyedScheduler (max 1 n) pendingLimit (Just . fstBatchKey) batchAction batchesStream + Async n -> + runKeyedScheduler (max 1 n) pendingLimit (Just . fstBatchKey) batchAction batchesStream + +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 + haltRef <- liftIO $ newIORef Nothing + + let batchesStream = Stream.fromList batches + processBatchesUntilDrained metricsHandle procId concurrency handler batchesStream haltRef + + maybeHalt <- liftIO $ readIORef haltRef + case maybeHalt of + Just reason -> throwIO (ProcessorHalt reason) + Nothing -> liftIO $ sampleMetrics metricsHandle diff --git a/src/Shibuya/Internal/Runner/Batcher.hs b/src/Shibuya/Internal/Runner/Batcher.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Batcher.hs @@ -0,0 +1,268 @@ +-- | __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 accumulation engine. +-- +-- Groups a stream of individual ingested messages into batches by batch key. +-- A batch is emitted when it reaches the configured size, when its per-key +-- timeout elapses, or when the input ends / the processor drains. This module +-- only /regroups/ messages: it never runs a handler and never touches an +-- an ack handle. Its correctness property is message conservation — every input +-- message appears in exactly one emitted batch, and per batch key arrival order +-- is preserved. +-- +-- The accumulation logic is a pure, deterministic core (see 'stepArrival', +-- 'stepTick', 'stepFlush') that takes the current time as an argument, so it is +-- property-tested with no threads or wall-clock. 'runBatcher' is a thin IO +-- wrapper that drives the core from a background consumer and a single timeout +-- ticker, buffering results in a bounded queue for backpressure. +module Shibuya.Internal.Runner.Batcher + ( -- * Pure accumulation core + Accum (..), + BatcherState (..), + emptyBatcherState, + ReadyBatch, + stepArrival, + stepTick, + stepFlush, + + -- * IO engine + runBatcher, + ) +where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.STM + ( TVar, + atomically, + newTVarIO, + readTVar, + readTVarIO, + retry, + writeTVar, + ) +import Control.Monad (when) +import Data.List.NonEmpty (NonEmpty) +import Data.List.NonEmpty qualified as NE +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Sequence (Seq (..), (><)) +import Data.Sequence qualified as Seq +import Data.Word (Word64) +import GHC.Clock (getMonotonicTimeNSec) +import Shibuya.Batch (BatchConfig (..), BatchInfo (..), BatchKey, BatchTrigger (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Core.Types (Envelope (..)) +import Shibuya.Prelude +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream +import UnliftIO (finally, throwIO) +import UnliftIO.Async (Async, async, cancel, waitCatch) + +-- | In-progress state for one batch key. +data Accum es msg = Accum + { -- | Messages accumulated so far, newest first (reversed arrival order). + messages :: ![Ingested es msg], + -- | How many messages are buffered (== length of 'messages'). Always >= 1. + count :: !Int, + -- | When the first message for this key arrived (drives the timeout). + firstArrivalAt :: !Word64, + -- | Partition of the first message, copied into the emitted batch info. + partition0 :: !(Maybe Text) + } + +-- | The engine's memory: one accumulator per active batch key. +newtype BatcherState es msg = BatcherState + { accums :: Map BatchKey (Accum es msg) + } + +-- | An empty state with no groups in progress. +emptyBatcherState :: BatcherState es msg +emptyBatcherState = BatcherState Map.empty + +-- | A finished group ready to hand downstream: metadata plus its messages in +-- arrival order. This is the exact value the execution stage (EP-18) consumes. +type ReadyBatch es msg = (BatchInfo, NonEmpty (Ingested es msg)) + +data DrainStep es msg + = DrainReady !(ReadyBatch es msg) + | DrainDone + +-- | Build a ReadyBatch from a completed accumulator. Reverses the buffered +-- messages back into arrival order; safe because count >= 1. +emitAccum :: BatchKey -> BatchTrigger -> Accum es msg -> ReadyBatch es msg +emitAccum key trig acc = + ( BatchInfo + { batchKey = key, + size = acc.count, + trigger = trig, + partition = acc.partition0 + }, + NE.fromList (reverse acc.messages) + ) + +-- | Fold one arriving message into the state. Emits a size-triggered batch iff +-- the message fills its key's accumulator to batchSize. +stepArrival :: + BatchConfig es msg -> + Word64 -> + Ingested es msg -> + BatcherState es msg -> + (BatcherState es msg, [ReadyBatch es msg]) +stepArrival cfg now ing (BatcherState accums) = + let key = cfg.batchKey ing.envelope + (ready, accums') = Map.alterF update key accums + in (BatcherState accums', ready) + where + update Nothing = + let acc = + Accum + { messages = [ing], + count = 1, + firstArrivalAt = now, + partition0 = ing.envelope.partition + } + in if cfg.batchSize <= 1 + then ([emitAccum (cfg.batchKey ing.envelope) TriggerSize acc], Nothing) + else ([], Just acc) + update (Just acc) = + let key = cfg.batchKey ing.envelope + acc' = acc {messages = ing : acc.messages, count = acc.count + 1} + in if acc'.count >= cfg.batchSize + then ([emitAccum key TriggerSize acc'], Nothing) + else ([], Just acc') + +-- | Emit every accumulator whose timeout has elapsed as of the supplied time. +stepTick :: + BatchConfig es msg -> + Word64 -> + BatcherState es msg -> + (BatcherState es msg, [ReadyBatch es msg]) +stepTick cfg now (BatcherState accums) = + let timedOut :: Accum es msg -> Bool + timedOut acc = now >= acc.firstArrivalAt && now - acc.firstArrivalAt >= nominalToNanos cfg.batchTimeout + (ripe, keep) = Map.partition timedOut accums + ready = [emitAccum k TriggerTimeout acc | (k, acc) <- Map.toList ripe] + in (BatcherState keep, ready) + +-- | Emit all remaining accumulators (end-of-input / drain). Leaves the state empty. +stepFlush :: + BatcherState es msg -> + (BatcherState es msg, [ReadyBatch es msg]) +stepFlush (BatcherState accums) = + (emptyBatcherState, [emitAccum k TriggerFlush acc | (k, acc) <- Map.toList accums]) + +-- | Run one pure step in the same STM transaction that appends emitted batches +-- to the pending output buffer. Admission waits while the pending buffer is at +-- capacity, so backpressure blocks without holding a separate mutex. One step +-- may append a burst larger than the remaining capacity; the bound is therefore +-- capacity plus one step's burst. +emitStep :: + TVar (BatcherState es msg, Seq (ReadyBatch es msg)) -> + Int -> + (BatcherState es msg -> (BatcherState es msg, [ReadyBatch es msg])) -> + IO () +emitStep stateVar capacity step = + atomically $ do + (state, pending) <- readTVar stateVar + when (Seq.length pending >= capacity) retry + let (state', ready) = step state + writeTVar stateVar (state', pending >< Seq.fromList ready) + +-- | Convert a 'NominalDiffTime' (seconds) into whole microseconds for +-- 'threadDelay', clamped to at least 1 so a misconfiguration cannot spin. +nominalToMicros :: NominalDiffTime -> Int +nominalToMicros d = max 1 (round (realToFrac d * 1e6 :: Double)) + +nominalToNanos :: NominalDiffTime -> Word64 +nominalToNanos d = round (realToFrac d * 1e9 :: Double) + +-- | Group a stream of individual messages into a stream of batches. +-- +-- @outputCapacity@ bounds how many finished batches may sit un-consumed (>= 1); +-- when full, the engine stalls, propagating backpressure upstream. @cfg@ carries +-- the size, timeout, key function, and tick interval. The @es@/@msg@ parameters +-- are phantom to the engine (it never runs an effect or acks a message). +runBatcher :: + Natural -> + BatchConfig es msg -> + Stream IO (Ingested es msg) -> + Stream IO (ReadyBatch es msg) +runBatcher outputCapacity cfg input = + Stream.bracketIO acquire release consume + where + tickMicros = nominalToMicros (fromMaybe cfg.batchTimeout cfg.tickInterval) + outputCapacityInt = max 1 (fromIntegral outputCapacity) + + acquire = do + stateVar <- newTVarIO (emptyBatcherState, Seq.empty) + doneVar <- newTVarIO False + + let onArrival ing = do + now <- getMonotonicTimeNSec + emitStep stateVar outputCapacityInt (stepArrival cfg now ing) + + -- Consume the whole input, flush the remainder, then mark done. + -- 'finally' guarantees doneVar is set even if the input stream throws, + -- so the output stream below always reaches the drain point. The drain + -- then observes this async with 'waitCatch' and rethrows any consumer + -- failure instead of reporting clean completion. + consumer = + ( do + Stream.fold Fold.drain (Stream.mapM onArrival input) + emitStep stateVar outputCapacityInt stepFlush + ) + `finally` atomically (writeTVar doneVar True) + + tickerLoop = do + threadDelay tickMicros + done <- readTVarIO doneVar + if done + then pure () + else do + now <- getMonotonicTimeNSec + emitStep stateVar outputCapacityInt (stepTick cfg now) + tickerLoop + + consumerA <- async consumer + tickerA <- async tickerLoop + pure (stateVar, doneVar, consumerA, tickerA) + + release (_stateVar, _doneVar, consumerA, tickerA) = do + cancel tickerA + cancel consumerA + + consume (stateVar, doneVar, consumerA, _tickerA) = drainQueue consumerA stateVar doneVar + +-- | Stream finished batches out of the bounded queue, ending when the consumer +-- has flushed everything (doneVar), the queue has drained, and the consumer +-- async is known to have completed successfully. +drainQueue :: + Async () -> + TVar (BatcherState es msg, Seq (ReadyBatch es msg)) -> + TVar Bool -> + Stream IO (ReadyBatch es msg) +drainQueue consumerA stateVar doneVar = Stream.unfoldrM step () + where + step _ = do + drainStep <- + atomically $ do + (state, pending) <- readTVar stateVar + case Seq.viewl pending of + rb Seq.:< rest -> do + writeTVar stateVar (state, rest) + pure (DrainReady rb) + Seq.EmptyL -> do + done <- readTVar doneVar + if done + then pure DrainDone + else retry + case drainStep of + DrainReady rb -> pure (Just (rb, ())) + DrainDone -> + waitCatch consumerA >>= \case + Left ex -> throwIO ex + Right () -> pure Nothing diff --git a/src/Shibuya/Internal/Runner/Finalize.hs b/src/Shibuya/Internal/Runner/Finalize.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Finalize.hs @@ -0,0 +1,50 @@ +-- | __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. +-- +-- Shared resilient finalization for framework-owned ack paths. +-- +-- Both single-message and batch processing call adapter-provided +-- 'AckHandle.finalize' through this module. A transient finalizer exception is +-- retried with the fixed schedule below; each retry uses the same already +-- resolved 'AckDecision'. +module Shibuya.Internal.Runner.Finalize + ( finalizeRetryDelaysMicros, + finalizeWithRetry, + ) +where + +import Effectful (Eff, IOE, liftIO, (:>)) +import OpenTelemetry.Trace.Core qualified as OTel +import Shibuya.Core.Ack (AckDecision) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested (..)) +import Shibuya.Telemetry.Effect (Tracing, recordException) +import UnliftIO (SomeException, catchAny) +import UnliftIO.Concurrent (threadDelay) + +-- | Retry delays, in microseconds, used after failed finalization attempts. +finalizeRetryDelaysMicros :: [Int] +finalizeRetryDelaysMicros = [10_000, 50_000, 250_000] + +-- | Call a message finalizer until it succeeds or the bounded retry budget is +-- exhausted. The resolved decision is never recomputed between attempts. +finalizeWithRetry :: + (IOE :> es, Tracing :> es) => + OTel.Span -> + Ingested es msg -> + AckDecision -> + Eff es (Either SomeException ()) +finalizeWithRetry traceSpan ingested decision = go finalizeRetryDelaysMicros + where + go delays = + catchAny + (Right <$> ingested.ack.finalize decision) + ( \ex -> do + recordException traceSpan ex + case delays of + [] -> pure (Left ex) + delay : rest -> do + liftIO $ threadDelay delay + go rest + ) diff --git a/src/Shibuya/Internal/Runner/Halt.hs b/src/Shibuya/Internal/Runner/Halt.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Halt.hs @@ -0,0 +1,23 @@ +-- | __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. +-- +-- Halt exception for processor termination. +-- Thrown when a handler returns AckHalt to stop processing. +module Shibuya.Internal.Runner.Halt + ( ProcessorHalt (..), + ) +where + +import Control.Exception (Exception) +import Shibuya.Core.Ack (HaltReason) +import Shibuya.Prelude + +-- | Exception thrown when processing should halt. +-- The supervisor catches this to handle graceful shutdown. +data ProcessorHalt = ProcessorHalt + { reason :: !HaltReason + } + deriving stock (Show, Generic) + +instance Exception ProcessorHalt diff --git a/src/Shibuya/Internal/Runner/Ingester.hs b/src/Shibuya/Internal/Runner/Ingester.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Ingester.hs @@ -0,0 +1,62 @@ +-- | __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. +-- +-- Ingester - reads from adapter stream and sends to inbox. +-- Provides backpressure via bounded inbox. +module Shibuya.Internal.Runner.Ingester + ( runIngester, + runIngesterWithMetrics, + ) +where + +import Control.Concurrent.NQE.Process (Inbox, inboxToMailbox, send) +import Effectful (Eff, IOE, liftIO, (:>)) +import Shibuya.Core.Ingested (Ingested) +import Shibuya.Core.Metrics (MetricsHandle, incrementReceived) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream + +-- | Run the ingester: reads from stream, sends to inbox. +-- Sends each element to the inbox with backpressure - if inbox is full, blocks. +-- Completes when the stream completes. +runIngester :: + (IOE :> es) => + -- | Source stream from adapter + Stream (Eff es) (Ingested es msg) -> + -- | Target inbox (bounded for backpressure) + Inbox (Ingested es msg) -> + Eff es () +runIngester source inbox = do + let mailbox = inboxToMailbox inbox + -- Send each stream element to the mailbox, then drain + -- This provides backpressure: if mailbox is full (bounded), send blocks + Stream.fold Fold.drain $ + Stream.mapM + (\msg -> send msg mailbox >> pure msg) + source + +-- | Run the ingester with metrics tracking. +-- Increments the received count for each message sent to inbox. +runIngesterWithMetrics :: + (IOE :> es) => + -- | Metrics handle (for updating received count) + MetricsHandle -> + -- | Source stream from adapter + Stream (Eff es) (Ingested es msg) -> + -- | Target inbox (bounded for backpressure) + Inbox (Ingested es msg) -> + Eff es () +runIngesterWithMetrics metricsHandle source inbox = do + let mailbox = inboxToMailbox inbox + Stream.fold Fold.drain $ + Stream.mapM + ( \msg -> do + -- Increment received count + liftIO $ incrementReceived metricsHandle + -- Send to inbox (blocks if full - backpressure) + liftIO $ send msg mailbox + pure msg + ) + source diff --git a/src/Shibuya/Internal/Runner/KeyedScheduler.hs b/src/Shibuya/Internal/Runner/KeyedScheduler.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/KeyedScheduler.hs @@ -0,0 +1,207 @@ +-- | __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. +-- +-- Bounded FIFO scheduler for work keyed by an optional partition key. +module Shibuya.Internal.Runner.KeyedScheduler + ( runKeyedScheduler, + ) +where + +import Control.Applicative ((<|>)) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Control.Concurrent.STM + ( STM, + TVar, + atomically, + modifyTVar', + newTVarIO, + readTVar, + readTVarIO, + retry, + writeTVar, + ) +import Data.Foldable (traverse_) +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Map.Strict qualified as Map +import Data.Sequence (Seq) +import Data.Sequence qualified as Seq +import Data.Set qualified as Set +import Data.Unique (Unique, newUnique) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream qualified as Stream +import UnliftIO (Async, SomeException, async, cancel, catchAny, finally, throwIO, withAsync) + +-- | Run every item in a stream through the worker action, at most +-- @maxConcurrency@ at a time. Items with the same @Just key@ run strictly in +-- stream order; @Nothing@ items have no mutual ordering constraint. The pending +-- buffer is bounded and blocks the reader when full, preserving upstream +-- backpressure. +runKeyedScheduler :: + (Ord key) => + Int -> + Int -> + (item -> Maybe key) -> + (item -> IO ()) -> + Stream.Stream IO item -> + IO () +runKeyedScheduler requestedConcurrency requestedPendingLimit itemKey itemAction itemStream = do + scheduler <- newTVarIO emptyKeyedSchedulerState + workers <- newTVarIO (Map.empty :: Map.Map Unique (Async ())) + let maxConcurrency = max 1 requestedConcurrency + pendingLimit = max 1 requestedPendingLimit + cancelWorkers = do + liveWorkers <- readTVarIO workers + traverse_ cancel (Map.elems liveWorkers) + + reader = + ( do + Stream.fold Fold.drain $ + Stream.mapM (enqueueItem pendingLimit scheduler) itemStream + atomically $ markInputDone scheduler Nothing + ) + `catchAny` \ex -> + atomically $ markInputDone scheduler (Just ex) + + loop = do + step <- atomically $ nextSchedulerStep maxConcurrency itemKey scheduler + case step of + SchedulerDone Nothing -> pure () + SchedulerDone (Just ex) -> throwIO ex + StartItem item -> do + workerId <- newUnique + startGate <- newEmptyMVar + worker <- async $ do + takeMVar startGate + runWorker scheduler workers workerId itemKey itemAction item + atomically $ modifyTVar' workers (Map.insert workerId worker) + putMVar startGate () + loop + + withAsync reader $ \_reader -> + loop `finally` cancelWorkers + +data KeyedSchedulerState key item = KeyedSchedulerState + { inputDone :: !Bool, + activeKeys :: !(Set.Set key), + running :: !Int, + pending :: !(Seq item), + firstFailure :: !(Maybe SomeException) + } + +data SchedulerStep item + = StartItem !item + | SchedulerDone !(Maybe SomeException) + +emptyKeyedSchedulerState :: KeyedSchedulerState key item +emptyKeyedSchedulerState = + KeyedSchedulerState + { inputDone = False, + activeKeys = Set.empty, + running = 0, + pending = Seq.empty, + firstFailure = Nothing + } + +runWorker :: + (Ord key) => + TVar (KeyedSchedulerState key item) -> + TVar (Map.Map Unique (Async ())) -> + Unique -> + (item -> Maybe key) -> + (item -> IO ()) -> + item -> + IO () +runWorker scheduler workers workerId itemKey itemAction item = do + resultRef <- newIORef Nothing + let runItem = + (itemAction item >> pure Nothing) + `catchAny` (pure . Just) + >>= writeIORef resultRef + cleanup = do + result <- readIORef resultRef + atomically $ finishItem scheduler itemKey item result + atomically $ modifyTVar' workers (Map.delete workerId) + runItem `finally` cleanup + +enqueueItem :: + Int -> + TVar (KeyedSchedulerState key item) -> + item -> + IO () +enqueueItem pendingLimit scheduler item = + atomically $ do + s <- readTVar scheduler + if Seq.length s.pending >= pendingLimit + then retry + else writeTVar scheduler s {pending = s.pending Seq.|> item} + +markInputDone :: + TVar (KeyedSchedulerState key item) -> + Maybe SomeException -> + STM () +markInputDone scheduler failure = + modifyTVar' scheduler $ \s -> + s + { inputDone = True, + firstFailure = s.firstFailure <|> failure + } + +nextSchedulerStep :: + (Ord key) => + Int -> + (item -> Maybe key) -> + TVar (KeyedSchedulerState key item) -> + STM (SchedulerStep item) +nextSchedulerStep maxConcurrency itemKey scheduler = do + s <- readTVar scheduler + case (s.running < maxConcurrency, popStartable itemKey s.activeKeys s.pending) of + (True, Just (item, rest)) -> do + writeTVar + scheduler + s + { activeKeys = maybe s.activeKeys (`Set.insert` s.activeKeys) (itemKey item), + running = s.running + 1, + pending = rest + } + pure (StartItem item) + _ + | s.inputDone && Seq.null s.pending && s.running == 0 -> + pure (SchedulerDone s.firstFailure) + | otherwise -> + retry + +finishItem :: + (Ord key) => + TVar (KeyedSchedulerState key item) -> + (item -> Maybe key) -> + item -> + Maybe SomeException -> + STM () +finishItem scheduler itemKey item failure = + modifyTVar' scheduler $ \s -> + s + { activeKeys = maybe s.activeKeys (`Set.delete` s.activeKeys) (itemKey item), + running = s.running - 1, + firstFailure = s.firstFailure <|> failure + } + +popStartable :: + (Ord key) => + (item -> Maybe key) -> + Set.Set key -> + Seq item -> + Maybe (item, Seq item) +popStartable itemKey active = go Seq.empty + where + go skipped items = + case Seq.viewl items of + Seq.EmptyL -> + Nothing + item Seq.:< rest -> + case itemKey item of + Just key + | key `Set.member` active -> + go (skipped Seq.|> item) rest + _ -> + Just (item, skipped <> rest) diff --git a/src/Shibuya/Internal/Runner/Master.hs b/src/Shibuya/Internal/Runner/Master.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Master.hs @@ -0,0 +1,190 @@ +-- | __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. +-- +-- Master process - central coordinator for queue processors. +-- Provides supervision, metrics collection, and control API. +-- +-- Architecture: +-- - Master is an NQE Process that handles control messages +-- - Holds a Supervisor for managing child processors +-- - Maintains TVar MetricsMap for O(1) metrics access +-- - Processors register their metrics TVars with the Master +module Shibuya.Internal.Runner.Master + ( -- * Master Handle + Master (..), + MasterState (..), + + -- * Control Messages + MasterMessage (..), + + -- * Starting the Master + startMaster, + stopMaster, + + -- * Introspection + getAllMetrics, + getAllMetricsIO, + getProcessorMetrics, + getProcessorMetricsIO, + + -- * Processor Management + registerProcessor, + unregisterProcessor, + ) +where + +import Control.Concurrent.NQE.Process + ( Inbox, + Listen, + Process (..), + newInbox, + receive, + ) +import Control.Concurrent.NQE.Supervisor (Strategy (..), Supervisor) +import Control.Concurrent.NQE.Supervisor qualified as Supervisor +import Control.Concurrent.STM + ( TVar, + atomically, + modifyTVar', + newTVarIO, + readTVar, + ) +import Control.Monad (forever) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Effectful (Eff, IOE, liftIO, (:>)) +import Shibuya.Core.Metrics + ( MetricsHandle, + MetricsMap, + ProcessorId, + ProcessorMetrics, + sampleMetrics, + ) +import Shibuya.Prelude +import UnliftIO (Async, async, cancel, link) + +-- | Messages for the master process. +data MasterMessage + = -- | Register a processor's metrics handle + RegisterProcessor !ProcessorId !MetricsHandle !(Listen ()) + | -- | Unregister a processor + UnregisterProcessor !ProcessorId !(Listen ()) + | -- | Shutdown all processors + Shutdown !(Listen ()) + +-- | Master state held in TVars. +data MasterState = MasterState + { -- | Map of processor IDs to their metrics handles + metrics :: !(TVar (Map ProcessorId MetricsHandle)), + -- | The supervisor managing child processors + supervisor :: !Supervisor, + -- | Whether child failures should be linked into the spawning thread. + -- Derived from the supervision strategy: True for KillAll/IgnoreGraceful + -- (failure must reach the application), False for IgnoreAll/Notify. + propagateFailures :: !Bool + } + deriving (Generic) + +-- | Master handle - provides access to the master process. +data Master = Master + { -- | The async handle for the master + handle :: !(Async ()), + -- | Direct access to master state + state :: !MasterState, + -- | Inbox for sending messages + inbox :: !(Inbox MasterMessage) + } + deriving (Generic) + +-- | Start the master process. +-- Returns a handle for interacting with the master. +-- The caller is responsible for calling stopMaster when done. +startMaster :: (IOE :> es) => Strategy -> Eff es Master +startMaster strategy = liftIO $ do + -- Create supervisor + sup <- Supervisor.supervisor strategy + + metricsMapVar <- newTVarIO Map.empty + let propagate = case strategy of + KillAll -> True + IgnoreGraceful -> True + IgnoreAll -> False + Notify _ -> False + masterState = MasterState metricsMapVar sup propagate + + masterInbox <- newInbox + + -- Start master loop + masterHandle <- async $ masterLoop masterState masterInbox + link masterHandle + + pure + Master + { handle = masterHandle, + state = masterState, + inbox = masterInbox + } + +-- | Stop the master and all child processors. +-- Cancels the supervisor first (which cancels all children via NQE's stopAll), +-- then cancels the master message loop. +stopMaster :: (IOE :> es) => Master -> Eff es () +stopMaster master = liftIO $ do + -- Cancel the supervisor first - this triggers NQE's stopAll which cancels all children + cancel (getProcessAsync master.state.supervisor) + -- Then cancel the master message loop + cancel master.handle + +-- | The master process main loop. +masterLoop :: MasterState -> Inbox MasterMessage -> IO () +masterLoop state inbox = forever $ do + msg <- receive inbox + handleMessage state msg + +-- | Handle a single master message. +handleMessage :: MasterState -> MasterMessage -> IO () +handleMessage state msg = case msg of + RegisterProcessor pid metricsHandle respond -> do + atomically $ do + modifyTVar' state.metrics $ Map.insert pid metricsHandle + atomically $ respond () + UnregisterProcessor pid respond -> do + atomically $ do + modifyTVar' state.metrics $ Map.delete pid + atomically $ respond () + Shutdown respond -> do + -- Clear all processors + atomically $ modifyTVar' state.metrics (const Map.empty) + atomically $ respond () + +-- | Get metrics for all processors. +getAllMetrics :: (IOE :> es) => Master -> Eff es MetricsMap +getAllMetrics = liftIO . getAllMetricsIO + +-- | Get metrics for all processors (IO version for web servers). +getAllMetricsIO :: Master -> IO MetricsMap +getAllMetricsIO master = do + handlesMap <- atomically $ readTVar master.state.metrics + traverse sampleMetrics handlesMap + +-- | Get metrics for a specific processor. +getProcessorMetrics :: (IOE :> es) => Master -> ProcessorId -> Eff es (Maybe ProcessorMetrics) +getProcessorMetrics master = liftIO . getProcessorMetricsIO master + +-- | Get metrics for a specific processor (IO version for web servers). +getProcessorMetricsIO :: Master -> ProcessorId -> IO (Maybe ProcessorMetrics) +getProcessorMetricsIO master pid = do + handlesMap <- atomically $ readTVar master.state.metrics + traverse sampleMetrics (Map.lookup pid handlesMap) + +-- | Register a processor with the master. +-- The processor should call this with its metrics handle. +registerProcessor :: (IOE :> es) => Master -> ProcessorId -> MetricsHandle -> Eff es () +registerProcessor master pid metricsHandle = + liftIO $ atomically $ modifyTVar' master.state.metrics $ Map.insert pid metricsHandle + +-- | Unregister a processor from the master. +unregisterProcessor :: (IOE :> es) => Master -> ProcessorId -> Eff es () +unregisterProcessor master pid = + liftIO $ atomically $ modifyTVar' master.state.metrics $ Map.delete pid diff --git a/src/Shibuya/Internal/Runner/Supervised.hs b/src/Shibuya/Internal/Runner/Supervised.hs new file mode 100644 --- /dev/null +++ b/src/Shibuya/Internal/Runner/Supervised.hs @@ -0,0 +1,683 @@ +-- | __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. +-- +-- Supervised runner - runs processors under NQE supervision with metrics. +-- This is the production runner with introspection and control. +-- +-- Architecture: +-- - Each processor runs as a child under the Master's supervisor +-- - Metrics are tracked in a TVar and registered with the Master +-- - Backpressure is provided via bounded inbox between stream and processor +module Shibuya.Internal.Runner.Supervised + ( -- * Running with Supervision + runSupervised, + runSupervisedBatch, + + -- * Standalone (without Master) + runWithMetrics, + runWithMetricsBatch, + + -- * Processor Handle + SupervisedProcessor (..), + + -- * Introspection + getMetrics, + getProcessorState, + isDone, + ) +where + +import Control.Concurrent.NQE.Process (Inbox, mailboxEmptySTM, newBoundedInbox, receiveSTM) +import Control.Concurrent.NQE.Supervisor (addChild) +import Control.Concurrent.STM + ( TVar, + atomically, + modifyTVar', + newTVarIO, + orElse, + readTVar, + readTVarIO, + retry, + writeTVar, + ) +import Control.Monad (when) +import Data.HashMap.Strict qualified as HashMap +import Data.IORef (IORef, atomicWriteIORef, newIORef, readIORef) +import Data.Text qualified as Text +import Effectful (Eff, IOE, Limit (..), Persistence (..), UnliftStrategy (..), liftIO, withEffToIO, (:>)) +import Effectful.Dispatch.Static (unsafeEff_) +import OpenTelemetry.Attributes (Attribute, toAttribute) +import OpenTelemetry.Trace.Core qualified as OTel +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Batch (BatchConfig, BatchHandler) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..), RetryDelay (..)) +import Shibuya.Core.Error (HandlerError (..), handlerErrorToText) +import Shibuya.Core.Ingested (Ingested (..), toMessage) +import Shibuya.Core.Metrics + ( AckDecisionMetric (..), + MetricsHandle (..), + ProcessorId (..), + ProcessorMetrics (..), + ProcessorState (..), + beginProcessing, + finishFinalizationFailure, + finishProcessing, + newMetricsHandle, + sampleMetrics, + ) +import Shibuya.Core.Metrics qualified as Metrics +import Shibuya.Core.Types (Envelope (..), MessageId (..)) +import Shibuya.Handler (Handler) +import Shibuya.Internal.Runner.BatchProcessor (processBatchesUntilDrained) +import Shibuya.Internal.Runner.Batcher (runBatcher) +import Shibuya.Internal.Runner.Finalize (finalizeWithRetry) +import Shibuya.Internal.Runner.Halt (ProcessorHalt (..)) +import Shibuya.Internal.Runner.Ingester (runIngesterWithMetrics) +import Shibuya.Internal.Runner.KeyedScheduler (runKeyedScheduler) +import Shibuya.Internal.Runner.Master (Master (..), MasterState (..), registerProcessor, unregisterProcessor) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +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, + attrMessagingMessageId, + attrMessagingOperation, + attrMessagingSystem, + attrShibuyaAckDecision, + attrShibuyaInflightCount, + attrShibuyaInflightMax, + attrShibuyaPartition, + consumerSpanArgs, + eventAckDecision, + eventHandlerCompleted, + eventHandlerStarted, + mkEvent, + processSpanName, + ) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream qualified as Stream +import Streamly.Data.Stream.Prelude qualified as StreamP +import UnliftIO (Async, catch, catchAny, displayException, finally, throwIO) +import UnliftIO qualified as UIO + +-- | Handle for a supervised processor. +-- Provides introspection into the running processor. +data SupervisedProcessor = SupervisedProcessor + { -- | Live metrics for this processor + metrics :: !MetricsHandle, + -- | The processor's ID + processorId :: !ProcessorId, + -- | Whether processing is complete + done :: !(TVar Bool), + -- | The async handle if running under supervision + child :: !(Maybe (Async ())) + } + +-- | Get current metrics for the processor. +getMetrics :: (IOE :> es) => SupervisedProcessor -> Eff es ProcessorMetrics +getMetrics sp = liftIO $ sampleMetrics sp.metrics + +-- | Get current state of the processor. +getProcessorState :: (IOE :> es) => SupervisedProcessor -> Eff es ProcessorState +getProcessorState sp = liftIO $ (.state) <$> sampleMetrics sp.metrics + +-- | Check if processing is done. +isDone :: (IOE :> es) => SupervisedProcessor -> Eff es Bool +isDone sp = liftIO $ readTVarIO sp.done + +-- | Run a processor under the Master's supervision with metrics tracking. +-- +-- Architecture: +-- 1. Creates a bounded inbox (using inboxSize for backpressure) +-- 2. Spawns ingester async (reads from adapter, sends to inbox) +-- 3. Runs processor loop (receives from inbox, calls handler) +-- 4. Registers metrics with Master, unregisters on completion +-- +-- Returns immediately with a handle for introspection. +runSupervised :: + (IOE :> es, Tracing :> es) => + Master -> + -- | Inbox size (for backpressure) + Natural -> + -- | Processor identifier + ProcessorId -> + -- | OrderingPolicy policy + OrderingPolicy -> + -- | Concurrency mode + Concurrency -> + -- | Queue adapter + Adapter es msg -> + -- | Message handler + Handler es msg -> + Eff es SupervisedProcessor +runSupervised master inboxSize procId ordering concurrency adapter handler = do + now <- liftIO getCurrentTime + + -- Initialize state + metricsHandle <- liftIO $ newMetricsHandle now + doneVar <- liftIO $ newTVarIO False + + -- Register with Master + registerProcessor master procId metricsHandle + + -- Add as supervised child using NQE's Supervisor + -- ConcUnlift Persistent allows the runInIO function to be used in the async child + supervisedChild <- withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> + addChild master.state.supervisor $ + runInIO + ( -- Catch ProcessorHalt to prevent propagation via link + -- (Halt is intentional, not a failure - other processors should continue) + ( runIngesterAndProcessor metricsHandle procId inboxSize ordering concurrency adapter handler + `catch` \(ProcessorHalt _) -> pure () -- Convert halt to graceful exit + ) + `finally` unregisterProcessor master procId + ) + `finally` atomically (writeTVar doneVar True) + + -- Link so exceptions propagate to the parent for strategies that request it. + when master.state.propagateFailures $ + unsafeEff_ $ + UIO.link supervisedChild + + pure + SupervisedProcessor + { metrics = metricsHandle, + processorId = procId, + done = doneVar, + child = Just supervisedChild + } + +-- | Run a processor with metrics but without Master supervision. +-- Useful for testing or simple single-processor setups. +-- This blocks until the stream is exhausted and all messages are processed +-- using serial processing with the same concurrent ingester/drainer shape as +-- the supervised runner. +runWithMetrics :: + (IOE :> es, Tracing :> es) => + -- | Inbox size (for backpressure) + Natural -> + -- | Processor identifier + ProcessorId -> + -- | Queue adapter + Adapter es msg -> + -- | Message handler + Handler es msg -> + Eff es SupervisedProcessor +runWithMetrics inboxSize procId adapter handler = do + now <- liftIO getCurrentTime + + -- Initialize state + metricsHandle <- liftIO $ newMetricsHandle now + doneVar <- liftIO $ newTVarIO False + + runIngesterAndProcessor metricsHandle procId inboxSize Unordered Serial adapter handler + `finally` liftIO (atomically (writeTVar doneVar True)) + + pure + SupervisedProcessor + { metrics = metricsHandle, + processorId = procId, + done = doneVar, + child = Nothing + } + +-- | Run ingester and processor with a bounded inbox. +-- Ingester reads from adapter stream, processor calls handler. +-- When stream exhausts, processor drains remaining messages and exits. +runIngesterAndProcessor :: + (IOE :> es, Tracing :> es) => + MetricsHandle -> + ProcessorId -> + Natural -> + OrderingPolicy -> + Concurrency -> + Adapter es msg -> + Handler es msg -> + Eff es () +runIngesterAndProcessor metricsHandle procId inboxSize ordering concurrency adapter handler = do + -- Create bounded inbox (this is where inboxSize is used for backpressure) + inbox <- liftIO $ newBoundedInbox inboxSize + + -- Signal when ingester completes (stream exhausted) + streamDoneVar <- liftIO $ newTVarIO False + + -- Run ingester async, processor in main thread + withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do + -- Ingester: run until stream exhausts, then signal done + -- Use finally to ensure streamDoneVar is always set, even if ingester fails + let ingesterWithSignal = + runInIO (runIngesterWithMetrics metricsHandle adapter.source inbox) + `finally` atomically (writeTVar streamDoneVar True) + + UIO.withAsync ingesterWithSignal $ \ingesterAsync -> do + -- Processor: process messages, exit when stream done and inbox empty + runInIO $ processUntilDrained metricsHandle procId ordering concurrency handler inbox streamDoneVar + UIO.waitCatch ingesterAsync >>= \case + Left ingesterErr -> do + now <- getCurrentTime + atomically $ + modifyTVar' metricsHandle.cold $ \m -> + m {Metrics.state = Failed (Text.pack (displayException ingesterErr)) now} + UIO.throwIO ingesterErr + Right () -> pure () + +-- | Run a batching processor under the Master's supervision with metrics. +-- +-- Identical in shape to 'runSupervised' but the inner loop accumulates messages +-- into batches (via the 'BatchConfig') and runs a 'BatchHandler' over each batch, +-- finalizing every message exactly once. On a batch-handler halt the child exits +-- gracefully (the processor halt is caught here), matching 'runSupervised'. +runSupervisedBatch :: + (IOE :> es, Tracing :> es) => + Master -> + -- | Inbox size (for backpressure) + Natural -> + -- | Processor identifier + ProcessorId -> + -- | Concurrency mode (bounds how many BATCHES run at once) + Concurrency -> + -- | Batch configuration + BatchConfig es msg -> + -- | Queue adapter + Adapter es msg -> + -- | Batch handler + BatchHandler es msg -> + Eff es SupervisedProcessor +runSupervisedBatch master inboxSize procId concurrency batchConfig adapter batchHandler = do + now <- liftIO getCurrentTime + + metricsHandle <- liftIO $ newMetricsHandle now + doneVar <- liftIO $ newTVarIO False + + registerProcessor master procId metricsHandle + + supervisedChild <- withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> + addChild master.state.supervisor $ + runInIO + ( ( runIngesterAndProcessorBatch + metricsHandle + procId + inboxSize + concurrency + batchConfig + adapter + batchHandler + `catch` \(ProcessorHalt _) -> pure () + ) + `finally` unregisterProcessor master procId + ) + `finally` atomically (writeTVar doneVar True) + + when master.state.propagateFailures $ + unsafeEff_ $ + UIO.link supervisedChild + + pure + SupervisedProcessor + { metrics = metricsHandle, + processorId = procId, + done = doneVar, + child = Just supervisedChild + } + +-- | Run a batching processor with metrics but without Master supervision. +-- Blocks until the adapter stream is exhausted and every accumulated batch has +-- been processed (including the end-of-input flush). Useful for tests. +runWithMetricsBatch :: + (IOE :> es, Tracing :> es) => + Natural -> + ProcessorId -> + Concurrency -> + BatchConfig es msg -> + Adapter es msg -> + BatchHandler es msg -> + Eff es SupervisedProcessor +runWithMetricsBatch inboxSize procId concurrency batchConfig adapter batchHandler = do + now <- liftIO getCurrentTime + + metricsHandle <- liftIO $ newMetricsHandle now + doneVar <- liftIO $ newTVarIO False + + runIngesterAndProcessorBatch + metricsHandle + procId + inboxSize + concurrency + batchConfig + adapter + batchHandler + `finally` liftIO (atomically (writeTVar doneVar True)) + + pure + SupervisedProcessor + { metrics = metricsHandle, + processorId = procId, + done = doneVar, + child = Nothing + } + +-- | Run ingester and batch processor with a bounded inbox. +-- The ingester reads from the adapter stream into the inbox exactly as in the +-- single-message path; 'inboxToStream' turns the inbox into a halt-aware, +-- stream-done-aware Stream; 'runBatcher' groups that into ready batches; and +-- 'processBatchesUntilDrained' runs the batch handler and finalizes each message +-- exactly once. When the adapter stream ends (including on graceful shutdown, +-- when 'Adapter.shutdown' ends 'source'), the ingester completes, sets +-- streamDoneVar, 'inboxToStream' terminates once the inbox is empty, the batcher +-- reaches end-of-input and flushes all pending partial batches with TriggerFlush, +-- and only then does 'processBatchesUntilDrained' return. The spawn sites +-- ('runSupervised', 'runSupervisedBatch', and 'runWithMetricsBatch') own +-- marking the processor done via 'finally'. +-- +-- 'processBatchesUntilDrained' (EP-18) only /sets/ the halt reference on a +-- batch-handler @AckHalt@ or exhausted finalization; it does not throw. So we +-- read the halt reference after it returns and throw a processor halt here, +-- mirroring the single-message +-- 'processUntilDrained'. +runIngesterAndProcessorBatch :: + (IOE :> es, Tracing :> es) => + MetricsHandle -> + ProcessorId -> + Natural -> + Concurrency -> + BatchConfig es msg -> + Adapter es msg -> + BatchHandler es msg -> + Eff es () +runIngesterAndProcessorBatch metricsHandle procId inboxSize concurrency batchConfig adapter batchHandler = do + inbox <- liftIO $ newBoundedInbox inboxSize + streamDoneVar <- liftIO $ newTVarIO False + haltRef <- liftIO $ newIORef Nothing + + withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do + let ingesterWithSignal = + runInIO (runIngesterWithMetrics metricsHandle adapter.source inbox) + `finally` atomically (writeTVar streamDoneVar True) + + UIO.withAsync ingesterWithSignal $ \ingesterAsync -> do + let inboxStream = inboxToStream inbox streamDoneVar haltRef + readyBatchStream = runBatcher inboxSize batchConfig inboxStream + batchProcessor = + runInIO $ do + processBatchesUntilDrained + metricsHandle + procId + concurrency + batchHandler + readyBatchStream + haltRef + maybeHalt <- liftIO (readIORef haltRef) + maybe (pure ()) (throwIO . ProcessorHalt) maybeHalt + batchProcessor `catchAny` \processorErr -> do + now <- getCurrentTime + atomically $ + modifyTVar' metricsHandle.cold $ \m -> + m {Metrics.state = Failed (Text.pack (displayException processorErr)) now} + UIO.throwIO processorErr + UIO.waitCatch ingesterAsync >>= \case + Left ingesterErr -> do + now <- getCurrentTime + atomically $ + modifyTVar' metricsHandle.cold $ \m -> + m {Metrics.state = Failed (Text.pack (displayException ingesterErr)) now} + UIO.throwIO ingesterErr + Right () -> pure () + +-- | Convert inbox to a stream for use with streamly. +-- Respects both the stream-done signal and halt flag. +-- +-- Uses STM to atomically check done/empty and receive, avoiding a race +-- condition where the processor could block on receive after the stream +-- has completed but before the done flag was checked. +inboxToStream :: + Inbox (Ingested es msg) -> + TVar Bool -> + IORef (Maybe HaltReason) -> + Stream.Stream IO (Ingested es msg) +inboxToStream inbox streamDoneVar haltRef = Stream.unfoldrM step () + where + step _ = do + -- Check halt flag first (outside STM since it's an IORef) + halted <- readIORef haltRef + case halted of + Just _ -> pure Nothing -- Stop reading + Nothing -> do + -- Atomically either receive a message or detect completion. + -- This avoids TOCTOU race where we check done/empty separately + -- and then block on receive after the stream has completed. + result <- + atomically $ + -- Try to receive a message + (Just <$> receiveSTM inbox) + `orElse` + -- Or check if we're done (stream exhausted and inbox empty) + ( do + done <- readTVar streamDoneVar + empty <- mailboxEmptySTM inbox + if done && empty + then pure Nothing + else retry -- Inbox empty but stream not done, wait for message + ) + pure $ fmap (,()) result + +-- | Process messages from inbox until stream is done and inbox is empty. +-- Supports Serial, Ahead, and Async concurrency modes. +processUntilDrained :: + (IOE :> es, Tracing :> es) => + MetricsHandle -> + ProcessorId -> + OrderingPolicy -> + Concurrency -> + Handler es msg -> + Inbox (Ingested es msg) -> + TVar Bool -> + Eff es () +processUntilDrained metricsHandle procId ordering concurrency handler inbox streamDoneVar = do + haltRef <- liftIO $ newIORef Nothing + + let maxConc = case concurrency of + Serial -> 1 + Ahead n -> n + Async n -> n + ProcessorId pidText = procId + spanName = processSpanName pidText + constantFrameworkAttrs = + HashMap.fromList + [ (attrMessagingSystem, toAttribute ("shibuya" :: Text)), + (attrMessagingDestinationName, toAttribute pidText), + (attrMessagingOperation, toAttribute ("process" :: Text)) + ] + + withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do + let inboxStream = inboxToStream inbox streamDoneVar haltRef + processAction = runInIO . processOne metricsHandle spanName constantFrameworkAttrs maxConc haltRef handler + partitioned n = + runKeyedScheduler + (max 1 n) + (max 2 (2 * max 1 n)) + (\ingested -> ingested.envelope.partition) + processAction + inboxStream + + case (ordering, concurrency) of + (_, Serial) -> + Stream.fold Fold.drain $ + Stream.mapM processAction inboxStream + (PartitionedInOrder, Ahead n) -> + partitioned n + (PartitionedInOrder, Async n) -> + partitioned n + (_, Ahead n) -> + Stream.fold Fold.drain $ + StreamP.parMapM (StreamP.maxThreads n . StreamP.maxBuffer n . StreamP.ordered True) processAction inboxStream + (_, Async n) -> + Stream.fold Fold.drain $ + StreamP.parMapM (StreamP.maxThreads n . StreamP.maxBuffer n) processAction inboxStream + + -- After draining, check if we halted + maybeHalt <- readIORef haltRef + case maybeHalt of + Just reason -> throwIO $ ProcessorHalt reason + Nothing -> pure () + +handlerStartedEvent :: OTel.NewEvent +handlerStartedEvent = mkEvent eventHandlerStarted [] +{-# NOINLINE handlerStartedEvent #-} + +-- | Process a single message with metrics tracking and tracing. +-- Thread-safe for concurrent execution. +processOne :: + (IOE :> es, Tracing :> es) => + MetricsHandle -> + Text -> + HashMap.HashMap Text Attribute -> + Int -> + IORef (Maybe HaltReason) -> + Handler es msg -> + Ingested es msg -> + Eff es () +processOne metricsHandle spanName constantFrameworkAttrs maxConc haltRef handler ingested = do + -- Extract parent context from message headers for distributed tracing + let parentCtx = ingested.envelope.traceContext >>= extractTraceContext + + withExtractedContext parentCtx $ + withSpan' spanName consumerSpanArgs $ \traceSpan -> do + -- Build the messaging.* attribute set: framework defaults first, + -- then the adapter's HashMap layered on top so adapter keys with + -- the same name override the framework's default. The explicit + -- 'HashMap.union' here is left-biased (left wins), which keeps + -- the precedence rule local and obvious instead of relying on + -- the order of repeated 'addAttribute' / 'addAttributes' calls + -- against the underlying mutable Span. + let MessageId msgIdText = ingested.envelope.messageId + frameworkAttrs = + HashMap.insert attrMessagingMessageId (toAttribute msgIdText) $ + case ingested.envelope.partition of + Just p -> HashMap.insert attrShibuyaPartition (toAttribute p) constantFrameworkAttrs + Nothing -> constantFrameworkAttrs + mergedAttrs = HashMap.union ingested.envelope.attributes frameworkAttrs + addAttributes traceSpan mergedAttrs + + -- Increment in-flight and add inflight attributes. + currentInflight <- liftIO $ beginProcessing metricsHandle maxConc + + addAttribute traceSpan attrShibuyaInflightCount currentInflight + addAttribute traceSpan attrShibuyaInflightMax maxConc + + -- Record handler start event + addEvent traceSpan handlerStartedEvent + + -- Call handler and finalizer separately. A handler exception is + -- substituted with immediate retry so the adapter always observes a + -- finalization decision for an ingested message. + handlerResult <- + catchAny + (Right <$> handler (toMessage ingested)) + ( \ex -> do + recordException traceSpan ex + pure (Left ex) + ) + let (decision, result) = case handlerResult of + Right d -> (d, Right d) + Left ex -> + ( AckRetry (RetryDelay 0), + Left $ HandlerException $ Text.pack $ show ex + ) + + when (isLeft handlerResult) $ + addEvent traceSpan $ + mkEvent + eventAckDecision + [(attrShibuyaAckDecision, OTel.toAttribute ("ack_retry" :: Text))] + + finalizeResult <- finalizeWithRetry traceSpan ingested decision + + -- Record completion event and set status + case finalizeResult of + Left _ -> do + addEvent traceSpan $ + mkEvent + eventAckDecision + [(attrShibuyaAckDecision, OTel.toAttribute ("finalization_failed" :: Text))] + setStatus traceSpan $ OTel.Error $ finalizationFailureText msgIdText + Right () -> case result of + Right decision' -> do + let decisionText = showAckDecision decision' + addEvent traceSpan $ + mkEvent + eventHandlerCompleted + [(attrShibuyaAckDecision, OTel.toAttribute decisionText)] + addAttribute traceSpan attrShibuyaAckDecision decisionText + + -- Set span status based on decision + case decision' of + AckOk -> setStatus traceSpan OTel.Ok + AckRetry _ -> setStatus traceSpan OTel.Ok + AckDeadLetter reason -> + setStatus traceSpan $ OTel.Error $ showDeadLetterReason reason + AckHalt reason -> + setStatus traceSpan $ OTel.Error $ showHaltReason reason + Left err -> do + addEvent traceSpan $ + mkEvent + eventAckDecision + [(attrShibuyaAckDecision, OTel.toAttribute ("error" :: Text))] + setStatus traceSpan $ OTel.Error $ handlerErrorToText err + + -- Decrement in-flight and update stats. + liftIO $ + case finalizeResult of + Left _ -> finishFinalizationFailure metricsHandle (finalizationFailureText msgIdText) + Right () -> finishProcessing metricsHandle (metricForResult result) + + -- Handle halt (set flag, don't throw - let stream drain) + case finalizeResult of + Left _ -> + liftIO $ atomicWriteIORef haltRef (Just (HaltFatal (finalizationFailureText msgIdText))) + Right () -> case result of + Right (AckHalt reason) -> liftIO $ atomicWriteIORef haltRef (Just reason) + _ -> pure () + where + isLeft :: Either a b -> Bool + isLeft (Left _) = True + isLeft (Right _) = False + + finalizationFailureText :: Text -> Text + finalizationFailureText msgIdText = + "finalization failed for message id: " <> msgIdText + + showAckDecision :: AckDecision -> Text + showAckDecision AckOk = "ack_ok" + showAckDecision (AckRetry _) = "ack_retry" + showAckDecision (AckDeadLetter _) = "ack_dead_letter" + showAckDecision (AckHalt _) = "ack_halt" + + showDeadLetterReason :: DeadLetterReason -> Text + showDeadLetterReason (PoisonPill t) = "poison_pill: " <> t + showDeadLetterReason (InvalidPayload t) = "invalid_payload: " <> t + showDeadLetterReason MaxRetriesExceeded = "max_retries_exceeded" + + showHaltReason :: HaltReason -> Text + showHaltReason (HaltOrderedStream t) = "halt_ordered_stream: " <> t + showHaltReason (HaltFatal t) = "halt_fatal: " <> t + +metricForResult :: Either HandlerError AckDecision -> Either Text AckDecisionMetric +metricForResult (Left err) = Left (handlerErrorToText err) +metricForResult (Right AckOk) = Right CountProcessed +metricForResult (Right (AckRetry _)) = Right CountProcessed +metricForResult (Right (AckDeadLetter _)) = Right CountFailed +metricForResult (Right (AckHalt reason)) = Right (CountHalt (haltReasonText reason)) + where + haltReasonText (HaltOrderedStream t) = t + haltReasonText (HaltFatal t) = t diff --git a/src/Shibuya/Policy.hs b/src/Shibuya/Policy.hs --- a/src/Shibuya/Policy.hs +++ b/src/Shibuya/Policy.hs @@ -1,8 +1,8 @@ --- | Ordering and concurrency policies. +-- | OrderingPolicy and concurrency policies. -- Runner policy that maps ordering guarantees to concurrency constraints. module Shibuya.Policy - ( -- * Ordering - Ordering (..), + ( -- * OrderingPolicy + OrderingPolicy (..), -- * Concurrency Concurrency (..), @@ -14,13 +14,14 @@ import Shibuya.Core.Error (PolicyError (..)) import Shibuya.Prelude -import Prelude hiding (Ordering) -- | Message ordering guarantees. -data Ordering +data OrderingPolicy = -- | Event-sourced subscriptions - must be Serial StrictInOrder - | -- | Kafka-style - parallel across partitions + | -- | Kafka-style ordering: messages with the same partition key are + -- processed and acknowledged in arrival order, while distinct partitions + -- may run concurrently. Messages without a partition key are unconstrained. PartitionedInOrder | -- | No ordering guarantees Unordered @@ -30,7 +31,11 @@ data Concurrency = -- | One message at a time Serial - | -- | Prefetch N, process in order + | -- | Process up to N messages concurrently. Stream results are yielded + -- downstream in input order, but handler execution and acknowledgement run + -- concurrently and may complete in any order. Since Shibuya discards the + -- per-message result, this ordered yielding is not observable as ordered + -- side effects or ordered acks. Ahead !Int | -- | Process N concurrently Async !Int @@ -38,7 +43,7 @@ -- | Validate policy combinations. -- Invariant: StrictInOrder => Serial -validatePolicy :: Ordering -> Concurrency -> Either PolicyError () +validatePolicy :: OrderingPolicy -> Concurrency -> Either PolicyError () validatePolicy StrictInOrder (Ahead _) = Left $ InvalidPolicyCombo "StrictInOrder requires Serial concurrency" validatePolicy StrictInOrder (Async _) = Left $ InvalidPolicyCombo "StrictInOrder requires Serial concurrency" validatePolicy _ _ = Right () diff --git a/src/Shibuya/Prelude.hs b/src/Shibuya/Prelude.hs --- a/src/Shibuya/Prelude.hs +++ b/src/Shibuya/Prelude.hs @@ -13,18 +13,11 @@ LocalTime, NominalDiffTime, getCurrentTime, - -- lens - module Control.Lens, - -- vector - Vector, ) where import "base" Control.Monad.IO.Class (MonadIO) import "base" GHC.Generics (Generic) import "base" Numeric.Natural (Natural) -import "generic-lens" Data.Generics.Labels () -import "lens" Control.Lens import "text" Data.Text (Text) import "time" Data.Time (Day, LocalTime, NominalDiffTime, UTCTime, getCurrentTime) -import "vector" Data.Vector (Vector) diff --git a/src/Shibuya/Runner/Halt.hs b/src/Shibuya/Runner/Halt.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Halt.hs +++ /dev/null @@ -1,19 +0,0 @@ --- | Halt exception for processor termination. --- Thrown when a handler returns AckHalt to stop processing. -module Shibuya.Runner.Halt - ( ProcessorHalt (..), - ) -where - -import Control.Exception (Exception) -import Shibuya.Core.Ack (HaltReason) -import Shibuya.Prelude - --- | Exception thrown when processing should halt. --- The supervisor catches this to handle graceful shutdown. -data ProcessorHalt = ProcessorHalt - { reason :: !HaltReason - } - deriving stock (Show, Generic) - -instance Exception ProcessorHalt diff --git a/src/Shibuya/Runner/Ingester.hs b/src/Shibuya/Runner/Ingester.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Ingester.hs +++ /dev/null @@ -1,60 +0,0 @@ --- | Ingester - reads from adapter stream and sends to inbox. --- Provides backpressure via bounded inbox. -module Shibuya.Runner.Ingester - ( runIngester, - runIngesterWithMetrics, - ) -where - -import Control.Concurrent.NQE.Process (Inbox, inboxToMailbox, send) -import Control.Concurrent.STM (TVar, atomically, modifyTVar') -import Effectful (Eff, IOE, liftIO, (:>)) -import Shibuya.Core.Ingested (Ingested) -import Shibuya.Runner.Metrics (ProcessorMetrics (..), incReceived) -import Streamly.Data.Fold qualified as Fold -import Streamly.Data.Stream (Stream) -import Streamly.Data.Stream qualified as Stream - --- | Run the ingester: reads from stream, sends to inbox. --- Sends each element to the inbox with backpressure - if inbox is full, blocks. --- Completes when the stream completes. -runIngester :: - (IOE :> es) => - -- | Source stream from adapter - Stream (Eff es) (Ingested es msg) -> - -- | Target inbox (bounded for backpressure) - Inbox (Ingested es msg) -> - Eff es () -runIngester source inbox = do - let mailbox = inboxToMailbox inbox - -- Send each stream element to the mailbox, then drain - -- This provides backpressure: if mailbox is full (bounded), send blocks - Stream.fold Fold.drain $ - Stream.mapM - (\msg -> send msg mailbox >> pure msg) - source - --- | Run the ingester with metrics tracking. --- Increments 'received' count for each message sent to inbox. -runIngesterWithMetrics :: - (IOE :> es) => - -- | Metrics TVar (for updating received count) - TVar ProcessorMetrics -> - -- | Source stream from adapter - Stream (Eff es) (Ingested es msg) -> - -- | Target inbox (bounded for backpressure) - Inbox (Ingested es msg) -> - Eff es () -runIngesterWithMetrics metricsVar source inbox = do - let mailbox = inboxToMailbox inbox - Stream.fold Fold.drain $ - Stream.mapM - ( \msg -> do - -- Increment received count - liftIO $ atomically $ modifyTVar' metricsVar $ \m -> - m {stats = incReceived m.stats} - -- Send to inbox (blocks if full - backpressure) - liftIO $ send msg mailbox - pure msg - ) - source diff --git a/src/Shibuya/Runner/Master.hs b/src/Shibuya/Runner/Master.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Master.hs +++ /dev/null @@ -1,195 +0,0 @@ --- | Master process - central coordinator for queue processors. --- Provides supervision, metrics collection, and control API. --- --- Architecture: --- - Master is an NQE Process that handles control messages --- - Holds a Supervisor for managing child processors --- - Maintains TVar MetricsMap for O(1) metrics access --- - Processors register their metrics TVars with the Master -module Shibuya.Runner.Master - ( -- * Master Handle - Master (..), - MasterState (..), - - -- * Control Messages - MasterMessage (..), - - -- * Starting the Master - startMaster, - stopMaster, - - -- * Introspection - getAllMetrics, - getAllMetricsIO, - getProcessorMetrics, - getProcessorMetricsIO, - - -- * Processor Management - registerProcessor, - unregisterProcessor, - ) -where - -import Control.Concurrent.NQE.Process - ( Inbox, - Listen, - Process (..), - newInbox, - query, - receive, - ) -import Control.Concurrent.NQE.Supervisor (Strategy (..), Supervisor) -import Control.Concurrent.NQE.Supervisor qualified as Supervisor -import Control.Concurrent.STM - ( TVar, - atomically, - modifyTVar', - newTVarIO, - readTVar, - ) -import Control.Monad (forever) -import Data.Map.Strict (Map) -import Data.Map.Strict qualified as Map -import Effectful (Eff, IOE, liftIO, (:>)) -import Shibuya.Prelude -import Shibuya.Runner.Metrics - ( MetricsMap, - ProcessorId, - ProcessorMetrics, - ) -import UnliftIO (Async, async, cancel, link) - --- | Messages for the master process. -data MasterMessage - = -- | Get metrics for all processors - GetAllMetrics !(Listen MetricsMap) - | -- | Get metrics for a specific processor - GetProcessorMetrics ProcessorId !(Listen (Maybe ProcessorMetrics)) - | -- | Register a processor's metrics TVar - RegisterProcessor !ProcessorId !(TVar ProcessorMetrics) !(Listen ()) - | -- | Unregister a processor - UnregisterProcessor !ProcessorId !(Listen ()) - | -- | Shutdown all processors - Shutdown !(Listen ()) - --- | Master state held in TVars. -data MasterState = MasterState - { -- | Map of processor IDs to their metrics TVars - metrics :: !(TVar (Map ProcessorId (TVar ProcessorMetrics))), - -- | The supervisor managing child processors - supervisor :: !Supervisor - } - deriving (Generic) - --- | Master handle - provides access to the master process. -data Master = Master - { -- | The async handle for the master - handle :: !(Async ()), - -- | Direct access to master state - state :: !MasterState, - -- | Inbox for sending messages - inbox :: !(Inbox MasterMessage) - } - deriving (Generic) - --- | Start the master process. --- Returns a handle for interacting with the master. --- The caller is responsible for calling stopMaster when done. -startMaster :: (IOE :> es) => Strategy -> Eff es Master -startMaster strategy = liftIO $ do - -- Create supervisor - sup <- Supervisor.supervisor strategy - - metricsMapVar <- newTVarIO Map.empty - let masterState = MasterState metricsMapVar sup - - masterInbox <- newInbox - - -- Start master loop - masterHandle <- async $ masterLoop masterState masterInbox - link masterHandle - - pure - Master - { handle = masterHandle, - state = masterState, - inbox = masterInbox - } - --- | Stop the master and all child processors. --- Cancels the supervisor first (which cancels all children via NQE's stopAll), --- then cancels the master message loop. -stopMaster :: (IOE :> es) => Master -> Eff es () -stopMaster master = liftIO $ do - -- Cancel the supervisor first - this triggers NQE's stopAll which cancels all children - cancel (getProcessAsync master.state.supervisor) - -- Then cancel the master message loop - cancel (master ^. #handle) - --- | The master process main loop. -masterLoop :: MasterState -> Inbox MasterMessage -> IO () -masterLoop state inbox = forever $ do - msg <- receive inbox - handleMessage state msg - --- | Handle a single master message. -handleMessage :: MasterState -> MasterMessage -> IO () -handleMessage state msg = case msg of - GetAllMetrics respond -> do - -- Read all metrics from their TVars - metricsMap <- atomically $ do - tvarsMap <- readTVar state.metrics - traverse readTVar tvarsMap - atomically $ respond metricsMap - GetProcessorMetrics pid respond -> do - result <- atomically $ do - tvarsMap <- readTVar state.metrics - case Map.lookup pid tvarsMap of - Nothing -> pure Nothing - Just tvar -> Just <$> readTVar tvar - atomically $ respond result - RegisterProcessor pid metricsTVar respond -> do - atomically $ do - modifyTVar' state.metrics $ Map.insert pid metricsTVar - atomically $ respond () - UnregisterProcessor pid respond -> do - atomically $ do - modifyTVar' state.metrics $ Map.delete pid - atomically $ respond () - Shutdown respond -> do - -- Clear all processors - atomically $ modifyTVar' state.metrics (const Map.empty) - atomically $ respond () - --- | Get metrics for all processors. -getAllMetrics :: (IOE :> es) => Master -> Eff es MetricsMap -getAllMetrics master = - liftIO $ - GetAllMetrics `query` master.inbox - --- | Get metrics for all processors (IO version for web servers). -getAllMetricsIO :: Master -> IO MetricsMap -getAllMetricsIO master = GetAllMetrics `query` master.inbox - --- | Get metrics for a specific processor. -getProcessorMetrics :: (IOE :> es) => Master -> ProcessorId -> Eff es (Maybe ProcessorMetrics) -getProcessorMetrics master pid = - liftIO $ - GetProcessorMetrics pid `query` master.inbox - --- | Get metrics for a specific processor (IO version for web servers). -getProcessorMetricsIO :: Master -> ProcessorId -> IO (Maybe ProcessorMetrics) -getProcessorMetricsIO master pid = GetProcessorMetrics pid `query` master.inbox - --- | Register a processor with the master. --- The processor should call this with its metrics TVar. -registerProcessor :: (IOE :> es) => Master -> ProcessorId -> TVar ProcessorMetrics -> Eff es () -registerProcessor master pid metricsTVar = - liftIO $ - RegisterProcessor pid metricsTVar `query` master.inbox - --- | Unregister a processor from the master. -unregisterProcessor :: (IOE :> es) => Master -> ProcessorId -> Eff es () -unregisterProcessor master pid = - liftIO $ - UnregisterProcessor pid `query` master.inbox diff --git a/src/Shibuya/Runner/Metrics.hs b/src/Shibuya/Runner/Metrics.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Metrics.hs +++ /dev/null @@ -1,167 +0,0 @@ --- | Metrics and state tracking for processors. --- Provides introspection into what's happening in the system. -module Shibuya.Runner.Metrics - ( -- * Processor State - ProcessorState (..), - ProcessorId (..), - - -- * In-Flight Tracking - InFlightInfo (..), - emptyInFlightInfo, - - -- * Stream Statistics - StreamStats (..), - emptyStreamStats, - - -- * Combined Metrics - ProcessorMetrics (..), - emptyProcessorMetrics, - - -- * Metrics Map - MetricsMap, - - -- * Metrics Updates - incReceived, - incDropped, - incProcessed, - incFailed, - ) -where - -import Data.Aeson (FromJSON (..), FromJSONKey (..), ToJSON (..), ToJSONKey (..), object, withObject, (.:)) -import Data.Aeson qualified as Aeson -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 dropped (backpressure) - dropped :: !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 0 - --- | Combined processor metrics. -data ProcessorMetrics = ProcessorMetrics - { -- | Current state - state :: !ProcessorState, - -- | Statistics - stats :: !StreamStats, - -- | 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, - startedAt = now - } - --- | Map of processor IDs to their metrics. -type MetricsMap = Map ProcessorId ProcessorMetrics - --- | Increment received count. -incReceived :: StreamStats -> StreamStats -incReceived = #received %~ (+ 1) - --- | Increment dropped count. -incDropped :: StreamStats -> StreamStats -incDropped = #dropped %~ (+ 1) - --- | Increment processed count. -incProcessed :: StreamStats -> StreamStats -incProcessed = #processed %~ (+ 1) - --- | Increment failed count. -incFailed :: StreamStats -> StreamStats -incFailed = #failed %~ (+ 1) diff --git a/src/Shibuya/Runner/Processor.hs b/src/Shibuya/Runner/Processor.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Processor.hs +++ /dev/null @@ -1,78 +0,0 @@ --- | Processor - receives from inbox, calls handler, finalizes ack. --- Runs in a loop until terminated. -module Shibuya.Runner.Processor - ( runProcessor, - runProcessorN, - drainInbox, - ) -where - -import Control.Concurrent.NQE.Process (Inbox, mailboxEmpty, receive) -import Control.Monad (forever, unless) -import Effectful (Eff, IOE, (:>)) -import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Handler (Handler) - --- | Run the processor forever: receive message, call handler, finalize ack. --- This runs until an exception is thrown or the thread is cancelled. -runProcessor :: - (IOE :> es) => - -- | Message handler - Handler es msg -> - -- | Source inbox - Inbox (Ingested es msg) -> - Eff es () -runProcessor handler inbox = forever $ processOne handler inbox - --- | Run the processor for exactly N messages. --- Useful for testing. -runProcessorN :: - (IOE :> es) => - -- | Number of messages to process - Int -> - -- | Message handler - Handler es msg -> - -- | Source inbox - Inbox (Ingested es msg) -> - Eff es () -runProcessorN n handler inbox = go n - where - go 0 = pure () - go remaining = do - processOne handler inbox - go (remaining - 1) - --- | Process a single message: receive, handle, finalize. -processOne :: - (IOE :> es) => - Handler es msg -> - Inbox (Ingested es msg) -> - Eff es () -processOne handler inbox = do - -- Receive next message (blocks if inbox empty) - ingested <- receive inbox - - -- Call the handler to get the ack decision - decision <- handler ingested - - -- Finalize the ack (commit, retry, dead-letter, or halt) - -- Note: AckHalt is handled by the supervised runner (Supervised.hs throws ProcessorHalt) - ingested.ack.finalize decision - --- | Drain all remaining messages from inbox until empty. --- Used after ingester completes to process buffered messages. -drainInbox :: - (IOE :> es) => - -- | Message handler - Handler es msg -> - -- | Source inbox - Inbox (Ingested es msg) -> - Eff es () -drainInbox handler inbox = go - where - go = do - empty <- mailboxEmpty inbox - unless empty $ do - processOne handler inbox - go diff --git a/src/Shibuya/Runner/Serial.hs b/src/Shibuya/Runner/Serial.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Serial.hs +++ /dev/null @@ -1,44 +0,0 @@ --- | Serial Runner - processes messages one at a time. --- Simple sequential runner for ordered processing. -module Shibuya.Runner.Serial - ( runSerial, - ) -where - -import Effectful (Eff, IOE, (:>)) -import Numeric.Natural (Natural) -import Shibuya.Adapter (Adapter (..)) -import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Handler (Handler) -import Streamly.Data.Fold qualified as Fold -import Streamly.Data.Stream qualified as Stream - --- | Run the serial runner. --- Processes each message from the adapter stream sequentially: --- 1. Pull message from stream --- 2. Call handler --- 3. Finalize ack --- 4. Repeat until stream is exhausted --- --- This is the simplest runner - no concurrency, no backpressure buffering. --- Suitable for ordered, low-throughput workloads. -runSerial :: - (IOE :> es) => - -- | Inbox size (unused in serial mode, for API compatibility) - Natural -> - -- | Queue adapter - Adapter es msg -> - -- | Message handler - Handler es msg -> - Eff es () -runSerial _inboxSize adapter handler = do - -- Process each message from the stream sequentially - Stream.fold Fold.drain $ - Stream.mapM processOne adapter.source - where - processOne ingested = do - -- Call the handler to get the ack decision - decision <- handler ingested - -- Finalize the ack - ingested.ack.finalize decision diff --git a/src/Shibuya/Runner/Supervised.hs b/src/Shibuya/Runner/Supervised.hs deleted file mode 100644 --- a/src/Shibuya/Runner/Supervised.hs +++ /dev/null @@ -1,516 +0,0 @@ --- | Supervised runner - runs processors under NQE supervision with metrics. --- This is the production runner with introspection and control. --- --- Architecture: --- - Each processor runs as a child under the Master's supervisor --- - Metrics are tracked in a TVar and registered with the Master --- - Backpressure is provided via bounded inbox between stream and processor -module Shibuya.Runner.Supervised - ( -- * Running with Supervision - runSupervised, - - -- * Standalone (without Master) - runWithMetrics, - - -- * Processor Handle - SupervisedProcessor (..), - - -- * Introspection - getMetrics, - getProcessorState, - isDone, - ) -where - -import Control.Concurrent.NQE.Process (Inbox, mailboxEmpty, mailboxEmptySTM, newBoundedInbox, receive, receiveSTM) -import Control.Concurrent.NQE.Supervisor (addChild) -import Control.Concurrent.STM - ( TVar, - atomically, - modifyTVar', - newTVarIO, - orElse, - readTVar, - readTVarIO, - retry, - writeTVar, - ) -import Control.Monad (unless) -import Data.HashMap.Strict qualified as HashMap -import Data.IORef (IORef, atomicWriteIORef, newIORef, readIORef) -import Data.Text qualified as Text -import Effectful (Eff, IOE, liftIO, withEffToIO, (:>)) -import Effectful.Dispatch.Static (unsafeEff_) -import Effectful.Internal.Unlift (Limit (..), Persistence (..), UnliftStrategy (..)) -import OpenTelemetry.Attributes (toAttribute) -import OpenTelemetry.Trace.Core qualified as OTel -import Shibuya.Adapter (Adapter (..)) -import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), HaltReason (..)) -import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Error (HandlerError (..), handlerErrorToText) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Core.Types (Envelope (..), MessageId (..)) -import Shibuya.Handler (Handler) -import Shibuya.Policy (Concurrency (..)) -import Shibuya.Prelude -import Shibuya.Runner.Halt (ProcessorHalt (..)) -import Shibuya.Runner.Ingester (runIngesterWithMetrics) -import Shibuya.Runner.Master (Master (..), MasterState (..), registerProcessor, unregisterProcessor) -import Shibuya.Runner.Metrics - ( InFlightInfo (..), - ProcessorId (..), - ProcessorMetrics (..), - ProcessorState (..), - emptyProcessorMetrics, - incFailed, - incProcessed, - ) -import Shibuya.Telemetry.Effect - ( Tracing, - addAttribute, - addAttributes, - addEvent, - recordException, - setStatus, - withExtractedContext, - withSpan', - ) -import Shibuya.Telemetry.Propagation (extractTraceContext) -import Shibuya.Telemetry.Semantic - ( attrMessagingDestinationName, - attrMessagingMessageId, - attrMessagingOperation, - attrMessagingSystem, - attrShibuyaAckDecision, - attrShibuyaInflightCount, - attrShibuyaInflightMax, - attrShibuyaPartition, - consumerSpanArgs, - eventAckDecision, - eventHandlerCompleted, - eventHandlerStarted, - mkEvent, - processSpanName, - ) -import Streamly.Data.Fold qualified as Fold -import Streamly.Data.Stream qualified as Stream -import Streamly.Data.Stream.Prelude qualified as StreamP -import UnliftIO (Async, catch, catchAny, displayException, finally, throwIO) -import UnliftIO qualified as UIO - --- | Handle for a supervised processor. --- Provides introspection into the running processor. -data SupervisedProcessor = SupervisedProcessor - { -- | Live metrics for this processor - metrics :: !(TVar ProcessorMetrics), - -- | The processor's ID - processorId :: !ProcessorId, - -- | Whether processing is complete - done :: !(TVar Bool), - -- | The async handle if running under supervision - child :: !(Maybe (Async ())) - } - --- | Get current metrics for the processor. -getMetrics :: (IOE :> es) => SupervisedProcessor -> Eff es ProcessorMetrics -getMetrics sp = liftIO $ readTVarIO sp.metrics - --- | Get current state of the processor. -getProcessorState :: (IOE :> es) => SupervisedProcessor -> Eff es ProcessorState -getProcessorState sp = liftIO $ (.state) <$> readTVarIO sp.metrics - --- | Check if processing is done. -isDone :: (IOE :> es) => SupervisedProcessor -> Eff es Bool -isDone sp = liftIO $ readTVarIO sp.done - --- | Run a processor under the Master's supervision with metrics tracking. --- --- Architecture: --- 1. Creates a bounded inbox (using inboxSize for backpressure) --- 2. Spawns ingester async (reads from adapter, sends to inbox) --- 3. Runs processor loop (receives from inbox, calls handler) --- 4. Registers metrics with Master, unregisters on completion --- --- Returns immediately with a handle for introspection. -runSupervised :: - (IOE :> es, Tracing :> es) => - Master -> - -- | Inbox size (for backpressure) - Natural -> - -- | Processor identifier - ProcessorId -> - -- | Concurrency mode - Concurrency -> - -- | Queue adapter - Adapter es msg -> - -- | Message handler - Handler es msg -> - Eff es SupervisedProcessor -runSupervised master inboxSize procId concurrency adapter handler = do - now <- liftIO getCurrentTime - - -- Initialize state - let initialMetrics = emptyProcessorMetrics now - metricsVar <- liftIO $ newTVarIO initialMetrics - doneVar <- liftIO $ newTVarIO False - - -- Register with Master - registerProcessor master procId metricsVar - - -- Add as supervised child using NQE's Supervisor - -- ConcUnlift Persistent allows the runInIO function to be used in the async child - supervisedChild <- withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> - addChild master.state.supervisor $ - runInIO $ - -- Catch ProcessorHalt to prevent propagation via link - -- (Halt is intentional, not a failure - other processors should continue) - ( runIngesterAndProcessor metricsVar procId doneVar inboxSize concurrency adapter handler - `catch` \(ProcessorHalt _) -> pure () -- Convert halt to graceful exit - ) - `finally` unregisterProcessor master procId - - -- Link so exceptions propagate to the parent - unsafeEff_ $ UIO.link supervisedChild - - pure - SupervisedProcessor - { metrics = metricsVar, - processorId = procId, - done = doneVar, - child = Just supervisedChild - } - --- | Run a processor with metrics but without Master supervision. --- Useful for testing or simple single-processor setups. --- This blocks until the stream is exhausted and all messages are processed. -runWithMetrics :: - (IOE :> es, Tracing :> es) => - -- | Inbox size (for backpressure) - Natural -> - -- | Processor identifier - ProcessorId -> - -- | Queue adapter - Adapter es msg -> - -- | Message handler - Handler es msg -> - Eff es SupervisedProcessor -runWithMetrics inboxSize procId adapter handler = do - now <- liftIO getCurrentTime - - -- Initialize state - let initialMetrics = emptyProcessorMetrics now - metricsVar <- liftIO $ newTVarIO initialMetrics - doneVar <- liftIO $ newTVarIO False - - -- Create bounded inbox - inbox <- liftIO $ newBoundedInbox inboxSize - - -- Run ingester to completion (all messages sent to inbox) - runIngesterWithMetrics metricsVar adapter.source inbox - - -- Drain remaining messages from inbox - drainInboxWithMetrics metricsVar procId handler inbox - - -- Mark done - liftIO $ atomically $ writeTVar doneVar True - - pure - SupervisedProcessor - { metrics = metricsVar, - processorId = procId, - done = doneVar, - child = Nothing - } - --- | Run ingester and processor with a bounded inbox. --- Ingester reads from adapter stream, processor calls handler. --- When stream exhausts, processor drains remaining messages and exits. -runIngesterAndProcessor :: - (IOE :> es, Tracing :> es) => - TVar ProcessorMetrics -> - ProcessorId -> - TVar Bool -> - Natural -> - Concurrency -> - Adapter es msg -> - Handler es msg -> - Eff es () -runIngesterAndProcessor metricsVar procId doneVar inboxSize concurrency adapter handler = do - -- Create bounded inbox (this is where inboxSize is used for backpressure) - inbox <- liftIO $ newBoundedInbox inboxSize - - -- Signal when ingester completes (stream exhausted) - streamDoneVar <- liftIO $ newTVarIO False - - -- Run ingester async, processor in main thread - withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do - -- Ingester: run until stream exhausts, then signal done - -- Use finally to ensure streamDoneVar is always set, even if ingester fails - let ingesterWithSignal = - runInIO (runIngesterWithMetrics metricsVar adapter.source inbox) - `finally` atomically (writeTVar streamDoneVar True) - - UIO.withAsync ingesterWithSignal $ \ingesterAsync -> do - -- Processor: process messages, exit when stream done and inbox empty - runInIO $ processUntilDrained metricsVar procId concurrency handler inbox streamDoneVar - UIO.poll ingesterAsync >>= \case - Just (Left ingesterErr) -> do - now <- getCurrentTime - atomically $ - modifyTVar' metricsVar $ \m -> - m & #state .~ Failed (Text.pack (displayException ingesterErr)) now - atomically $ writeTVar doneVar True - UIO.throwIO ingesterErr - _ -> pure () - - -- Mark done when processor exits - liftIO $ atomically $ writeTVar doneVar True - --- | Convert inbox to a stream for use with streamly. --- Respects both the stream-done signal and halt flag. --- --- Uses STM to atomically check done/empty and receive, avoiding a race --- condition where the processor could block on receive after the stream --- has completed but before the done flag was checked. -inboxToStream :: - Inbox (Ingested es msg) -> - TVar Bool -> - IORef (Maybe HaltReason) -> - Stream.Stream IO (Ingested es msg) -inboxToStream inbox streamDoneVar haltRef = Stream.unfoldrM step () - where - step _ = do - -- Check halt flag first (outside STM since it's an IORef) - halted <- readIORef haltRef - case halted of - Just _ -> pure Nothing -- Stop reading - Nothing -> do - -- Atomically either receive a message or detect completion. - -- This avoids TOCTOU race where we check done/empty separately - -- and then block on receive after the stream has completed. - result <- - atomically $ - -- Try to receive a message - (Just <$> receiveSTM inbox) - `orElse` - -- Or check if we're done (stream exhausted and inbox empty) - ( do - done <- readTVar streamDoneVar - empty <- mailboxEmptySTM inbox - if done && empty - then pure Nothing - else retry -- Inbox empty but stream not done, wait for message - ) - pure $ fmap (,()) result - --- | Process messages from inbox until stream is done and inbox is empty. --- Supports Serial, Ahead, and Async concurrency modes. -processUntilDrained :: - (IOE :> es, Tracing :> es) => - TVar ProcessorMetrics -> - ProcessorId -> - Concurrency -> - Handler es msg -> - Inbox (Ingested es msg) -> - TVar Bool -> - Eff es () -processUntilDrained metricsVar procId concurrency handler inbox streamDoneVar = do - haltRef <- liftIO $ newIORef Nothing - - let maxConc = case concurrency of - Serial -> 1 - Ahead n -> n - Async n -> n - - withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> do - let inboxStream = inboxToStream inbox streamDoneVar haltRef - processAction = runInIO . processOne metricsVar procId maxConc haltRef handler - - case concurrency of - Serial -> - Stream.fold Fold.drain $ - Stream.mapM processAction inboxStream - Ahead n -> - Stream.fold Fold.drain $ - StreamP.parMapM (StreamP.maxBuffer n . StreamP.ordered True) processAction inboxStream - Async n -> - Stream.fold Fold.drain $ - StreamP.parMapM (StreamP.maxBuffer n) processAction inboxStream - - -- After draining, check if we halted - maybeHalt <- readIORef haltRef - case maybeHalt of - Just reason -> throwIO $ ProcessorHalt reason - Nothing -> pure () - --- | Drain inbox until empty, with metrics tracking. --- Used for testing with finite streams. -drainInboxWithMetrics :: - (IOE :> es, Tracing :> es) => - TVar ProcessorMetrics -> - ProcessorId -> - Handler es msg -> - Inbox (Ingested es msg) -> - Eff es () -drainInboxWithMetrics metricsVar procId handler inbox = do - haltRef <- liftIO $ newIORef Nothing - go haltRef - where - go haltRef = do - empty <- liftIO $ mailboxEmpty inbox - unless empty $ do - ingested <- liftIO $ receive inbox - processOne metricsVar procId 1 haltRef handler ingested - -- Check if halted - halted <- liftIO $ readIORef haltRef - case halted of - Just reason -> throwIO $ ProcessorHalt reason - Nothing -> go haltRef - --- | Process a single message with metrics tracking and tracing. --- Thread-safe for concurrent execution. -processOne :: - (IOE :> es, Tracing :> es) => - TVar ProcessorMetrics -> - ProcessorId -> - Int -> - IORef (Maybe HaltReason) -> - Handler es msg -> - Ingested es msg -> - Eff es () -processOne metricsVar procId maxConc haltRef handler ingested = do - -- Extract parent context from message headers for distributed tracing - let parentCtx = ingested.envelope.traceContext >>= extractTraceContext - ProcessorId pidText = procId - - withExtractedContext parentCtx $ - withSpan' (processSpanName pidText) consumerSpanArgs $ \traceSpan -> do - -- Build the messaging.* attribute set: framework defaults first, - -- then the adapter's HashMap layered on top so adapter keys with - -- the same name override the framework's default. The explicit - -- 'HashMap.union' here is left-biased (left wins), which keeps - -- the precedence rule local and obvious instead of relying on - -- the order of repeated 'addAttribute' / 'addAttributes' calls - -- against the underlying mutable Span. - let MessageId msgIdText = ingested.envelope.messageId - frameworkAttrs = - HashMap.fromList $ - [ (attrMessagingSystem, toAttribute ("shibuya" :: Text)), - (attrMessagingDestinationName, toAttribute pidText), - (attrMessagingOperation, toAttribute ("process" :: Text)), - (attrMessagingMessageId, toAttribute msgIdText) - ] - <> case ingested.envelope.partition of - Just p -> [(attrShibuyaPartition, toAttribute p)] - Nothing -> [] - mergedAttrs = HashMap.union ingested.envelope.attributes frameworkAttrs - addAttributes traceSpan mergedAttrs - - -- Increment in-flight and add inflight attributes - now <- liftIO getCurrentTime - currentInflight <- liftIO $ atomically $ do - modifyTVar' metricsVar $ \m -> - let current = case m.state of - Processing info _ -> info.inFlight - _ -> 0 - in m & #state .~ Processing (InFlightInfo (current + 1) maxConc) now - m <- readTVar metricsVar - pure $ case m.state of - Processing info _ -> info.inFlight - _ -> 1 - - addAttribute traceSpan attrShibuyaInflightCount currentInflight - addAttribute traceSpan attrShibuyaInflightMax maxConc - - -- Record handler start event - addEvent traceSpan (mkEvent eventHandlerStarted []) - - -- Call handler and finalize - result <- - catchAny - ( do - decision <- handler ingested - ingested.ack.finalize decision - pure (Right decision) - ) - ( \ex -> do - recordException traceSpan ex - pure $ Left $ HandlerException $ Text.pack $ show ex - ) - - -- Record completion event and set status - case result of - Right decision -> do - let decisionText = showAckDecision decision - addEvent traceSpan $ - mkEvent - eventHandlerCompleted - [(attrShibuyaAckDecision, OTel.toAttribute decisionText)] - addAttribute traceSpan attrShibuyaAckDecision decisionText - - -- Set span status based on decision - case decision of - AckOk -> setStatus traceSpan OTel.Ok - AckRetry _ -> setStatus traceSpan OTel.Ok - AckDeadLetter reason -> - setStatus traceSpan $ OTel.Error $ showDeadLetterReason reason - AckHalt reason -> - setStatus traceSpan $ OTel.Error $ showHaltReason reason - Left err -> do - addEvent traceSpan $ - mkEvent - eventAckDecision - [(attrShibuyaAckDecision, OTel.toAttribute ("error" :: Text))] - setStatus traceSpan $ OTel.Error $ handlerErrorToText err - - -- Decrement in-flight and update stats - now' <- liftIO getCurrentTime - liftIO $ atomically $ modifyTVar' metricsVar $ decrementAndUpdate result now' - - -- Handle halt (set flag, don't throw - let stream drain) - case result of - Right (AckHalt reason) -> liftIO $ atomicWriteIORef haltRef (Just reason) - _ -> pure () - where - showAckDecision :: AckDecision -> Text - showAckDecision AckOk = "ack_ok" - showAckDecision (AckRetry _) = "ack_retry" - showAckDecision (AckDeadLetter _) = "ack_dead_letter" - showAckDecision (AckHalt _) = "ack_halt" - - showDeadLetterReason :: DeadLetterReason -> Text - showDeadLetterReason (PoisonPill t) = "poison_pill: " <> t - showDeadLetterReason (InvalidPayload t) = "invalid_payload: " <> t - showDeadLetterReason MaxRetriesExceeded = "max_retries_exceeded" - - showHaltReason :: HaltReason -> Text - showHaltReason (HaltOrderedStream t) = "halt_ordered_stream: " <> t - showHaltReason (HaltFatal t) = "halt_fatal: " <> t - --- | Decrement in-flight count and update stats based on result. -decrementAndUpdate :: - Either HandlerError AckDecision -> - UTCTime -> - ProcessorMetrics -> - ProcessorMetrics -decrementAndUpdate result now m = - let newState = case m.state of - Processing info _ -> - if info.inFlight <= 1 - then Idle - else Processing (info {inFlight = info.inFlight - 1}) now - other -> other - newStats = case result of - Right AckOk -> incProcessed m.stats - Right (AckRetry _) -> incProcessed m.stats - Right (AckDeadLetter _) -> incFailed m.stats - Right (AckHalt _) -> m.stats -- Mark as failed with halt message - Left _ -> incFailed m.stats - finalState = case result of - Right (AckHalt reason) -> Failed (haltReasonText reason) now - Left err -> Failed (handlerErrorToText err) now - _ -> newState - in m {state = finalState, stats = newStats} - where - haltReasonText (HaltOrderedStream t) = t - haltReasonText (HaltFatal t) = t diff --git a/src/Shibuya/Stream.hs b/src/Shibuya/Stream.hs --- a/src/Shibuya/Stream.hs +++ b/src/Shibuya/Stream.hs @@ -5,7 +5,7 @@ pollingStream, -- * Transformations - batchStream, + chunksOf, filterStream, ) where @@ -33,9 +33,10 @@ Nothing -> liftIO (threadDelay interval) >> pure Nothing Just msg -> pure (Just msg) --- | Batch messages into chunks. -batchStream :: (Monad m) => Int -> Stream m msg -> Stream m [msg] -batchStream n = Stream.foldMany (Fold.take n Fold.toList) +-- | Group a stream into lists of at most n elements. +-- Unrelated to the batch-processing API in "Shibuya.Batch". +chunksOf :: (Monad m) => Int -> Stream m msg -> Stream m [msg] +chunksOf n = Stream.foldMany (Fold.take n Fold.toList) -- | Filter messages. filterStream :: (Monad m) => (msg -> Bool) -> Stream m msg -> Stream m msg diff --git a/src/Shibuya/Telemetry.hs b/src/Shibuya/Telemetry.hs --- a/src/Shibuya/Telemetry.hs +++ b/src/Shibuya/Telemetry.hs @@ -4,28 +4,26 @@ -- -- == Quick Start -- --- 1. Enable tracing in your config: +-- 1. Initialize a tracer and run your app under 'runTracing' +-- (or 'runTracingNoop' to disable tracing): -- -- @ --- config = defaultAppConfig --- { tracing = defaultTracingConfig { enabled = True } --- } +-- provider <- OTel.initializeTracerProvider +-- let tracer = OTel.makeTracer provider instrumentationLibrary OTel.tracerOptions +-- runEff $ runTracing tracer app -- @ -- --- 2. Set environment variables for the OTLP exporter: +-- 2. Configure the OTLP exporter with environment variables: -- -- @ -- export OTEL_SERVICE_NAME="my-service" -- export OTEL_EXPORTER_OTLP_ENDPOINT="http://localhost:4317" -- @ -- --- 3. Spans are automatically created for message processing. --- See "Shibuya.Telemetry.Effect" for custom instrumentation. +-- 3. Spans are created automatically per message; see +-- "Shibuya.Telemetry.Effect" for custom instrumentation. module Shibuya.Telemetry - ( -- * Configuration - module Shibuya.Telemetry.Config, - - -- * Effect + ( -- * Effect module Shibuya.Telemetry.Effect, -- * Context Propagation @@ -36,7 +34,6 @@ ) where -import Shibuya.Telemetry.Config import Shibuya.Telemetry.Effect import Shibuya.Telemetry.Propagation import Shibuya.Telemetry.Semantic diff --git a/src/Shibuya/Telemetry/Config.hs b/src/Shibuya/Telemetry/Config.hs deleted file mode 100644 --- a/src/Shibuya/Telemetry/Config.hs +++ /dev/null @@ -1,33 +0,0 @@ --- | Tracing configuration for Shibuya. --- Most configuration is handled via OTEL_* environment variables --- by the hs-opentelemetry SDK. This module provides the minimal --- application-level config. -module Shibuya.Telemetry.Config - ( TracingConfig (..), - defaultTracingConfig, - ) -where - -import Data.Text (Text) -import GHC.Generics (Generic) - --- | Tracing configuration. --- The SDK reads most settings from OTEL_* environment variables. -data TracingConfig = TracingConfig - { -- | Master switch for tracing. When False, runTracingNoop is used. - enabled :: !Bool, - -- | Service name (fallback if OTEL_SERVICE_NAME not set) - serviceName :: !Text, - -- | Service version for span attributes - serviceVersion :: !(Maybe Text) - } - deriving stock (Eq, Show, Generic) - --- | Default tracing config with tracing disabled. -defaultTracingConfig :: TracingConfig -defaultTracingConfig = - TracingConfig - { enabled = False, - serviceName = "shibuya", - serviceVersion = Nothing - } diff --git a/src/Shibuya/Telemetry/Effect.hs b/src/Shibuya/Telemetry/Effect.hs --- a/src/Shibuya/Telemetry/Effect.hs +++ b/src/Shibuya/Telemetry/Effect.hs @@ -34,7 +34,7 @@ OTel.NewEvent (..), OTel.Tracer, OTel.defaultSpanArguments, - OTel.toAttribute, + toAttribute, ) where @@ -43,14 +43,13 @@ import Data.HashMap.Strict (HashMap) import Data.HashMap.Strict qualified as HashMap import Data.Text (Text) -import Effectful (Dispatch (..), DispatchOf, Eff, Effect, IOE, liftIO, withEffToIO, (:>)) +import Effectful (Dispatch (..), DispatchOf, Eff, Effect, IOE, Limit (..), Persistence (..), UnliftStrategy (..), liftIO, withEffToIO, (:>)) import Effectful.Dispatch.Static ( SideEffects (..), StaticRep, evalStaticRep, getStaticRep, ) -import Effectful.Internal.Unlift (Limit (..), Persistence (..), UnliftStrategy (..)) import OpenTelemetry.Attributes (Attribute, ToAttribute) import OpenTelemetry.Attributes qualified as Attr import OpenTelemetry.Context qualified as Ctx @@ -95,7 +94,7 @@ Eff es a runTracing t = evalStaticRep (TracingRep t True) --- | Run with tracing disabled (zero overhead). +-- | Run with tracing disabled (near-zero overhead). -- All tracing operations become no-ops. runTracingNoop :: (IOE :> es) => @@ -160,10 +159,7 @@ withSpan' name args f = do TracingRep {tracer, tracingEnabled} <- getStaticRep if not tracingEnabled - then do - -- Create a dropped span for the callback - dummySpan <- liftIO mkDummySpan - f dummySpan + then f dummySpan else withEffToIO (ConcUnlift Persistent Unlimited) $ \runInIO -> OTel.inSpan' tracer name args $ \traceSpan -> runInIO (f traceSpan) @@ -234,6 +230,10 @@ then pure () else liftIO $ OTel.setStatus traceSpan status +-- | Convert a value to an OpenTelemetry attribute. +toAttribute :: (ToAttribute a) => a -> Attribute +toAttribute = OTel.toAttribute + -------------------------------------------------------------------------------- -- Context Operations -------------------------------------------------------------------------------- @@ -289,24 +289,23 @@ -- Internal Helpers -------------------------------------------------------------------------------- --- | Create a dummy/dropped span for use when tracing is disabled. +-- | Shared non-recording span used when tracing is disabled. -- All operations on this span are no-ops. -mkDummySpan :: IO OTel.Span -mkDummySpan = do - -- Create a FrozenSpan with invalid (all-zeros) IDs - -- This creates a "Dropped" span that no-ops all operations - let dummyTraceId = case OTel.Id.bytesToTraceId (BS.replicate 16 0) of - Right tid -> tid - Left _ -> error "mkDummySpan: failed to create dummy TraceId" - dummySpanId = case OTel.Id.bytesToSpanId (BS.replicate 8 0) of - Right sid -> sid - Left _ -> error "mkDummySpan: failed to create dummy SpanId" - dummySpanContext = - OTel.SpanContext - { OTel.traceFlags = OTel.defaultTraceFlags, - OTel.isRemote = False, - OTel.traceId = dummyTraceId, - OTel.spanId = dummySpanId, - OTel.traceState = OTel.TraceState.empty - } - pure $ OTel.wrapSpanContext dummySpanContext +dummySpan :: OTel.Span +dummySpan = OTel.wrapSpanContext dummySpanContext + where + dummyTraceId = case OTel.Id.bytesToTraceId (BS.replicate 16 0) of + Right tid -> tid + Left _ -> error "dummySpan: failed to create dummy TraceId" + dummySpanId = case OTel.Id.bytesToSpanId (BS.replicate 8 0) of + Right sid -> sid + Left _ -> error "dummySpan: failed to create dummy SpanId" + dummySpanContext = + OTel.SpanContext + { OTel.traceFlags = OTel.defaultTraceFlags, + OTel.isRemote = False, + OTel.traceId = dummyTraceId, + OTel.spanId = dummySpanId, + OTel.traceState = OTel.TraceState.empty + } +{-# NOINLINE dummySpan #-} diff --git a/src/Shibuya/Telemetry/Propagation.hs b/src/Shibuya/Telemetry/Propagation.hs --- a/src/Shibuya/Telemetry/Propagation.hs +++ b/src/Shibuya/Telemetry/Propagation.hs @@ -75,7 +75,7 @@ -- and the failing-consumer's trace shows up linked to the resulting -- message in the downstream trace store. The lower-level -- 'injectTraceContext' is still exported for callers that already --- hold a 'Span' handle from inside a 'withSpan''. +-- hold a span handle from inside a @withSpan'@ callback. currentTraceHeaders :: (Tracing :> es, IOE :> es) => Eff es (Maybe TraceHeaders) diff --git a/src/Shibuya/Telemetry/Semantic.hs b/src/Shibuya/Telemetry/Semantic.hs --- a/src/Shibuya/Telemetry/Semantic.hs +++ b/src/Shibuya/Telemetry/Semantic.hs @@ -22,11 +22,16 @@ attrShibuyaInflightMax, attrShibuyaAckDecision, attrShibuyaPartition, + attrShibuyaBatchKey, + attrShibuyaBatchSize, + attrShibuyaBatchTrigger, -- * Event Names eventHandlerStarted, eventHandlerCompleted, eventAckDecision, + eventBatchStarted, + eventBatchCompleted, -- * SpanArguments Helpers consumerSpanArgs, @@ -139,6 +144,18 @@ attrShibuyaPartition :: Text attrShibuyaPartition = "shibuya.partition" +-- | The batch grouping key (@shibuya.batch.key@). +attrShibuyaBatchKey :: Text +attrShibuyaBatchKey = "shibuya.batch.key" + +-- | The number of messages in the batch (@shibuya.batch.size@). +attrShibuyaBatchSize :: Text +attrShibuyaBatchSize = "shibuya.batch.size" + +-- | Why the batch was emitted: size, timeout, or flush (@shibuya.batch.trigger@). +attrShibuyaBatchTrigger :: Text +attrShibuyaBatchTrigger = "shibuya.batch.trigger" + -------------------------------------------------------------------------------- -- Event Names -- @@ -161,6 +178,14 @@ eventAckDecision :: Text eventAckDecision = "shibuya.ack.decision" +-- | Event recorded when batch-handler execution starts (@shibuya.batch.started@). +eventBatchStarted :: Text +eventBatchStarted = "shibuya.batch.started" + +-- | Event recorded when a batch has been fully finalized (@shibuya.batch.completed@). +eventBatchCompleted :: Text +eventBatchCompleted = "shibuya.batch.completed" + -------------------------------------------------------------------------------- -- SpanArguments Helpers -------------------------------------------------------------------------------- @@ -190,7 +215,7 @@ -------------------------------------------------------------------------------- -- | Create a NewEvent with the given name and attributes. --- Use 'toAttribute' to convert values to Attribute. +-- Use @toAttribute@ to convert values to Attribute. -- -- Example: -- diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -2,10 +2,17 @@ module Main (main) where +import Shibuya.App.BatchSpec qualified +import Shibuya.App.LifecycleSpec qualified +import Shibuya.Batch.ReliabilitySpec qualified +import Shibuya.BatchSpec qualified import Shibuya.Core.AckSpec qualified import Shibuya.Core.RetrySpec qualified import Shibuya.Core.TypesSpec qualified import Shibuya.PolicySpec qualified +import Shibuya.Runner.BatchProcessorSpec qualified +import Shibuya.Runner.BatcherSpec qualified +import Shibuya.Runner.PartitionOrderingSpec qualified import Shibuya.Runner.SupervisedSpec qualified import Shibuya.RunnerSpec qualified import Shibuya.Telemetry.EffectSpec qualified @@ -15,11 +22,18 @@ main :: IO () main = hspec $ do + Shibuya.App.BatchSpec.spec + Shibuya.App.LifecycleSpec.spec + Shibuya.Batch.ReliabilitySpec.spec + Shibuya.BatchSpec.spec describe "Shibuya.Core.Types" Shibuya.Core.TypesSpec.spec describe "Shibuya.Core.Ack" Shibuya.Core.AckSpec.spec describe "Shibuya.Core.Retry" Shibuya.Core.RetrySpec.spec describe "Shibuya.Policy" Shibuya.PolicySpec.spec describe "Shibuya.Runner" Shibuya.RunnerSpec.spec + Shibuya.Runner.BatcherSpec.spec + Shibuya.Runner.BatchProcessorSpec.spec + Shibuya.Runner.PartitionOrderingSpec.spec Shibuya.Runner.SupervisedSpec.spec Shibuya.Telemetry.EffectSpec.spec Shibuya.Telemetry.PropagationSpec.spec diff --git a/test/Shibuya/App/BatchSpec.hs b/test/Shibuya/App/BatchSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/App/BatchSpec.hs @@ -0,0 +1,164 @@ +module Shibuya.App.BatchSpec (spec) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.STM (TVar, atomically, check, newTVarIO, readTVar, writeTVar) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) +import Data.List (nub) +import Data.Text qualified as Text +import Data.Time (UTCTime (..), fromGregorian) +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Adapter.Mock (TrackingAck (..), newTrackingAck, trackingAckHandle) +import Shibuya.App + ( AppError (..), + QueueProcessor (..), + defaultAppConfig, + mkBatchProcessor, + runApp, + stopApp, + waitApp, + ) +import Shibuya.Batch + ( BatchConfig (..), + BatchHandler, + BatchInfo (..), + BatchTrigger (..), + ackAllOk, + defaultBatchConfig, + ) +import Shibuya.Core.Ingested (Ingested, mkIngested) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import Streamly.Data.Stream qualified as Stream +import Test.Hspec + +spec :: Spec +spec = describe "Shibuya.App.Batch" $ do + describe "runApp with a BatchingProcessor" $ do + -- Milestone 1: a bad batch config is rejected by runApp with AppBatchConfigError. + it "rejects a batch config with size 0 (AppBatchConfigError)" $ do + result <- runEff $ runTracingNoop $ do + sizeRef <- liftIO $ newIORef [] + tracking <- newTrackingAck + messages <- createTrackedMessages tracking 1 + let adapter = listAdapter' messages + badConfig = defaultBatchConfig {batchSize = 0} + processor = mkBatchProcessor adapter (recordingHandler sizeRef) badConfig + runApp defaultAppConfig [(ProcessorId "bad", processor)] + case result of + Left (AppBatchConfigError _) -> pure () + Left e -> expectationFailure ("expected AppBatchConfigError, got: " ++ show e) + Right _ -> expectationFailure "expected Left, got a running AppHandle" + + it "rejects PartitionedInOrder with concurrent batching" $ do + result <- runEff $ runTracingNoop $ do + sizeRef <- liftIO $ newIORef [] + tracking <- newTrackingAck + messages <- createTrackedMessages tracking 1 + let adapter = listAdapter' messages + processor = + BatchingProcessor + adapter + (recordingHandler sizeRef) + defaultBatchConfig + PartitionedInOrder + (Async 2) + runApp defaultAppConfig [(ProcessorId "bad-policy", processor)] + case result of + Left (AppPolicyError _) -> pure () + Left e -> expectationFailure ("expected AppPolicyError, got: " ++ show e) + Right _ -> expectationFailure "expected Left, got a running AppHandle" + + -- Milestone 2: all N messages are acked exactly once, seen in batches. + it "processes all messages in batches and acks each exactly once" $ do + (sizes, ids) <- runEff $ runTracingNoop $ do + tracking <- newTrackingAck + sizeRef <- liftIO $ newIORef [] + messages <- createTrackedMessages tracking 10 + let adapter = listAdapter' messages + config = defaultBatchConfig {batchSize = 4, batchTimeout = 60} + processor = mkBatchProcessor adapter (recordingHandler sizeRef) config + res <- runApp defaultAppConfig [(ProcessorId "batch", processor)] + case res of + Left err -> liftIO $ ioError (userError (show err)) + Right appHandle -> do + waitApp appHandle + ss <- liftIO $ readIORef sizeRef + decs <- liftIO $ readIORef tracking.trackedDecisions + pure (ss, map fst decs) + sum (map fst sizes) `shouldBe` 10 + length ids `shouldBe` 10 -- normal path: each message finalized once + length (nub ids) `shouldBe` 10 -- and all distinct + + -- Milestone 3: a partial batch (< batchSize, timeout far off) is flushed on shutdown. + it "flushes a pending partial batch on graceful shutdown" $ do + (sizes, ids, triggers) <- runEff $ runTracingNoop $ do + tracking <- newTrackingAck + sizeRef <- liftIO $ newIORef [] + gate <- liftIO $ newTVarIO False + messages <- createTrackedMessages tracking 3 + let adapter = gatedAdapter gate messages + -- Never fills (size 100) and never times out during the test (60s). + config = defaultBatchConfig {batchSize = 100, batchTimeout = 60} + processor = mkBatchProcessor adapter (recordingHandler sizeRef) config + res <- runApp defaultAppConfig [(ProcessorId "flush", processor)] + case res of + Left err -> liftIO $ ioError (userError (show err)) + Right appHandle -> do + -- Let the 3 messages be ingested and accumulated. + liftIO $ threadDelay 200000 + stopApp appHandle -- ends source -> EOF flush -> waits for done + ss <- liftIO $ readIORef sizeRef + decs <- liftIO $ readIORef tracking.trackedDecisions + pure (ss, map fst decs, map snd ss) + sum (map fst sizes) `shouldBe` 3 + length ids `shouldBe` 3 + length (nub ids) `shouldBe` 3 + all (== TriggerFlush) triggers `shouldBe` True + +-- Handler that records (batch size, trigger) and acks everything OK. +recordingHandler :: + (IOE :> es) => IORef [(Int, BatchTrigger)] -> BatchHandler es String +recordingHandler ref info _batch = do + liftIO $ modifyIORef' ref ((info.size, info.trigger) :) + pure ackAllOk + +testTime :: UTCTime +testTime = UTCTime (fromGregorian 2024 1 1) 0 + +createTrackedMessages :: (IOE :> es) => TrackingAck -> Int -> Eff es [Ingested es String] +createTrackedMessages tracking n = mapM mk [1 .. n] + where + mk i = do + let msgId = MessageId $ "msg-" <> (if i < 10 then "0" else "") <> Text.pack (show i) + env = + (mkEnvelope msgId ("message-" <> show i)) + { cursor = Just (CursorInt i), + enqueuedAt = Just testTime + } + pure $ mkIngested env (trackingAckHandle tracking msgId) + +-- Finite adapter: ends as soon as its list is exhausted. +listAdapter' :: [Ingested es String] -> Adapter es String +listAdapter' messages = + Adapter {adapterName = "test:batch", source = Stream.fromList messages, shutdown = pure ()} + +-- Gated adapter: emits the messages, then blocks until 'shutdown' opens the gate, +-- so a partial batch stays pending until stopApp triggers the flush. Implemented +-- as an unfold that yields each message and, once the list is exhausted, blocks +-- on the gate before ending the stream. +gatedAdapter :: (IOE :> es) => TVar Bool -> [Ingested es String] -> Adapter es String +gatedAdapter gate messages = + Adapter + { adapterName = "test:gated", + source = Stream.unfoldrM step messages, + shutdown = liftIO $ atomically $ writeTVar gate True + } + where + step [] = do + waitGate + pure Nothing + step (m : ms) = pure (Just (m, ms)) + waitGate = liftIO $ atomically $ readTVar gate >>= check diff --git a/test/Shibuya/App/LifecycleSpec.hs b/test/Shibuya/App/LifecycleSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/App/LifecycleSpec.hs @@ -0,0 +1,310 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Shibuya.App.LifecycleSpec (spec) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.NQE.Supervisor (Strategy (..)) +import Control.Concurrent.STM (readTVarIO) +import Data.IORef (modifyIORef', newIORef, readIORef) +import Data.Map.Strict qualified as Map +import Data.Text qualified as Text +import Data.Time (UTCTime (..), diffUTCTime, fromGregorian, getCurrentTime) +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.App + ( AppConfig (..), + QueueProcessor, + ShutdownConfig (..), + SupervisionStrategy (..), + defaultAppConfig, + mkBatchProcessor, + mkProcessor, + runApp, + stopAppGracefully, + waitApp, + ) +import Shibuya.Batch (BatchConfig (..), ackAll, defaultBatchConfig) +import Shibuya.Core.Ack (AckDecision (..), HaltReason (..)) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested, mkIngested) +import Shibuya.Core.Metrics (ProcessorId (..), ProcessorMetrics (..), ProcessorState (..), sampleMetrics) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Internal.App (AppHandle (..)) +import Shibuya.Internal.Runner.Master (startMaster, stopMaster) +import Shibuya.Internal.Runner.Supervised (SupervisedProcessor (..), runSupervised) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +import Shibuya.Telemetry.Effect (Tracing, runTracingNoop) +import Streamly.Data.Stream qualified as Stream +import Test.Hspec +import UnliftIO qualified as UIO + +spec :: Spec +spec = describe "Shibuya.App lifecycle" $ do + it "waitApp returns after a handler halts" $ do + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messages <- createTestMessages 10 + let handler _ = pure $ AckHalt (HaltFatal "stop") + processor = mkProcessor (testAdapter messages) handler + app <- runAppOrFail IgnoreFailures 10 [(ProcessorId "halt", processor)] + waitApp app + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + pure () + + result `shouldBe` Just () + + it "stopAppGracefully returns True promptly after halt" $ do + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messages <- createTestMessages 10 + let handler _ = pure $ AckHalt (HaltFatal "stop") + processor = mkProcessor (testAdapter messages) handler + app <- runAppOrFail IgnoreFailures 10 [(ProcessorId "halt-stop", processor)] + liftIO $ threadDelay 200_000 + startedAt <- liftIO getCurrentTime + drained <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + finishedAt <- liftIO getCurrentTime + pure (drained, diffUTCTime finishedAt startedAt) + + case result of + Just (drained, elapsed) -> do + drained `shouldBe` True + elapsed `shouldSatisfy` (< 2) + Nothing -> expectationFailure "stopAppGracefully timed out" + + it "cancellation sets done" $ do + doneAfterStop <- runEff $ runTracingNoop $ do + master <- startMaster IgnoreAll + sp <- runSupervised master 10 (ProcessorId "cancelled") Unordered Serial infiniteAdapter alwaysAckOk + liftIO $ threadDelay 50_000 + stopMaster master + liftIO $ readTVarIO sp.done + + doneAfterStop `shouldBe` True + + it "waitApp returns after a batch handler halts" $ do + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messages <- createTestMessages 10 + let handler _info _msgs = pure $ ackAll (AckHalt (HaltFatal "stop")) + config = defaultBatchConfig {batchSize = 2, batchTimeout = 0.1} + processor = mkBatchProcessor (testAdapter messages) handler config + app <- runAppOrFail IgnoreFailures 10 [(ProcessorId "batch-halt", processor)] + waitApp app + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + pure () + + result `shouldBe` Just () + + it "graceful completion under StopAllOnFailure does not kill siblings" $ do + countBRef <- newIORef (0 :: Int) + + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messagesA <- createTestMessages 3 + messagesB <- createTestMessages 30 + let handlerA _ = pure AckOk + handlerB _ = do + liftIO $ threadDelay 10_000 + liftIO $ modifyIORef' countBRef (+ 1) + pure AckOk + procA = mkProcessor (testAdapter messagesA) handlerA + procB = mkProcessor (testAdapter messagesB) handlerB + app <- + runAppOrFail + StopAllOnFailure + 10 + [(ProcessorId "complete-A", procA), (ProcessorId "complete-B", procB)] + waitApp app + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + liftIO $ readIORef countBRef + + result `shouldBe` Just 30 + + it "halt under StopAllOnFailure does not kill siblings" $ do + countARef <- newIORef (0 :: Int) + countBRef <- newIORef (0 :: Int) + + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messagesA <- createTestMessages 10 + messagesB <- createTestMessages 30 + let handlerA _ = do + count <- liftIO $ readIORef countARef + liftIO $ modifyIORef' countARef (+ 1) + if count >= 1 + then pure $ AckHalt (HaltFatal "A stops") + else pure AckOk + handlerB _ = do + liftIO $ threadDelay 10_000 + liftIO $ modifyIORef' countBRef (+ 1) + pure AckOk + procA = mkProcessor (testAdapter messagesA) handlerA + procB = mkProcessor (testAdapter messagesB) handlerB + app <- + runAppOrFail + StopAllOnFailure + 10 + [(ProcessorId "halt-A", procA), (ProcessorId "halt-B", procB)] + waitApp app + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + liftIO $ readIORef countBRef + + result `shouldBe` Just 30 + + it "failure under StopAllOnFailure kills siblings and propagates" $ do + countBRef <- newIORef (0 :: Int) + + result <- + UIO.withAsync + ( runEff $ + runTracingNoop $ do + messagesB <- createTestMessages 50 + let handlerA _ = pure AckOk + handlerB _ = do + liftIO $ threadDelay 20_000 + liftIO $ modifyIORef' countBRef (+ 1) + pure AckOk + procA = mkProcessor (failingAfterAdapter 3 "Adapter A source failed!") handlerA + procB = mkProcessor (testAdapter messagesB) handlerB + app <- + runAppOrFail + StopAllOnFailure + 10 + [(ProcessorId "fail-A", procA), (ProcessorId "fail-B", procB)] + liftIO $ threadDelay 500_000 + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + pure () + ) + UIO.waitCatch + + case result of + Left err -> + Text.pack (show err) `shouldSatisfy` Text.isInfixOf "Adapter A source failed" + Right () -> + expectationFailure "Expected adapter A source failure to propagate" + + countB <- readIORef countBRef + countB `shouldSatisfy` (< 50) + + it "IgnoreFailures isolates a failing processor" $ do + countBRef <- newIORef (0 :: Int) + + result <- + UIO.timeout 5_000_000 $ + runEff $ + runTracingNoop $ do + messagesB <- createTestMessages 20 + let handlerA _ = pure AckOk + handlerB _ = do + liftIO $ threadDelay 10_000 + liftIO $ modifyIORef' countBRef (+ 1) + pure AckOk + procA = mkProcessor (failingAfterAdapter 3 "Adapter A source failed!") handlerA + procB = mkProcessor (testAdapter messagesB) handlerB + app <- + runAppOrFail + IgnoreFailures + 10 + [(ProcessorId "ignore-fail-A", procA), (ProcessorId "ignore-fail-B", procB)] + waitApp app + metricsA <- processorMetrics app (ProcessorId "ignore-fail-A") + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + countB <- liftIO $ readIORef countBRef + pure (metricsA, countB) + + case result of + Nothing -> expectationFailure "IgnoreFailures run timed out" + Just (metricsA, countB) -> do + countB `shouldBe` 20 + case metricsA.state of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "Adapter A source failed" + other -> expectationFailure $ "Expected failed processor metrics, got: " ++ show other + +runAppOrFail :: + (IOE :> es, Tracing :> es) => + SupervisionStrategy -> + Int -> + [(ProcessorId, QueueProcessor es)] -> + Eff es (AppHandle es) +runAppOrFail strategy inboxSize processors = do + result <- runApp defaultAppConfig {strategy = strategy, inboxSize = inboxSize} processors + case result of + Left err -> liftIO $ expectationFailure ("runApp failed: " <> show err) >> error "unreachable" + Right app -> pure app + +processorMetrics :: + (IOE :> es) => + AppHandle es -> + ProcessorId -> + Eff es ProcessorMetrics +processorMetrics app pid = + case app of + AppHandle {processors = processorsMap} -> + case Map.lookup pid processorsMap of + Nothing -> liftIO $ expectationFailure ("missing processor: " <> show pid) >> error "unreachable" + Just (SupervisedProcessor {metrics = metricsHandle}, _) -> liftIO $ sampleMetrics metricsHandle + +alwaysAckOk :: (Applicative f) => a -> f AckDecision +alwaysAckOk _ = pure AckOk + +testTime :: UTCTime +testTime = UTCTime (fromGregorian 2024 1 1) 0 + +createTestMessages :: (IOE :> es) => Int -> Eff es [Ingested es String] +createTestMessages n = traverse createTestMessage [1 .. n] + +createTestMessage :: (IOE :> es) => Int -> Eff es (Ingested es String) +createTestMessage i = do + let msgId = MessageId $ "msg-" <> (if i < 10 then "0" else "") <> Text.pack (show i) + env = + (mkEnvelope msgId ("message-" <> show i)) + { cursor = Just (CursorInt i), + enqueuedAt = Just testTime + } + pure $ mkIngested env (AckHandle $ \_ -> pure ()) + +testAdapter :: [Ingested es String] -> Adapter es String +testAdapter messages = + Adapter + { adapterName = "test:list", + source = Stream.fromList messages, + shutdown = pure () + } + +failingAfterAdapter :: (IOE :> es) => Int -> Text.Text -> Adapter es String +failingAfterAdapter goodCount failureText = + Adapter + { adapterName = "test:failing", + source = Stream.unfoldrM step (0 :: Int), + shutdown = pure () + } + where + step n + | n < goodCount = do + msg <- createTestMessage (n + 1) + pure (Just (msg, n + 1)) + | otherwise = error (Text.unpack failureText) + +infiniteAdapter :: (IOE :> es) => Adapter es String +infiniteAdapter = + Adapter + { adapterName = "test:infinite", + source = Stream.unfoldrM step (1 :: Int), + shutdown = pure () + } + where + step n = do + liftIO $ threadDelay 5_000 + msg <- createTestMessage n + pure (Just (msg, n + 1)) diff --git a/test/Shibuya/Batch/ReliabilitySpec.hs b/test/Shibuya/Batch/ReliabilitySpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Batch/ReliabilitySpec.hs @@ -0,0 +1,625 @@ +module Shibuya.Batch.ReliabilitySpec (spec) where + +import Control.Concurrent.STM + ( TVar, + atomically, + check, + modifyTVar', + newTVarIO, + readTVar, + readTVarIO, + writeTVar, + ) +import Control.Monad (forM_) +import Data.Foldable (toList) +import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef) +import Data.List (isInfixOf) +import Data.List.NonEmpty (NonEmpty) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Set qualified as Set +import Data.Text (Text) +import Data.Text qualified as Text +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Shibuya.Adapter (Adapter (..)) +import Shibuya.Adapter.Mock + ( TrackingAck, + getTrackedDecisions, + listAdapter, + mkTrackedIngested, + newTrackingAck, + trackedListAdapter, + ) +import Shibuya.App + ( AppConfig (..), + QueueProcessor (..), + ShutdownConfig (..), + defaultAppConfig, + mkBatchProcessor, + runApp, + stopAppGracefully, + ) +import Shibuya.Batch + ( BatchAck, + BatchConfig (..), + BatchHandler, + BatchInfo (..), + BatchKey (..), + BatchTrigger (..), + ackAllOk, + ackExcept, + defaultBatchConfig, + ) +import Shibuya.Batch.TestHarness +import Shibuya.Core.Ack + ( AckDecision (..), + DeadLetterReason (..), + HaltReason (..), + RetryDelay (..), + ) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested, Message (..), mkIngested) +import Shibuya.Core.Metrics + ( BatchStats (..), + ProcessorId (..), + ProcessorMetrics (..), + ProcessorState (..), + StreamStats (..), + sampleMetrics, + ) +import Shibuya.Core.Types (Envelope (..), MessageId (..)) +import Shibuya.Internal.App (AppHandle (..)) +import Shibuya.Internal.Runner.Supervised (SupervisedProcessor (..)) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +import Shibuya.Telemetry.Effect (Tracing, runTracingNoop) +import Streamly.Data.Stream qualified as Stream +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Monadic (assert, monadicIO, monitor, run) +import UnliftIO (throwIO) +import UnliftIO.Concurrent (threadDelay) + +tshowT :: Int -> Text +tshowT = Text.pack . show + +-- | Run @runApp@ and fail the test loudly if it returns a Left (avoids relying on +-- a @MonadFail@ instance for @Eff@). +runAppOrFail :: + (IOE :> es, Tracing :> es) => + Int -> + [(ProcessorId, QueueProcessor es)] -> + Eff es (AppHandle es) +runAppOrFail inbox procs = do + res <- runApp defaultAppConfig {inboxSize = inbox} procs + case res of + Right app -> pure app + Left e -> liftIO $ ioError (userError ("runApp failed: " <> show e)) + +-- | Read a processor's live metrics directly from its 'SupervisedProcessor' +-- (via the 'AppHandle'), NOT from the Master registry. A finished or halted +-- processor unregisters its metrics from the Master in its @finally@, so +-- @getAppMetrics@ would return nothing for it; the metrics handle itself persists +-- and holds the final counters/state. +metricsFor :: (IOE :> es) => AppHandle es -> ProcessorId -> Eff es ProcessorMetrics +metricsFor app pid = + case Map.lookup pid app.processors of + Just (sp, _) -> liftIO $ sampleMetrics sp.metrics + Nothing -> liftIO $ ioError (userError ("no processor " <> show pid)) + +-------------------------------------------------------------------------------- +-- Shared handlers / builders +-------------------------------------------------------------------------------- + +-- | An observed batch: the info the framework passed and the ids in the batch. +data ObservedBatch = ObservedBatch + { info :: !BatchInfo, + ids :: ![MessageId] + } + deriving stock (Show) + +-- | A handler that records every batch it sees and then returns the given +-- 'BatchAck' (computed from the batch). +recordingHandler :: + (IOE :> es) => + IORef [ObservedBatch] -> + (BatchInfo -> NonEmpty (Message es Int) -> BatchAck) -> + BatchHandler es Int +recordingHandler ref decide bi msgs = do + let obs = ObservedBatch bi [ing.envelope.messageId | ing <- toList msgs] + liftIO $ atomicModifyIORef' ref (\xs -> (obs : xs, ())) + pure (decide bi msgs) + +fixedEnvelopes :: Int -> [Envelope Int] +fixedEnvelopes n = [mkEnvelope i (BatchKey "default") i | i <- [1 .. n]] + +keyedEnvelopes :: [BatchKey] -> Int -> [Envelope Int] +keyedEnvelopes keys n = + [mkEnvelope i (keys !! ((i - 1) `mod` length keys)) i | i <- [1 .. n]] + +-- | Round-robin key membership for scenario #9: msg-1,3,5 belong to @ka@; +-- msg-2,4,6 belong to @kb@. +idBelongsToKey :: BatchKey -> MessageId -> Bool +idBelongsToKey (BatchKey "ka") (MessageId t) = odd (msgNum t) +idBelongsToKey (BatchKey "kb") (MessageId t) = even (msgNum t) +idBelongsToKey _ _ = True + +msgNum :: Text -> Int +msgNum t = read (drop 4 (Text.unpack t)) -- "msg-" + +-- | A gated tracked adapter: yields tracked-ingested messages then blocks until +-- 'shutdown' opens the gate, keeping the stream open so the timeout ticker (not an +-- immediate end-of-input flush) is what emits an accumulated partial batch. +gatedTrackedAdapter :: + (IOE :> es) => TVar Bool -> TrackingAck -> [Envelope Int] -> Adapter es Int +gatedTrackedAdapter gate tracking envs = + Adapter + { adapterName = "test:gated-tracked", + source = Stream.unfoldrM step (map (mkTrackedIngested tracking) envs), + shutdown = liftIO $ atomically $ writeTVar gate True + } + where + step [] = do + liftIO $ atomically $ readTVar gate >>= check + pure Nothing + step (m : ms) = pure (Just (m, ms)) + +-------------------------------------------------------------------------------- +-- Property driver +-------------------------------------------------------------------------------- + +scenarioConfig :: BatchScenario -> BatchConfig es Int +scenarioConfig s = + defaultBatchConfig + { batchSize = s.batchSize, + batchTimeout = fromIntegral s.batchTimeoutMs / 1000, + batchKey = scenarioBatchKey + } + +-- | A batch handler that finalizes each message with the intended per-message +-- decision. For every message whose intended decision is a dead-letter, it emits +-- an override; everything else falls back to 'AckOk'. Exercises both the override +-- and fallback paths. +intendedHandler :: Map MessageId AckDecision -> BatchHandler es Int +intendedHandler intended _info msgs = + pure $ + ackExcept + [ (mid, d) + | ing <- toList msgs, + let mid = ing.envelope.messageId, + Just d <- [Map.lookup mid intended], + d /= AckOk + ] + +-- | Run a scenario end-to-end through the real @runApp@ path, flush the tail via a +-- graceful shutdown, and return the tracked acknowledgements plus final metrics. +runScenario :: + BatchScenario -> + IO ([(MessageId, AckDecision)], ProcessorMetrics) +runScenario s = runEff $ runTracingNoop $ do + tracking <- newTrackingAck + let pid = ProcessorId "batch" + adapter = trackedListAdapter tracking (scenarioEnvelopes s) + proc = mkBatchProcessor adapter (intendedHandler (scenarioIntended s)) (scenarioConfig s) + app <- runAppOrFail 100 [(pid, proc)] + _drained <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + m <- metricsFor app pid + tracked <- getTrackedDecisions tracking + pure (tracked, m) + +-------------------------------------------------------------------------------- +-- Spec +-------------------------------------------------------------------------------- + +spec :: Spec +spec = describe "Shibuya.Batch reliability" $ do + describe "successful-finalization property" $ do + it "finalizes every normal-path message once with the intended decision" $ + withMaxSuccess 50 $ + forAll genScenario $ \s -> monadicIO $ do + (tracked, metrics) <- run (runScenario s) + let expected = scenarioIntended s + case finalizedExactlyOnce tracked expected of + Left err -> do + monitor (counterexample ("successful-finalization violated: " <> err)) + assert False + Right () -> pure () + monitor + ( counterexample + ( "accounting: processed=" + <> show metrics.stats.processed + <> " failed=" + <> show metrics.stats.failed + <> " n=" + <> show s.msgCount + ) + ) + assert (metrics.stats.processed + metrics.stats.failed == s.msgCount) + assert (metrics.batch.batchedMessages == s.msgCount) + + -- Non-vacuity of the checker itself: feed it perturbed tracked lists and prove + -- it fires. Without these, a checker that always returns Right () would pass + -- the property above vacuously. + describe "finalizedExactlyOnce checker (non-vacuity)" $ do + let mid i = MessageId ("msg-" <> tshowT i) + okFor is = Map.fromList [(mid i, AckOk) | i <- is] + failsWith needle = either (needle `isInfixOf`) (const False) + it "accepts a complete, once-each tracked list regardless of order" $ + finalizedExactlyOnce [(mid 2, AckOk), (mid 1, AckOk)] (okFor [1, 2]) + `shouldBe` Right () + it "rejects a double-finalized message" $ + finalizedExactlyOnce [(mid 1, AckOk), (mid 2, AckOk), (mid 1, AckOk)] (okFor [1, 2]) + `shouldSatisfy` failsWith "finalized more than once" + it "rejects a missing finalization" $ + finalizedExactlyOnce [(mid 1, AckOk)] (okFor [1, 2]) + `shouldSatisfy` failsWith "id set mismatch" + it "rejects an unexpected finalization" $ + finalizedExactlyOnce [(mid 1, AckOk), (mid 2, AckOk)] (okFor [1]) + `shouldSatisfy` failsWith "id set mismatch" + it "rejects a finalization with the wrong decision" $ + finalizedExactlyOnce [(mid 1, AckRetry (RetryDelay 1))] (okFor [1]) + `shouldSatisfy` failsWith "wrong decision" + + describe "scenarios" $ do + -- #2 Timeout flush (via a gated adapter so the ticker, not EOF, flushes). + it "flushes a partial batch on timeout" $ do + observedRef <- newIORef ([] :: [ObservedBatch]) + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, observed) <- runEff $ runTracingNoop $ do + gate <- liftIO $ newTVarIO False + let pid = ProcessorId "timeout" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 100, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler = recordingHandler observedRef (\_ _ -> ackAllOk) + adapter = gatedTrackedAdapter gate tracking (fixedEnvelopes 3) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + liftIO $ threadDelay 400000 -- 400 ms > 100 ms timeout: ticker flushes + t0 <- getTrackedDecisions tracking + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + o <- liftIO $ readIORef observedRef + pure (t0, o) + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 3 :: Int]]) + `shouldBe` Right () + map (.info.trigger) observed `shouldSatisfy` elem TriggerTimeout + + -- #3 Partial failure. + it "acks a partial-failure batch: exactly the failed messages dead-letter" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, metrics) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "partial" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 5, batchTimeout = 0.1, batchKey = scenarioBatchKey} + fails = + [ (MessageId "msg-2", AckDeadLetter MaxRetriesExceeded), + (MessageId "msg-4", AckDeadLetter (PoisonPill "bad")) + ] + handler _ _ = pure (ackExcept fails) + adapter = trackedListAdapter tracking (fixedEnvelopes 5) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + t <- getTrackedDecisions tracking + m <- metricsFor app pid + pure (t, m) + let expected = + Map.fromList + [ (MessageId "msg-1", AckOk), + (MessageId "msg-2", AckDeadLetter MaxRetriesExceeded), + (MessageId "msg-3", AckOk), + (MessageId "msg-4", AckDeadLetter (PoisonPill "bad")), + (MessageId "msg-5", AckOk) + ] + finalizedExactlyOnce tracked expected `shouldBe` Right () + metrics.stats.failed `shouldBe` 2 + metrics.stats.processed `shouldBe` 3 + + -- #4 Batch-handler exception -> whole-batch AckRetry, app survives. + it "on batch-handler exception, every message retries and the app survives" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, drained) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "exc" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 4, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler _ _ = error "boom in batch" + adapter = trackedListAdapter tracking (fixedEnvelopes 4) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + d <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + t <- getTrackedDecisions tracking + pure (t, d) + let expected = Map.fromList [(MessageId ("msg-" <> tshowT i), AckRetry (RetryDelay 0)) | i <- [1 .. 4 :: Int]] + finalizedExactlyOnce tracked expected `shouldBe` Right () + drained `shouldBe` True + + -- #4b Batcher consumer exception -> processor failure, not clean completion. + it "propagates a throwing batchKey as processor failure without double finalizing" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, mState, done, drained) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "batch-key-failure" + cfg = + (defaultBatchConfig @_ @Int) + { batchSize = 10, + batchTimeout = 30, + batchKey = \env -> + if env.messageId == MessageId "msg-3" + then error "boom in batchKey" + else scenarioBatchKey env + } + handler _ _ = pure ackAllOk + adapter = trackedListAdapter tracking (fixedEnvelopes 5) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + liftIO $ threadDelay 300000 + m <- metricsFor app pid + doneState <- case Map.lookup pid app.processors of + Just (sp, _) -> liftIO $ readTVarIO sp.done + Nothing -> liftIO $ ioError (userError "missing batch-key-failure processor") + d <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + t <- getTrackedDecisions tracking + pure (t, m.state, doneState, d) + case mState of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "boom in batchKey" + other -> expectationFailure ("expected Failed, got: " <> show other) + done `shouldBe` True + drained `shouldBe` True + tracked `shouldBe` [] + + -- #5 Transient finalizer retry. + it "retries a transient finalizer failure and records one success" $ do + attemptRef <- newIORef (0 :: Int) + trackRef <- newIORef ([] :: [(MessageId, AckDecision)]) + (tracked, attempts) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "retry" + mid = MessageId "msg-1" + flakyHandle = AckHandle $ \d -> do + n <- liftIO $ atomicModifyIORef' attemptRef (\c -> (c + 1, c)) + if n < 2 + then liftIO $ throwIO (userError "transient") + else liftIO $ modifyIORef' trackRef ((mid, d) :) + ing = mkIngested (mkEnvelope 1 (BatchKey "default") 1) flakyHandle + cfg = (defaultBatchConfig @_ @Int) {batchSize = 1, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler _ _ = pure ackAllOk + proc = mkBatchProcessor (listAdapter [ing]) handler cfg + app <- runAppOrFail 100 [(pid, proc)] + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + t <- liftIO $ readIORef trackRef + a <- liftIO $ readIORef attemptRef + pure (t, a) + tracked `shouldBe` [(MessageId "msg-1", AckOk)] + attempts `shouldBe` 3 + + -- #6 Permanent finalizer failure -> fail loud, name the id, still attempt others. + it "exhausts permanent finalizer failures and halts loudly with the failed MessageId" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, mState) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "perm" + failHandle = AckHandle $ \_ -> liftIO $ throwIO (userError "permanent") + ing1 = mkIngested (mkEnvelope 1 (BatchKey "default") 1) failHandle + ing2 = mkTrackedIngested tracking (mkEnvelope 2 (BatchKey "default") 2) + cfg = (defaultBatchConfig @_ @Int) {batchSize = 2, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler _ _ = pure ackAllOk + proc = mkBatchProcessor (listAdapter [ing1, ing2]) handler cfg + app <- runAppOrFail 100 [(pid, proc)] + liftIO $ threadDelay 700000 -- allow the [10,50,250]ms retry schedule to exhaust + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + t <- getTrackedDecisions tracking + m <- metricsFor app pid + pure (t, m.state) + case mState of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "msg-1" + other -> expectationFailure ("expected Failed, got: " <> show other) + -- The other message was still attempted and successfully finalized. + tracked `shouldBe` [(MessageId "msg-2", AckOk)] + + -- #7 Halt in a batch, with isolation of a second processor. + it "halt in one batch finalizes that batch, halts its processor, and spares others" $ do + trackingA <- runEff $ runTracingNoop newTrackingAck + trackingB <- runEff $ runTracingNoop newTrackingAck + (trackedA, trackedB, mA) <- runEff $ runTracingNoop $ do + let pidA = ProcessorId "halt-A" + pidB = ProcessorId "ok-B" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 3, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handlerA _ _ = pure (ackExcept [(MessageId "msg-2", AckHalt (HaltFatal "halt in batch"))]) + handlerB _ _ = pure (ackExcept []) + procA = mkBatchProcessor (trackedListAdapter trackingA (fixedEnvelopes 3)) handlerA cfg + procB = mkBatchProcessor (trackedListAdapter trackingB (fixedEnvelopes 3)) handlerB cfg + app <- runAppOrFail 100 [(pidA, procA), (pidB, procB)] + liftIO $ threadDelay 300000 + mA' <- metricsFor app pidA + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 1}) app + tA <- getTrackedDecisions trackingA + tB <- getTrackedDecisions trackingB + pure (tA, tB, mA') + finalizedExactlyOnce + trackedA + ( Map.fromList + [ (MessageId "msg-1", AckOk), + (MessageId "msg-2", AckHalt (HaltFatal "halt in batch")), + (MessageId "msg-3", AckOk) + ] + ) + `shouldBe` Right () + case mA.state of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "halt in batch" + other -> expectationFailure ("expected A Failed, got: " <> show other) + finalizedExactlyOnce + trackedB + (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 3 :: Int]]) + `shouldBe` Right () + + -- #8 Graceful drain flush (gated adapter so the flush comes from shutdown). + it "flushes a pending partial batch on graceful shutdown" $ do + observedRef <- newIORef ([] :: [ObservedBatch]) + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, drained, observed) <- runEff $ runTracingNoop $ do + gate <- liftIO $ newTVarIO False + let pid = ProcessorId "drain" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 100, batchTimeout = 30, batchKey = scenarioBatchKey} + handler = recordingHandler observedRef (\_ _ -> ackAllOk) + adapter = gatedTrackedAdapter gate tracking (fixedEnvelopes 3) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + liftIO $ threadDelay 100000 -- accumulate (no size/timeout flush) + d <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + t <- getTrackedDecisions tracking + o <- liftIO $ readIORef observedRef + pure (t, d, o) + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 3 :: Int]]) + `shouldBe` Right () + drained `shouldBe` True + map (.info.trigger) observed `shouldSatisfy` elem TriggerFlush + + -- #9 Multiple batch keys accumulate independently. + it "accumulates independent per-key batches and acks all exactly once" $ do + observedRef <- newIORef ([] :: [ObservedBatch]) + tracking <- runEff $ runTracingNoop newTrackingAck + (tracked, observed) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "multi-key" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 3, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler = recordingHandler observedRef (\_ _ -> ackAllOk) + adapter = trackedListAdapter tracking (keyedEnvelopes [BatchKey "ka", BatchKey "kb"] 6) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 100 [(pid, proc)] + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + t <- getTrackedDecisions tracking + o <- liftIO $ readIORef observedRef + pure (t, o) + forM_ observed $ \ob -> + all (idBelongsToKey ob.info.batchKey) ob.ids `shouldBe` True + map (.info.batchKey) observed `shouldSatisfy` (\ks -> BatchKey "ka" `elem` ks && BatchKey "kb" `elem` ks) + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 6 :: Int]]) + `shouldBe` Right () + + -- #10 Per-key FIFO under concurrency (Async 2). + it "preserves per-key FIFO while allowing cross-key concurrency" $ do + (violated, tracked) <- runEff $ runTracingNoop $ do + active <- liftIO $ newTVarIO Set.empty + violatedVar <- liftIO $ newTVarIO False + tracking <- newTrackingAck + let pid = ProcessorId "fifo" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 1, batchTimeout = 0.1, batchKey = scenarioBatchKey} + handler info _ = do + liftIO $ atomically $ do + a <- readTVar active + if info.batchKey `Set.member` a + then writeTVar violatedVar True + else pure () + modifyTVar' active (Set.insert info.batchKey) + liftIO $ threadDelay 30000 + liftIO $ atomically $ modifyTVar' active (Set.delete info.batchKey) + pure ackAllOk + keys = [BatchKey "same", BatchKey "same", BatchKey "same", BatchKey "other", BatchKey "other"] + adapter = trackedListAdapter tracking (keyedEnvelopes keys 5) + proc = + BatchingProcessor + { adapter = adapter, + batchHandler = handler, + batchConfig = cfg, + ordering = Unordered, + concurrency = Async 2 + } + app <- runAppOrFail 100 [(pid, proc)] + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 5}) app + v <- liftIO $ readTVarIO violatedVar + t <- getTrackedDecisions tracking + pure (v, t) + violated `shouldBe` False + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 5 :: Int]]) + `shouldBe` Right () + + -- #11 Backpressure liveness (limited: a no-loss proxy, not a memory measurement). + it "under a tiny inbox and a slow handler, loses no messages (backpressure liveness; see limitation)" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + tracked <- runEff $ runTracingNoop $ do + let pid = ProcessorId "backpressure" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 4, batchTimeout = 0.05, batchKey = scenarioBatchKey} + handler _ _ = do + liftIO $ threadDelay 5000 -- 5 ms per batch + pure ackAllOk + adapter = trackedListAdapter tracking (fixedEnvelopes 20) + proc = mkBatchProcessor adapter handler cfg + app <- runAppOrFail 2 [(pid, proc)] -- inbox size 2 + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 10}) app + getTrackedDecisions tracking + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 20 :: Int]]) + `shouldBe` Right () + + -- #12 Keyed scheduler pending bound. + it "bounds upstream pulls while keyed batch handlers are blocked" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + pulledRef <- newIORef (0 :: Int) + (pulledWhileBlocked, tracked) <- runEff $ runTracingNoop $ do + gate <- liftIO $ newTVarIO False + let pid = ProcessorId "bounded-keyed" + totalMessages = 50 + inboxSize = 2 + pendingLimit = max 2 (2 * 2) + allowedPulls = inboxSize + inboxSize + pendingLimit + 2 + 10 + source = Stream.unfoldrM (pullStep tracking pulledRef totalMessages) (1 :: Int) + adapter = + Adapter + { adapterName = "test:counting-source", + source = source, + shutdown = pure () + } + cfg = (defaultBatchConfig @_ @Int) {batchSize = 1, batchTimeout = 30, batchKey = scenarioBatchKey} + handler _ _ = do + liftIO $ atomically $ readTVar gate >>= check + pure ackAllOk + proc = + BatchingProcessor + { adapter = adapter, + batchHandler = handler, + batchConfig = cfg, + ordering = Unordered, + concurrency = Async 2 + } + app <- runAppOrFail inboxSize [(pid, proc)] + liftIO $ threadDelay 300000 + pulled <- liftIO $ readIORef pulledRef + liftIO $ atomically $ writeTVar gate True + _ <- stopAppGracefully (ShutdownConfig {drainTimeout = 10}) app + t <- getTrackedDecisions tracking + liftIO $ pulled `shouldSatisfy` (<= allowedPulls) + pure (pulled, t) + pulledWhileBlocked `shouldSatisfy` (< 50) + finalizedExactlyOnce tracked (Map.fromList [(MessageId ("msg-" <> tshowT i), AckOk) | i <- [1 .. 50 :: Int]]) + `shouldBe` Right () + + -- #13 Forced shutdown must not leak keyed scheduler workers. + it "does not finalize additional messages after forced shutdown returns" $ do + tracking <- runEff $ runTracingNoop newTrackingAck + (drained, trackedBefore, trackedAfter) <- runEff $ runTracingNoop $ do + let pid = ProcessorId "shutdown-no-leak" + cfg = (defaultBatchConfig @_ @Int) {batchSize = 1, batchTimeout = 30, batchKey = scenarioBatchKey} + handler _ _ = do + liftIO $ threadDelay 200000 + pure ackAllOk + proc = + BatchingProcessor + { adapter = trackedListAdapter tracking (fixedEnvelopes 30), + batchHandler = handler, + batchConfig = cfg, + ordering = Unordered, + concurrency = Async 2 + } + app <- runAppOrFail 2 [(pid, proc)] + liftIO $ threadDelay 50000 + d <- stopAppGracefully (ShutdownConfig {drainTimeout = 0.05}) app + t1 <- getTrackedDecisions tracking + liftIO $ threadDelay 300000 + t2 <- getTrackedDecisions tracking + pure (d, t1, t2) + drained `shouldBe` False + trackedAfter `shouldBe` trackedBefore + Set.size (Set.fromList (map fst trackedAfter)) `shouldBe` length trackedAfter + +pullStep :: + (IOE :> es) => + TrackingAck -> + IORef Int -> + Int -> + Int -> + Eff es (Maybe (Ingested es Int, Int)) +pullStep tracking pulledRef totalMessages i + | i > totalMessages = pure Nothing + | otherwise = do + liftIO $ modifyIORef' pulledRef (+ 1) + pure (Just (mkTrackedIngested tracking (mkEnvelope i (BatchKey "default") i), i + 1)) diff --git a/test/Shibuya/Batch/TestHarness.hs b/test/Shibuya/Batch/TestHarness.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Batch/TestHarness.hs @@ -0,0 +1,153 @@ +module Shibuya.Batch.TestHarness + ( -- * Scenario model + BatchScenario (..), + genScenario, + + -- * Envelope builders + mkEnvelope, + scenarioEnvelopes, + scenarioBatchKey, + scenarioMsgIds, + scenarioIntended, + + -- * Exactly-once checking + finalizedExactlyOnce, + ) +where + +import Data.List (sort) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as Text +import Shibuya.Batch (BatchKey (..)) +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..)) +import Shibuya.Core.Types qualified as Core +import Test.QuickCheck + +-- | A randomized batch schedule. Each message index @i@ in @[1 .. msgCount]@ has a +-- distinct 'MessageId' @"msg-i"@, an assigned 'BatchKey' (@keyOf ! i@), and an +-- intended finalized decision (@outcomeOf ! i@, always 'AckOk' or an +-- 'AckDeadLetter'). @batchSize@ and @batchTimeoutMs@ configure the batcher. +data BatchScenario = BatchScenario + { msgCount :: !Int, + batchSize :: !Int, + batchTimeoutMs :: !Int, + keyOf :: !(Map Int BatchKey), + outcomeOf :: !(Map Int AckDecision) + } + deriving stock (Show) + +tshow :: (Show a) => a -> Text +tshow = Text.pack . show + +-- | Generate a randomized scenario: 1..40 messages, a batch size that may be +-- smaller than, equal to, or larger than the message count (so some runs never +-- flush by size), a short timeout, 1..4 batch keys, and a per-message outcome that +-- is mostly 'AckOk' with an occasional dead-letter. +genScenario :: Gen BatchScenario +genScenario = do + n <- choose (1, 40) + bs <- choose (1, n + 5) + toMs <- choose (20, 200) + numKeys <- choose (1, min 4 n) + keys <- vectorOf n (chooseKey numKeys) + outs <- vectorOf n genOutcome + pure + BatchScenario + { msgCount = n, + batchSize = bs, + batchTimeoutMs = toMs, + keyOf = Map.fromList (zip [1 ..] keys), + outcomeOf = Map.fromList (zip [1 ..] outs) + } + where + chooseKey k = do + j <- choose (1, k) + pure (BatchKey ("k" <> tshow (j :: Int))) + genOutcome = + frequency + [ (3, pure AckOk), + (1, AckDeadLetter <$> elements [MaxRetriesExceeded, PoisonPill "qc", InvalidPayload "qc"]) + ] + +instance Arbitrary BatchScenario where + arbitrary = genScenario + +-- | Build one envelope carrying a payload, with the given id number and batch key +-- stamped into the 'partition' field (so a pure @batchKey@ config function can +-- recover it). +mkEnvelope :: Int -> BatchKey -> Int -> Envelope Int +mkEnvelope i key payload = + (Core.mkEnvelope (MessageId ("msg-" <> tshow i)) payload) + { cursor = Just (CursorInt i), + partition = Just key.unBatchKey + } + +-- | The envelopes for a scenario, in index order, each carrying its assigned key. +scenarioEnvelopes :: BatchScenario -> [Envelope Int] +scenarioEnvelopes s = + [ mkEnvelope i (keyFor i) i + | i <- [1 .. s.msgCount] + ] + where + keyFor i = fromMaybe (BatchKey "default") (Map.lookup i s.keyOf) + +-- | The pure @batchKey@ function to hand to 'BatchConfig': recover the key from the +-- envelope's partition, defaulting to @"default"@. +scenarioBatchKey :: Envelope Int -> BatchKey +scenarioBatchKey env = BatchKey (fromMaybe "default" env.partition) + +-- | The set of message ids a scenario produces. +scenarioMsgIds :: BatchScenario -> [MessageId] +scenarioMsgIds s = [MessageId ("msg-" <> tshow i) | i <- [1 .. s.msgCount]] + +-- | The intended finalized decision per message id. +scenarioIntended :: BatchScenario -> Map MessageId AckDecision +scenarioIntended s = + Map.fromList + [ (MessageId ("msg-" <> tshow i), fromMaybe AckOk (Map.lookup i s.outcomeOf)) + | i <- [1 .. s.msgCount] + ] + +-- | The successful-finalization checker. Given the raw tracked list of +-- @(messageId, decision)@ pairs (as recorded by a 'Shibuya.Adapter.Mock.TrackingAck', +-- which appends one entry per @finalize@ call) and the expected finalized decision +-- per id, return @Right ()@ iff every expected id appears exactly once and carries +-- its expected decision, and no unexpected id appears. Otherwise return a +-- human-readable explanation for use as a QuickCheck counterexample or an HSpec +-- failure message. +finalizedExactlyOnce :: + [(MessageId, AckDecision)] -> + Map MessageId AckDecision -> + Either String () +finalizedExactlyOnce tracked expected + | not (null dupes) = + Left ("finalized more than once: " <> show dupes) + | trackedIds /= expectedIds = + Left + ( "id set mismatch: missing=" + <> show (expectedIds `minus` trackedIds) + <> " extra=" + <> show (trackedIds `minus` expectedIds) + ) + | not (null wrong) = + Left ("wrong decision: " <> show wrong) + | otherwise = Right () + where + counts :: Map MessageId Int + counts = Map.fromListWith (+) [(mid, 1 :: Int) | (mid, _) <- tracked] + dupes = [mid | (mid, c) <- Map.toList counts, c /= 1] + trackedIds = sort (Map.keys counts) + expectedIds = sort (Map.keys expected) + minus xs ys = filter (`notElem` ys) xs + -- Since each id appears exactly once (checked above via dupes), a simple + -- lookup of the single decision is well-defined here. + got = Map.fromList tracked + wrong = + [ (mid, want, Map.lookup mid got) + | (mid, want) <- Map.toList expected, + Map.lookup mid got /= Just want + ] diff --git a/test/Shibuya/BatchSpec.hs b/test/Shibuya/BatchSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/BatchSpec.hs @@ -0,0 +1,57 @@ +module Shibuya.BatchSpec (spec) where + +import Data.Map.Strict qualified as Map +import Shibuya.Batch +import Shibuya.Core.Ack (AckDecision (..), DeadLetterReason (..), RetryDelay (..)) +import Shibuya.Core.Types (MessageId (..)) +import Test.Hspec + +-- | A concrete 'BatchConfig' so field access and 'validateBatchConfig' resolve +-- without ambiguous @es@/@msg@ type variables. +cfg :: BatchConfig '[] Int +cfg = defaultBatchConfig + +spec :: Spec +spec = describe "Shibuya.Batch" $ do + describe "defaultBatchConfig" $ do + it "has size 100 and 1s timeout" $ do + cfg.batchSize `shouldBe` 100 + cfg.batchTimeout `shouldBe` 1 + it "routes every message to the default key" $ + cfg.batchKey undefined `shouldBe` defaultBatchKey + + describe "validateBatchConfig" $ do + it "accepts the default config" $ + validateBatchConfig cfg `shouldBe` Right () + it "rejects size 0" $ + validateBatchConfig cfg {batchSize = 0} + `shouldBe` Left (BatchSizeNotPositive 0) + it "rejects non-positive timeout" $ + validateBatchConfig cfg {batchTimeout = 0} + `shouldBe` Left (BatchTimeoutNotPositive 0) + it "rejects non-positive tick interval" $ + validateBatchConfig cfg {tickInterval = Just 0} + `shouldBe` Left (TickIntervalNotPositive 0) + + describe "BatchAck smart constructors" $ do + it "ackAllOk falls back to AckOk with no overrides" $ do + ackAllOk.fallback `shouldBe` AckOk + Map.null ackAllOk.decisions `shouldBe` True + it "ackAll sets the fallback for everything" $ do + let a = ackAll (AckRetry (RetryDelay 5)) + Map.null a.decisions `shouldBe` True + a.fallback `shouldBe` AckRetry (RetryDelay 5) + it "ackExcept keeps AckOk fallback and records overrides" $ do + let a = ackExcept [(MessageId "m1", AckDeadLetter MaxRetriesExceeded)] + a.fallback `shouldBe` AckOk + Map.lookup (MessageId "m1") a.decisions + `shouldBe` Just (AckDeadLetter MaxRetriesExceeded) + it "failMessages dead-letters the listed ids" $ do + let a = failMessages [(MessageId "bad", PoisonPill "nope")] + a.fallback `shouldBe` AckOk + Map.lookup (MessageId "bad") a.decisions + `shouldBe` Just (AckDeadLetter (PoisonPill "nope")) + it "withFallback uses the given fallback" $ do + let a = withFallback (AckDeadLetter MaxRetriesExceeded) [(MessageId "ok", AckOk)] + a.fallback `shouldBe` AckDeadLetter MaxRetriesExceeded + Map.lookup (MessageId "ok") a.decisions `shouldBe` Just AckOk diff --git a/test/Shibuya/Core/RetrySpec.hs b/test/Shibuya/Core/RetrySpec.hs --- a/test/Shibuya/Core/RetrySpec.hs +++ b/test/Shibuya/Core/RetrySpec.hs @@ -2,12 +2,11 @@ module Shibuya.Core.RetrySpec (spec) where -import Data.HashMap.Strict qualified as HashMap import Data.Time (nominalDiffTimeToSeconds) import Effectful (runEff) import Shibuya.Core.Ack (AckDecision (..), RetryDelay (..)) import Shibuya.Core.Retry -import Shibuya.Core.Types (Attempt (..), Envelope (..), MessageId (..)) +import Shibuya.Core.Types (Attempt (..), Envelope (..), MessageId (..), mkEnvelope) import Test.Hspec import Test.QuickCheck @@ -16,17 +15,7 @@ testEnvelope :: Maybe Attempt -> Envelope () testEnvelope a = - Envelope - { messageId = MessageId "test", - cursor = Nothing, - partition = Nothing, - enqueuedAt = Nothing, - traceContext = Nothing, - headers = Nothing, - attempt = a, - attributes = HashMap.empty, - payload = () - } + (mkEnvelope (MessageId "test") ()) {attempt = a} spec :: Spec spec = do @@ -55,9 +44,9 @@ `shouldSatisfy` (\x -> x > 3.99 && x <= 4.0) it "delay never exceeds the clamped baseExp" $ property $ \(NonNegative n) -> - forAll (choose (0.0, 0.999999)) $ \sample -> + forAll (choose (0.0, 0.999999)) $ \jitterSample -> let attemptN = Attempt (fromIntegral (n :: Int)) - delay = secondsOf (exponentialBackoffPure p attemptN sample) + delay = secondsOf (exponentialBackoffPure p attemptN jitterSample) in delay >= 0 && delay <= 300 describe "EqualJitter" $ do @@ -69,9 +58,9 @@ `shouldSatisfy` (\x -> x > 3.99 && x <= 4.0) it "delay is in [baseExp/2, baseExp]" $ property $ \(NonNegative n) -> - forAll (choose (0.0, 0.999999)) $ \sample -> + forAll (choose (0.0, 0.999999)) $ \jitterSample -> let attemptN = Attempt (fromIntegral (n :: Int)) - delay = secondsOf (exponentialBackoffPure p attemptN sample) + delay = secondsOf (exponentialBackoffPure p attemptN jitterSample) in delay >= 0 && delay <= 300 describe "Attempt 0" $ do @@ -80,8 +69,8 @@ `shouldBe` 1.0 it "with FullJitter is in [0, base)" $ property $ - forAll (choose (0.0, 0.999999)) $ \sample -> - let delay = secondsOf (exponentialBackoffPure (defaultBackoffPolicy {jitter = FullJitter}) (Attempt 0) sample) + forAll (choose (0.0, 0.999999)) $ \jitterSample -> + let delay = secondsOf (exponentialBackoffPure (defaultBackoffPolicy {jitter = FullJitter}) (Attempt 0) jitterSample) in delay >= 0 && delay < 1.0 describe "retryWithBackoff" $ do diff --git a/test/Shibuya/Core/TypesSpec.hs b/test/Shibuya/Core/TypesSpec.hs --- a/test/Shibuya/Core/TypesSpec.hs +++ b/test/Shibuya/Core/TypesSpec.hs @@ -2,7 +2,6 @@ module Shibuya.Core.TypesSpec (spec) where -import Data.HashMap.Strict qualified as HashMap import Data.Time (UTCTime (..), fromGregorian) import Shibuya.Core.Types import Test.Hspec @@ -84,22 +83,27 @@ mapped = fmap show env mapped.headers `shouldBe` Just [("content-type", "application/json")] + it "mkEnvelope defaults optional metadata" $ do + let env = mkEnvelope (MessageId "test-id") ("payload" :: String) + env.messageId `shouldBe` MessageId "test-id" + env.cursor `shouldBe` Nothing + env.partition `shouldBe` Nothing + env.enqueuedAt `shouldBe` Nothing + env.traceContext `shouldBe` Nothing + env.headers `shouldBe` Nothing + env.attempt `shouldBe` Nothing + env.payload `shouldBe` "payload" + it "defaults headers to Nothing in the test helper" $ do (testEnvelope (1 :: Int)).headers `shouldBe` Nothing -- Test helper testEnvelope :: msg -> Envelope msg testEnvelope msg = - Envelope - { messageId = MessageId "test-id", - cursor = Just (CursorInt 42), + (mkEnvelope (MessageId "test-id") msg) + { cursor = Just (CursorInt 42), partition = Just "partition-0", - enqueuedAt = Just testTime, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = msg + enqueuedAt = Just testTime } testTime :: UTCTime diff --git a/test/Shibuya/PolicySpec.hs b/test/Shibuya/PolicySpec.hs --- a/test/Shibuya/PolicySpec.hs +++ b/test/Shibuya/PolicySpec.hs @@ -4,11 +4,10 @@ import Shibuya.Policy import Test.Hspec -import Prelude hiding (Ordering) spec :: Spec spec = do - describe "Ordering" $ do + describe "OrderingPolicy" $ do it "all constructors are distinguishable" $ do StrictInOrder `shouldNotBe` PartitionedInOrder PartitionedInOrder `shouldNotBe` Unordered @@ -63,7 +62,32 @@ it "allows Async" $ do validatePolicy Unordered (Async 10) `shouldBe` Right () + describe "validatePolicy matrix" $ do + let ok = True + rejected = False + cases = + [ (StrictInOrder, Serial, ok), + (StrictInOrder, Ahead 4, rejected), + (StrictInOrder, Async 4, rejected), + (PartitionedInOrder, Serial, ok), + (PartitionedInOrder, Ahead 4, ok), + (PartitionedInOrder, Async 4, ok), + (Unordered, Serial, ok), + (Unordered, Ahead 4, ok), + (Unordered, Async 4, ok) + ] + mapM_ + ( \(ordering, concurrency, expected) -> + it (show ordering <> " + " <> show concurrency) $ + isRight (validatePolicy ordering concurrency) `shouldBe` expected + ) + cases + -- Helper isLeft :: Either a b -> Bool isLeft (Left _) = True isLeft (Right _) = False + +isRight :: Either a b -> Bool +isRight (Right _) = True +isRight (Left _) = False diff --git a/test/Shibuya/Runner/BatchProcessorSpec.hs b/test/Shibuya/Runner/BatchProcessorSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Runner/BatchProcessorSpec.hs @@ -0,0 +1,323 @@ +module Shibuya.Runner.BatchProcessorSpec (spec) where + +import Control.Concurrent.STM + ( atomically, + modifyTVar', + newTVarIO, + readTVar, + readTVarIO, + writeTVar, + ) +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 (ProcessorHalt (..)) +import Shibuya.Adapter.Mock + ( TrackingAck (..), + getTrackedDecisions, + newTrackingAck, + trackingAckHandle, + ) +import Shibuya.Batch + ( BatchInfo (..), + BatchKey (..), + BatchTrigger (..), + ackAllOk, + ackExcept, + defaultBatchKey, + withFallback, + ) +import Shibuya.Core.Ack + ( AckDecision (..), + DeadLetterReason (..), + HaltReason (..), + RetryDelay (..), + ) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested, Message (..), mkIngested) +import Shibuya.Core.Metrics + ( BatchStats (..), + ProcessorId (..), + ProcessorMetrics (..), + StreamStats (..), + ) +import Shibuya.Core.Types (Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Internal.Runner.BatchProcessor (runBatchesWithMetrics) +import Shibuya.Policy (Concurrency (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import Test.Hspec +import UnliftIO (throwIO, try) +import UnliftIO.Concurrent (threadDelay) + +spec :: Spec +spec = describe "Shibuya.Internal.Runner.BatchProcessor" $ do + 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)] + + 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 (ProcessorHalt (HaltFatal msg)) -> + msg `shouldSatisfy` Text.isInfixOf "perm-1" + Left other -> expectationFailure ("unexpected halt reason: " <> show other) + Right () -> expectationFailure "expected ProcessorHalt 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 diff --git a/test/Shibuya/Runner/BatcherSpec.hs b/test/Shibuya/Runner/BatcherSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Runner/BatcherSpec.hs @@ -0,0 +1,229 @@ +module Shibuya.Runner.BatcherSpec (spec) where + +import Control.Concurrent (threadDelay) +import Data.List (sort) +import Data.List.NonEmpty qualified as NE +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Word (Word64) +import Effectful (Effect) +import Shibuya.Batch + ( BatchConfig (..), + BatchInfo (..), + BatchKey (..), + BatchTrigger (..), + ) +import Shibuya.Core.AckHandle (AckHandle (..)) +import Shibuya.Core.Ingested (Ingested (..), mkIngested) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Internal.Runner.Batcher + ( ReadyBatch, + emptyBatcherState, + runBatcher, + stepArrival, + stepFlush, + stepTick, + ) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream qualified as Stream +import Test.Hspec +import Test.QuickCheck + +-- The engine is parameterized by an effect stack and payload; the pure core +-- treats them as phantom. We pick es = 'E' (an empty effect stack, kind pinned +-- to '[Effect]') and msg = String for the tests. + +type E = ('[] :: [Effect]) + +tshow :: Int -> Text +tshow = Text.pack . show + +-- | Build a no-op-ack Ingested with a unique id, a sequence cursor, and a +-- partition equal to its batch key. +mkIng :: Int -> BatchKey -> Ingested E String +mkIng i (BatchKey k) = + mkIngested + ( (mkEnvelope (MessageId ("m-" <> tshow i)) ("payload-" <> show i)) + { cursor = Just (CursorInt i), + partition = Just k + } + ) + (AckHandle (\_ -> pure ())) + +-- Key the config on the message's partition (Just k -> BatchKey k). +partitionKeyConfig :: Int -> BatchConfig E String +partitionKeyConfig sz = + BatchConfig + { batchSize = sz, + batchTimeout = 5, + batchKey = \env -> BatchKey (fromMaybe "default" env.partition), + tickInterval = Nothing + } + +secondsToNanos :: Double -> Word64 +secondsToNanos seconds = round (seconds * 1e9) + +-- | Run a synthetic event list through the pure core and collect emitted batches. +-- Events carry their own virtual time; a final flush drains the remainder so the +-- conservation property can compare the full input against the full output. +data Ev + = EvArr Int BatchKey -- an arrival: unique id, key + | EvTick -- a ticker scan + deriving stock (Show) + +runEvents :: BatchConfig E String -> [(Double, Ev)] -> [ReadyBatch E String] +runEvents cfg = go emptyBatcherState + where + go st [] = snd (stepFlush st) + go st ((secs, ev) : rest) = + let now = secondsToNanos secs + (st', out) = case ev of + EvArr i k -> stepArrival cfg now (mkIng i k) st + EvTick -> stepTick cfg now st + in out ++ go st' rest + +batchIds :: [ReadyBatch E String] -> [MessageId] +batchIds bs = [ing.envelope.messageId | (_, ne) <- bs, ing <- NE.toList ne] + +batchCursors :: NE.NonEmpty (Ingested E String) -> [Int] +batchCursors ne = [c | ing <- NE.toList ne, Just (CursorInt c) <- [ing.envelope.cursor]] + +spec :: Spec +spec = describe "Shibuya.Internal.Runner.Batcher" $ do + describe "pure core (deterministic, no threads)" $ do + it "emits a size-triggered batch of exactly batchSize" $ do + let cfg = partitionKeyConfig 2 + evs = [(0, EvArr 0 "a"), (0, EvArr 1 "a")] + out = runEvents cfg evs + map (\(info, _) -> (info.trigger, info.size)) out + `shouldBe` [(TriggerSize, 2)] + batchIds out `shouldBe` [MessageId "m-0", MessageId "m-1"] + + it "keeps different keys in independent groups" $ do + let cfg = partitionKeyConfig 2 + evs = [(0, EvArr 0 "a"), (0, EvArr 1 "b"), (0, EvArr 2 "a"), (0, EvArr 3 "b")] + out = runEvents cfg evs + -- Two size batches, one per key; both size 2. + map (\(info, _) -> (info.batchKey, info.size, info.trigger)) out + `shouldMatchList` [(BatchKey "a", 2, TriggerSize), (BatchKey "b", 2, TriggerSize)] + + it "flushes a partial group at end of input with TriggerFlush" $ do + let cfg = partitionKeyConfig 10 + evs = [(0, EvArr 0 "a"), (0, EvArr 1 "a")] + out = runEvents cfg evs + map (\(info, _) -> (info.trigger, info.size)) out `shouldBe` [(TriggerFlush, 2)] + + it "emits a timeout batch once batchTimeout has elapsed" $ do + let cfg = partitionKeyConfig 10 -- large size so only timeout fires + -- first message at t=0, tick at t=5 (== batchTimeout) fires it + evs = [(0, EvArr 0 "a"), (5, EvTick)] + out = runEvents cfg evs + map (\(info, _) -> (info.trigger, info.size)) out `shouldBe` [(TriggerTimeout, 1)] + + it "does not emit on a tick before the timeout has elapsed" $ do + let cfg = partitionKeyConfig 10 + evs = [(0, EvArr 0 "a"), (4, EvTick)] -- 4s < 5s timeout + -- only the final flush emits + out = runEvents cfg evs + map (\(info, _) -> info.trigger) out `shouldBe` [TriggerFlush] + + it "copies the first message's partition into BatchInfo" $ do + let cfg = partitionKeyConfig 2 + out = runEvents cfg [(0, EvArr 0 "tenant-7"), (0, EvArr 1 "tenant-7")] + map (\(info, _) -> info.partition) out `shouldBe` [Just "tenant-7"] + + describe "message conservation (property)" $ do + it "every input message appears in exactly one emitted batch" $ + property prop_conservation + + it "within each emitted batch, cursors are strictly ascending (per-key FIFO)" $ + property prop_fifo + + it "every TriggerSize batch has size == batchSize" $ + property prop_sizeTrigger + + describe "IO engine runBatcher" $ do + it "conserves messages over a finite stream" $ do + let cfg = partitionKeyConfig 3 + ings = [mkIng i (BatchKey (if even i then "a" else "b")) | i <- [0 .. 19]] + out <- Stream.fold Fold.toList (runBatcher 8 cfg (Stream.fromList ings)) + sort (batchIds out) `shouldBe` sort [ing.envelope.messageId | ing <- ings] + + it "emits full size-3 batches for a single key" $ do + let cfg = partitionKeyConfig 3 + ings = [mkIng i "a" | i <- [0 .. 6]] -- 7 messages + out <- Stream.fold Fold.toList (runBatcher 8 cfg (Stream.fromList ings)) + map (\(info, _) -> (info.trigger, info.size)) out + `shouldBe` [(TriggerSize, 3), (TriggerSize, 3), (TriggerFlush, 1)] + + it "emits a timeout batch on a slow stream before it ends" $ do + -- A stream that yields message 0, then stalls > batchTimeout before 1. + let cfg = + (partitionKeyConfig 100) -- big size => only timeout/flush can fire + { batchTimeout = 0.1, -- 100 ms + tickInterval = Just 0.02 -- scan every 20 ms + } + slow = + Stream.unfoldrM + ( \n -> + if n >= 2 + then pure Nothing + else do + -- delay before the SECOND message so message 0 times out + if n == 1 then threadDelay 250000 else pure () + pure (Just (mkIng n "a", n + 1)) + ) + (0 :: Int) + out <- Stream.fold Fold.toList (runBatcher 8 cfg slow) + any (\(info, _) -> info.trigger == TriggerTimeout) out `shouldBe` True + sort (batchIds out) `shouldBe` [MessageId "m-0", MessageId "m-1"] + +-- QuickCheck generators and properties ------------------------------------- + +-- | A random schedule: arrivals with unique ids and random keys, interspersed +-- with ticks, plus a random batchSize. Time advances 1s per event; batchTimeout +-- is fixed at 5s in partitionKeyConfig, so ticks can fire real timeout batches. +genSchedule :: Gen (Int, [(Double, Ev)]) +genSchedule = do + sz <- choose (1, 8) + n <- choose (0, 60) + keys <- vectorOf n (elements ["a", "b", "c"]) + let arrivals = zipWith EvArr [0 ..] (map BatchKey keys) + -- Interleave ticks randomly. + evs <- interleaveTicks arrivals + let timed = zip (map fromIntegral [0 :: Int ..]) evs + pure (sz, timed) + where + interleaveTicks [] = pure [] + interleaveTicks (a : as) = do + addTick <- arbitrary + rest <- interleaveTicks as + pure (if addTick then a : EvTick : rest else a : rest) + +arrivalIds :: [(Double, Ev)] -> [MessageId] +arrivalIds evs = [MessageId ("m-" <> tshow i) | (_, EvArr i _) <- evs] + +prop_conservation :: Property +prop_conservation = forAll genSchedule $ \(sz, evs) -> + let cfg = partitionKeyConfig sz + out = runEvents cfg evs + in sort (batchIds out) === sort (arrivalIds evs) + +prop_fifo :: Property +prop_fifo = forAll genSchedule $ \(sz, evs) -> + let cfg = partitionKeyConfig sz + out = runEvents cfg evs + in conjoin [strictlyAscending (batchCursors ne) | (_, ne) <- out] + where + strictlyAscending xs = xs === sort xs + +prop_sizeTrigger :: Property +prop_sizeTrigger = forAll genSchedule $ \(sz, evs) -> + let cfg = partitionKeyConfig sz + out = runEvents cfg evs + in conjoin + [ info.size === sz + | (info, _) <- out, + info.trigger == TriggerSize + ] diff --git a/test/Shibuya/Runner/PartitionOrderingSpec.hs b/test/Shibuya/Runner/PartitionOrderingSpec.hs new file mode 100644 --- /dev/null +++ b/test/Shibuya/Runner/PartitionOrderingSpec.hs @@ -0,0 +1,171 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} + +module Shibuya.Runner.PartitionOrderingSpec (spec) where + +import Control.Concurrent.NQE.Supervisor (Strategy (..)) +import Control.Concurrent.STM (atomically, check, readTVar) +import Data.IORef (atomicModifyIORef', newIORef, readIORef) +import Data.List (nub, sort) +import Data.Map.Strict qualified as Map +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as Text +import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Shibuya.Adapter.Mock (getTrackedDecisions, newTrackingAck, trackedListAdapter) +import Shibuya.Core.Ack (AckDecision (..)) +import Shibuya.Core.Ingested (Message (..)) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..)) +import Shibuya.Core.Types qualified as Core +import Shibuya.Internal.Runner.Master (startMaster, stopMaster) +import Shibuya.Internal.Runner.Supervised (SupervisedProcessor (..), runSupervised) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) +import Shibuya.Telemetry.Effect (runTracingNoop) +import Test.Hspec +import Test.QuickCheck +import UnliftIO.Concurrent (threadDelay) + +data Payload = Payload + { delayMicros :: !Int + } + deriving stock (Eq, Show) + +data PartitionCase = PartitionCase + { caseMessages :: ![(Maybe Text, Int)], + caseConcurrency :: !Concurrency + } + deriving stock (Show) + +instance Arbitrary PartitionCase where + arbitrary = do + messageCount <- chooseInt (1, 24) + messages <- + vectorOf messageCount $ + (,) + <$> elements [Nothing, Just "p0", Just "p1", Just "p2"] + <*> chooseInt (0, 2000) + concurrency <- + oneof + [ Ahead <$> chooseInt (2, 8), + Async <$> chooseInt (2, 8) + ] + pure PartitionCase {caseMessages = messages, caseConcurrency = concurrency} + +spec :: Spec +spec = describe "Shibuya.Runner.PartitionOrdering" $ do + it "finalizes each partition in arrival order and exactly once" $ + property $ + withMaxSuccess 30 $ \(PartitionCase messages concurrency) -> + ioProperty $ do + finalized <- runPartitioned messages concurrency + let ids = makeIds (length messages) + partitionById = Map.fromList (zip ids (map fst messages)) + finalizedIds = map fst finalized + partitions = nub [p | Just p <- map fst messages] + finalizedFor p = + [ msgId + | msgId <- finalizedIds, + Map.lookup msgId partitionById == Just (Just p) + ] + arrivedFor p = + [ msgId + | (msgId, Just p') <- zip ids (map fst messages), + p == p' + ] + pure $ + counterexample ("finalized=" <> show finalizedIds) $ + sort finalizedIds === sort ids + .&&. conjoin [finalizedFor p === arrivedFor p | p <- partitions] + + it "respects the global concurrency bound" $ do + maxInFlightRef <- newIORef (0 :: Int) + currentInFlightRef <- newIORef (0 :: Int) + + _ <- + runPartitionedWithHandler + (zip (map Just ["p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"]) (repeat 50_000)) + (Async 3) + ( \_ -> do + cur <- + liftIO $ + atomicModifyIORef' currentInFlightRef (\n -> let n' = n + 1 in (n', n')) + liftIO $ atomicModifyIORef' maxInFlightRef (\m -> (max m cur, ())) + liftIO $ threadDelay 50_000 + liftIO $ atomicModifyIORef' currentInFlightRef (\n -> (n - 1, ())) + pure AckOk + ) + + maxInFlight <- readIORef maxInFlightRef + maxInFlight `shouldSatisfy` (<= 3) + + it "does not let a slow partition block fast partitions" $ do + finalized <- + runPartitioned + [ (Just "slow", 80_000), + (Just "fast", 1_000), + (Just "slow", 80_000), + (Just "fast", 1_000), + (Just "slow", 80_000), + (Just "fast", 1_000) + ] + (Async 4) + + let orderedIds = map fst finalized + lastSlowIndex = fromMaybe maxBound $ elemIndex' "msg-4" orderedIds + fastIndexes = + map + (`elemIndexOrMax` orderedIds) + ["msg-1", "msg-3", "msg-5"] + all (< lastSlowIndex) fastIndexes `shouldBe` True + +runPartitioned :: + [(Maybe Text, Int)] -> + Concurrency -> + IO [(MessageId, AckDecision)] +runPartitioned messages concurrency = + runPartitionedWithHandler messages concurrency $ \ingested -> do + let Message {envelope = Envelope {payload = Payload delay}} = ingested + liftIO $ threadDelay delay + pure AckOk + +runPartitionedWithHandler :: + [(Maybe Text, Int)] -> + Concurrency -> + (forall es. (IOE :> es) => Message es Payload -> Eff es AckDecision) -> + IO [(MessageId, AckDecision)] +runPartitionedWithHandler messages concurrency handler = + runEff $ runTracingNoop $ do + tracking <- newTrackingAck + let envs = zipWith mkEnvelope [0 :: Int ..] messages + adapter = trackedListAdapter tracking envs + master <- startMaster IgnoreAll + sp <- runSupervised master 50 (ProcessorId "partitioned") PartitionedInOrder concurrency adapter handler + liftIO $ atomically $ readTVar sp.done >>= check + decisions <- getTrackedDecisions tracking + stopMaster master + pure (reverse decisions) + +mkEnvelope :: Int -> (Maybe Text, Int) -> Envelope Payload +mkEnvelope i (partition, delayMicros) = + (Core.mkEnvelope (makeId i) Payload {delayMicros}) + { cursor = Just (CursorInt i), + partition = partition + } + +makeIds :: Int -> [MessageId] +makeIds n = map makeId [0 .. n - 1] + +makeId :: Int -> MessageId +makeId i = MessageId ("msg-" <> Text.pack (show i)) + +elemIndexOrMax :: MessageId -> [MessageId] -> Int +elemIndexOrMax needle haystack = fromMaybe maxBound (elemIndex' needle haystack) + +elemIndex' :: (Eq a) => a -> [a] -> Maybe Int +elemIndex' needle = go 0 + where + go _ [] = Nothing + go i (x : xs) + | x == needle = Just i + | otherwise = go (i + 1) xs diff --git a/test/Shibuya/Runner/SupervisedSpec.hs b/test/Shibuya/Runner/SupervisedSpec.hs --- a/test/Shibuya/Runner/SupervisedSpec.hs +++ b/test/Shibuya/Runner/SupervisedSpec.hs @@ -3,47 +3,58 @@ module Shibuya.Runner.SupervisedSpec (spec) where import Control.Concurrent.NQE.Supervisor (Strategy (..)) -import Control.Concurrent.STM (readTVarIO) -import Control.Monad (forM, replicateM) -import Data.HashMap.Strict qualified as HashMap +import Control.Concurrent.STM (atomically, check, readTVar, readTVarIO) +import Control.Monad (forM, forM_, replicateM) import Data.IORef (IORef, atomicModifyIORef', atomicWriteIORef, modifyIORef', newIORef, readIORef) +import Data.Map.Strict qualified as Map import Data.Text qualified as Text import Data.Time (UTCTime (..), fromGregorian) import Effectful (Eff, IOE, liftIO, runEff, (:>)) +import Shibuya (ProcessorHalt (..)) import Shibuya.Adapter (Adapter (..)) +import Shibuya.Adapter.Mock + ( getTrackedDecisions, + mkTrackedIngested, + newTrackingAck, + trackedListAdapter, + ) import Shibuya.App - ( ShutdownConfig (..), - SupervisionStrategy (..), + ( AppConfig (..), + ShutdownConfig (..), + defaultAppConfig, mkProcessor, runApp, stopAppGracefully, ) -import Shibuya.Core (ProcessorHalt (..)) -import Shibuya.Core.Ack (AckDecision (..), HaltReason (..)) +import Shibuya.Batch.TestHarness (finalizedExactlyOnce) +import Shibuya.Core.Ack (AckDecision (..), HaltReason (..), RetryDelay (..)) import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..)) -import Shibuya.Handler (Handler) -import Shibuya.Policy (Concurrency (..)) -import Shibuya.Runner.Master - ( getAllMetrics, - startMaster, - stopMaster, - ) -import Shibuya.Runner.Metrics +import Shibuya.Core.Ingested (Ingested, Message (..), mkIngested) +import Shibuya.Core.Metrics ( InFlightInfo (..), ProcessorId (..), ProcessorMetrics (..), ProcessorState (..), StreamStats (..), + sampleMetrics, ) -import Shibuya.Runner.Supervised +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Handler (Handler) +import Shibuya.Internal.Runner.Master + ( getAllMetrics, + getAllMetricsIO, + getProcessorMetricsIO, + startMaster, + stopMaster, + ) +import Shibuya.Internal.Runner.Supervised ( SupervisedProcessor (..), getMetrics, isDone, runSupervised, runWithMetrics, ) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) import Shibuya.Telemetry.Effect (runTracingNoop) import Streamly.Data.Stream qualified as Stream import Test.Hspec @@ -53,7 +64,7 @@ spec :: Spec spec = do - describe "Shibuya.Runner.Master" $ do + describe "Shibuya.Internal.Runner.Master" $ do it "starts and stops cleanly" $ do result <- runEff $ runTracingNoop $ do master <- startMaster IgnoreAll @@ -69,7 +80,17 @@ pure m metrics `shouldSatisfy` null - describe "Shibuya.Runner.Supervised" $ do + it "metrics queries return after stopMaster" $ do + master <- runEff $ runTracingNoop $ startMaster IgnoreAll + runEff $ runTracingNoop $ stopMaster master + + allMetrics <- UIO.timeout 1_000_000 (getAllMetricsIO master) + processorMetrics <- UIO.timeout 1_000_000 (getProcessorMetricsIO master (ProcessorId "x")) + + allMetrics `shouldSatisfy` maybe False null + processorMetrics `shouldBe` Just Nothing + + describe "Shibuya.Internal.Runner.Supervised" $ do describe "runWithMetrics" $ do it "processes messages and tracks metrics" $ do (finalMetrics, processedMsgs) <- runEff $ runTracingNoop $ do @@ -128,6 +149,33 @@ done `shouldBe` True + it "completes when the stream is longer than the inbox" $ do + processedRef <- newIORef (0 :: Int) + + result <- + UIO.timeout 10_000_000 $ + runEff $ + runTracingNoop $ do + messages <- createTestMessages 100 + + let handler _ = do + liftIO $ modifyIORef' processedRef (+ 1) + pure AckOk + adapter = testAdapter messages + procId = ProcessorId "longer-than-inbox" + + sp <- runWithMetrics 5 procId adapter handler + metrics <- getMetrics sp + processed <- liftIO $ readIORef processedRef + pure (metrics, processed) + + case result of + Nothing -> expectationFailure "runWithMetrics timed out" + Just (finalMetrics, processed) -> do + processed `shouldBe` 100 + finalMetrics.stats.received `shouldBe` 100 + finalMetrics.stats.processed `shouldBe` 100 + describe "AckHalt behavior" $ do it "stops processing when handler returns AckHalt" $ do processedCount <- do @@ -169,11 +217,11 @@ master <- startMaster IgnoreAll -- runSupervised catches ProcessorHalt, so we can check metrics - sp <- runSupervised master 10 (ProcessorId "halt-metrics") Serial adapter handler + sp <- runSupervised master 10 (ProcessorId "halt-metrics") Unordered Serial adapter handler -- Wait for processor to halt liftIO $ threadDelay 100000 -- 100ms - m <- liftIO $ readTVarIO sp.metrics + m <- liftIO $ sampleMetrics sp.metrics stopMaster master pure m @@ -207,8 +255,8 @@ master <- startMaster IgnoreAll - _spA <- runSupervised master 10 (ProcessorId "A") Serial adapterA handlerA - _spB <- runSupervised master 10 (ProcessorId "B") Serial adapterB handlerB + _spA <- runSupervised master 10 (ProcessorId "A") Unordered Serial adapterA handlerA + _spB <- runSupervised master 10 (ProcessorId "B") Unordered Serial adapterB handlerB -- Wait for both to complete liftIO $ threadDelay 500000 -- 500ms @@ -238,7 +286,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 10 (ProcessorId "ahead") (Ahead 3) adapter handler + _ <- runSupervised master 10 (ProcessorId "ahead") Unordered (Ahead 3) adapter handler -- Wait for completion liftIO $ threadDelay 300000 -- 300ms @@ -271,7 +319,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages -- Use Ahead with max 3 concurrent - _sp <- runSupervised master 10 (ProcessorId "ahead-limit") (Ahead 3) adapter handler + _sp <- runSupervised master 10 (ProcessorId "ahead-limit") Unordered (Ahead 3) adapter handler liftIO $ threadDelay 1000000 -- 1s to let it complete stopMaster master @@ -294,7 +342,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 10 (ProcessorId "async") (Async 5) adapter handler + _ <- runSupervised master 10 (ProcessorId "async") Unordered (Async 5) adapter handler -- Wait for completion liftIO $ threadDelay 300000 -- 300ms @@ -321,7 +369,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _sp <- runSupervised master 10 (ProcessorId "async-limit") (Async 3) adapter handler + _sp <- runSupervised master 10 (ProcessorId "async-limit") Unordered (Async 3) adapter handler liftIO $ threadDelay 1000000 -- 1s stopMaster master @@ -350,7 +398,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _sp <- runSupervised master 10 (ProcessorId "halt-concurrent") (Async 3) adapter handler + _sp <- runSupervised master 10 (ProcessorId "halt-concurrent") Unordered (Async 3) adapter handler liftIO $ threadDelay 500000 -- 500ms stopMaster master @@ -378,7 +426,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _sp <- runSupervised master 10 (ProcessorId "halt-stop-read") (Async 3) adapter handler + _sp <- runSupervised master 10 (ProcessorId "halt-stop-read") Unordered (Async 3) adapter handler liftIO $ threadDelay 500000 -- 500ms stopMaster master @@ -411,7 +459,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 10 (ProcessorId "error-handling") (Async 3) adapter handler + _ <- runSupervised master 10 (ProcessorId "error-handling") Unordered (Async 3) adapter handler liftIO $ threadDelay 500000 -- 500ms stopMaster master @@ -424,6 +472,87 @@ let successes = [n | Right n <- results] length successes `shouldSatisfy` (>= 3) + it "finalizes throwing handlers with immediate retry" $ do + (metrics, tracked) <- runEff $ runTracingNoop $ do + tracking <- newTrackingAck + let envelopes = createTestEnvelopes 5 + adapter = trackedListAdapter tracking envelopes + throwingIds = [MessageId "msg-2", MessageId "msg-4"] + handler ingested = + if ingested.envelope.messageId `elem` throwingIds + then error "handler failed for conservation test" + else pure AckOk + + sp <- runWithMetrics 10 (ProcessorId "handler-exception-finalize") adapter handler + decisions <- getTrackedDecisions tracking + finalMetrics <- getMetrics sp + pure (finalMetrics, decisions) + + let expected = + Map.fromList + [ (MessageId "msg-1", AckOk), + (MessageId "msg-2", AckRetry (RetryDelay 0)), + (MessageId "msg-3", AckOk), + (MessageId "msg-4", AckRetry (RetryDelay 0)), + (MessageId "msg-5", AckOk) + ] + finalizedExactlyOnce tracked expected `shouldBe` Right () + metrics.stats.processed `shouldBe` 3 + metrics.stats.failed `shouldBe` 2 + + it "retries transient finalizer failures on the single-message path" $ do + (attempts, tracked) <- runEff $ runTracingNoop $ do + attemptsRef <- liftIO $ newIORef (0 :: Int) + decisionsRef <- liftIO $ newIORef ([] :: [(MessageId, AckDecision)]) + let env = createTestEnvelope 1 + ingested = + mkIngested + env + ( AckHandle $ \decision -> + liftIO $ do + attempt <- + atomicModifyIORef' + attemptsRef + (\n -> let n' = n + 1 in (n', n')) + if attempt <= 2 + then ioError (userError "transient finalizer failure") + else modifyIORef' decisionsRef ((env.messageId, decision) :) + ) + adapter = testAdapter [ingested] + + _sp <- runWithMetrics 10 (ProcessorId "transient-finalizer") adapter alwaysAckOk + attemptCount <- liftIO $ readIORef attemptsRef + decisions <- liftIO $ readIORef decisionsRef + pure (attemptCount, decisions) + + attempts `shouldBe` 3 + finalizedExactlyOnce tracked (Map.singleton (MessageId "msg-1") AckOk) `shouldBe` Right () + + it "halts loudly when finalizer retry is exhausted" $ do + (metrics, done, tracked) <- runEff $ runTracingNoop $ do + tracking <- newTrackingAck + let env1 = createTestEnvelope 1 + env2 = createTestEnvelope 2 + failing = + mkIngested env1 (AckHandle $ \_ -> liftIO $ ioError (userError "permanent finalizer failure")) + trackedSecond = mkTrackedIngested tracking env2 + adapter = testAdapter [failing, trackedSecond] + + master <- startMaster IgnoreAll + sp <- runSupervised master 10 (ProcessorId "permanent-finalizer") Unordered (Async 2) adapter alwaysAckOk + liftIO $ threadDelay 700000 + finalMetrics <- liftIO $ sampleMetrics sp.metrics + doneState <- liftIO $ readTVarIO sp.done + decisions <- getTrackedDecisions tracking + stopMaster master + pure (finalMetrics, doneState, decisions) + + done `shouldBe` True + case metrics.state of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "msg-1" + other -> expectationFailure $ "Expected Failed state naming msg-1, got: " ++ show other + finalizedExactlyOnce tracked (Map.singleton (MessageId "msg-2") AckOk) `shouldBe` Right () + describe "Metrics tracking" $ do it "tracks in-flight count during concurrent processing" $ do maxInFlightObserved <- newIORef (0 :: Int) @@ -437,12 +566,12 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - sp <- runSupervised master 10 (ProcessorId "metrics-test") (Async 3) adapter handler + sp <- runSupervised master 10 (ProcessorId "metrics-test") Unordered (Async 3) adapter handler -- Check metrics while processing liftIO $ do threadDelay 50000 -- 50ms - should be in the middle of processing - metrics <- readTVarIO sp.metrics + metrics <- sampleMetrics sp.metrics case metrics.state of Processing info _ -> modifyIORef' maxInFlightObserved (max info.inFlight) _ -> pure () @@ -464,12 +593,12 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - sp <- runSupervised master 10 (ProcessorId "max-conc") (Ahead 7) adapter handler + sp <- runSupervised master 10 (ProcessorId "max-conc") Unordered (Ahead 7) adapter handler -- Check metrics while processing result <- liftIO $ do threadDelay 25000 -- 25ms - metrics <- readTVarIO sp.metrics + metrics <- sampleMetrics sp.metrics case metrics.state of Processing info _ -> pure $ Just info.maxConcurrency _ -> pure Nothing @@ -503,7 +632,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 10 (ProcessorId "ahead-order") (Ahead 5) adapter handler + _ <- runSupervised master 10 (ProcessorId "ahead-order") Unordered (Ahead 5) adapter handler liftIO $ threadDelay 500000 -- 500ms stopMaster master @@ -513,11 +642,11 @@ length completeOrder `shouldBe` 5 -- The key property: all messages were processed - -- (Ordering of side effects may vary, but stream output is ordered) + -- (OrderingPolicy of side effects may vary, but stream output is ordered) describe "Robustness" $ do describe "Adapter source exceptions" $ do - it "drains already-ingested messages, marks Failed, and propagates the ingester exception" $ do + it "drains already-ingested messages, marks Failed, and propagates the ingester exception under IgnoreGraceful" $ do processedRef <- newIORef (0 :: Int) spRef <- newIORef (Nothing :: Maybe SupervisedProcessor) @@ -545,8 +674,8 @@ liftIO $ modifyIORef' processedRef (+ 1) pure AckOk - master <- startMaster IgnoreAll - sp <- runSupervised master 10 (ProcessorId "failing-adapter") Serial adapter handler + master <- startMaster IgnoreGraceful + sp <- runSupervised master 10 (ProcessorId "failing-adapter") Unordered Serial adapter handler liftIO $ atomicWriteIORef spRef (Just sp) -- Wait for the linked ingester failure to reach this thread. @@ -571,11 +700,98 @@ Just sp -> do done <- readTVarIO sp.done done `shouldBe` True - metrics <- readTVarIO sp.metrics + metrics <- sampleMetrics sp.metrics case metrics.state of Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "Network failure mid-stream" other -> expectationFailure $ "Expected Failed state, got: " ++ show other + it "drains already-ingested messages and marks Failed without propagating under IgnoreAll" $ do + processedRef <- newIORef (0 :: Int) + spRef <- newIORef (Nothing :: Maybe SupervisedProcessor) + + result <- + UIO.withAsync + ( runEff $ runTracingNoop $ do + let failingSource = Stream.unfoldrM step (0 :: Int) + where + step n + | n < 3 = do + msg <- liftIO $ createSingleMessage (n + 1) + pure $ Just (msg, n + 1) + | n == 3 = error "Network failure mid-stream" + | otherwise = pure Nothing + + let adapter = + Adapter + { adapterName = "test:failing-source", + source = failingSource, + shutdown = pure () + } + + let handler _ = do + liftIO $ modifyIORef' processedRef (+ 1) + pure AckOk + + master <- startMaster IgnoreAll + sp <- runSupervised master 10 (ProcessorId "failing-adapter-ignore") Unordered Serial adapter handler + liftIO $ atomicWriteIORef spRef (Just sp) + + liftIO $ threadDelay 200000 -- 200ms + stopMaster master + ) + UIO.waitCatch + + case result of + Left err -> + expectationFailure $ "Expected IgnoreAll to isolate the ingester exception, got: " ++ show err + Right () -> + pure () + + processed <- readIORef processedRef + processed `shouldBe` 3 + + mSp <- readIORef spRef + case mSp of + Nothing -> expectationFailure "Expected supervised processor handle" + Just sp -> do + done <- readTVarIO sp.done + done `shouldBe` True + metrics <- sampleMetrics sp.metrics + case metrics.state of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "Network failure mid-stream" + other -> expectationFailure $ "Expected Failed state, got: " ++ show other + + it "never drops an ingester failure (repeated)" $ do + results <- forM [(1 :: Int) .. 25] $ \i -> + runEff $ runTracingNoop $ do + let failingSource = Stream.unfoldrM (\() -> error "boom immediately") () + adapter = + Adapter + { adapterName = "test:failing-immediate", + source = failingSource, + shutdown = pure () + } + handler _ = pure AckOk + + master <- startMaster IgnoreAll + sp <- runSupervised master 10 (ProcessorId $ "failing-immediate-" <> Text.pack (show i)) Unordered Serial adapter handler + + doneResult <- + liftIO $ + UIO.timeout 2_000_000 $ + atomically $ do + done <- readTVar sp.done + check done + metrics <- liftIO $ sampleMetrics sp.metrics + stopMaster master + pure (i, doneResult, metrics) + + forM_ results $ \(i, doneResult, metrics) -> do + doneResult `shouldBe` Just () + case metrics.state of + Failed msg _ -> msg `shouldSatisfy` Text.isInfixOf "boom immediately" + other -> expectationFailure $ "Iteration " ++ show i ++ " expected Failed state, got: " ++ show other + describe "Rapid start/stop cycles" $ do it "handles rapid start/stop without resource leaks" $ do -- Run 50 rapid cycles @@ -586,7 +802,7 @@ let adapter = testAdapter messages handler _ = pure AckOk - sp <- runSupervised master 10 (ProcessorId "rapid") Serial adapter handler + sp <- runSupervised master 10 (ProcessorId "rapid") Unordered Serial adapter handler -- Minimal wait liftIO $ threadDelay 10000 -- 10ms @@ -613,7 +829,7 @@ liftIO $ atomicModifyIORef' countRef (\n -> (n + 1, ())) liftIO $ threadDelay 5000 -- 5ms pure AckOk - runSupervised master 10 (ProcessorId $ "concurrent-" <> Text.pack (show i)) Serial adapter handler + runSupervised master 10 (ProcessorId $ "concurrent-" <> Text.pack (show i)) Unordered Serial adapter handler -- Wait briefly then stop liftIO $ threadDelay 100000 -- 100ms @@ -623,8 +839,8 @@ count <- readIORef countRef count `shouldSatisfy` (> 0) - describe "KillAll supervision strategy" $ do - it "stops all processors when adapter source fails with KillAll strategy" $ do + describe "IgnoreGraceful supervision strategy (StopAllOnFailure)" $ do + it "stops all processors when adapter source fails with IgnoreGraceful strategy" $ do countARef <- newIORef (0 :: Int) countBRef <- newIORef (0 :: Int) @@ -663,11 +879,11 @@ let adapterB = testAdapter messagesB - -- Use KillAll strategy - master <- startMaster KillAll + -- Use the NQE strategy backing StopAllOnFailure. + master <- startMaster IgnoreGraceful - _spA <- runSupervised master 10 (ProcessorId "killall-A") Serial adapterA handlerA - _spB <- runSupervised master 10 (ProcessorId "killall-B") Serial adapterB handlerB + _spA <- runSupervised master 10 (ProcessorId "killall-A") Unordered Serial adapterA handlerA + _spB <- runSupervised master 10 (ProcessorId "killall-B") Unordered Serial adapterB handlerB -- Wait for A to fail and trigger KillAll liftIO $ threadDelay 500000 -- 500ms @@ -712,7 +928,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 10 (ProcessorId "multi-halt") (Async 5) adapter handler + _ <- runSupervised master 10 (ProcessorId "multi-halt") Unordered (Async 5) adapter handler liftIO $ threadDelay 500000 -- 500ms stopMaster master @@ -739,7 +955,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _sp <- runSupervised master 100 (ProcessorId "load-test") (Async 10) adapter handler + _sp <- runSupervised master 100 (ProcessorId "load-test") Unordered (Async 10) adapter handler -- Wait for completion liftIO $ threadDelay 2000000 -- 2s @@ -770,7 +986,7 @@ master <- startMaster IgnoreAll let adapter = testAdapter messages - _ <- runSupervised master 50 (ProcessorId "mixed-load") (Async 5) adapter handler + _ <- runSupervised master 50 (ProcessorId "mixed-load") Unordered (Async 5) adapter handler liftIO $ threadDelay 1000000 -- 1s stopMaster master @@ -800,7 +1016,7 @@ let adapter = testAdapter messages processor = mkProcessor adapter handler - result <- runApp IgnoreFailures 10 [(ProcessorId "drain-test", processor)] + result <- runApp defaultAppConfig {inboxSize = 10} [(ProcessorId "drain-test", processor)] case result of Left _ -> pure False Right appHandle -> do @@ -832,7 +1048,7 @@ let adapter = testAdapter messages processor = mkProcessor adapter handler - result <- runApp IgnoreFailures 5 [(ProcessorId "timeout-test", processor)] + result <- runApp defaultAppConfig {inboxSize = 5} [(ProcessorId "timeout-test", processor)] case result of Left _ -> pure True -- Treat error as "drained" for simplicity Right appHandle -> do @@ -858,7 +1074,7 @@ let adapter = testAdapter messages processor = mkProcessor adapter handler - result <- runApp IgnoreFailures 10 [(ProcessorId "quick-test", processor)] + result <- runApp defaultAppConfig {inboxSize = 10} [(ProcessorId "quick-test", processor)] case result of Left _ -> pure False Right appHandle -> do @@ -874,55 +1090,31 @@ testTime :: UTCTime testTime = UTCTime (fromGregorian 2026 1 1) 0 +createTestEnvelope :: Int -> Envelope String +createTestEnvelope i = + (mkEnvelope (MessageId $ "msg-" <> Text.pack (show i)) ("message-" <> show i)) + { cursor = Just (CursorInt i), + enqueuedAt = Just testTime + } + +createTestEnvelopes :: Int -> [Envelope String] +createTestEnvelopes n = map createTestEnvelope [1 .. n] + createTestMessages :: (IOE :> es) => Int -> Eff es [Ingested es String] createTestMessages n = mapM createMessage [1 .. n] where createMessage i = do - let msgId = MessageId $ "msg-" <> Text.pack (show i) - env = - Envelope - { messageId = msgId, - cursor = Just (CursorInt i), - partition = Nothing, - enqueuedAt = Just testTime, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = "message-" <> show i - } + let env = createTestEnvelope i ackHandle = AckHandle $ \_ -> pure () - pure $ - Ingested - { envelope = env, - ack = ackHandle, - lease = Nothing - } + pure $ mkIngested env ackHandle -- | Create a single message (for use in streaming contexts) -- Polymorphic over effect stack for use with tracing. createSingleMessage :: (IOE :> es) => Int -> IO (Ingested es String) createSingleMessage i = do - let msgId = MessageId $ "msg-" <> Text.pack (show i) - env = - Envelope - { messageId = msgId, - cursor = Just (CursorInt i), - partition = Nothing, - enqueuedAt = Just testTime, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = "message-" <> show i - } + let env = createTestEnvelope i ackHandle = AckHandle $ \_ -> pure () - pure $ - Ingested - { envelope = env, - ack = ackHandle, - lease = Nothing - } + pure $ mkIngested env ackHandle testAdapter :: [Ingested es String] -> Adapter es String testAdapter messages = diff --git a/test/Shibuya/RunnerSpec.hs b/test/Shibuya/RunnerSpec.hs --- a/test/Shibuya/RunnerSpec.hs +++ b/test/Shibuya/RunnerSpec.hs @@ -2,30 +2,50 @@ module Shibuya.RunnerSpec (spec) where -import Data.HashMap.Strict qualified as HashMap import Data.IORef (IORef, modifyIORef', newIORef, readIORef) import Data.Text qualified as Text import Data.Time (UTCTime (..), fromGregorian) import Effectful (Eff, IOE, liftIO, runEff, (:>)) import Shibuya.Adapter (Adapter (..)) import Shibuya.Adapter.Mock (TrackingAck (..), newTrackingAck, trackingAckHandle) -import Shibuya.App (AppError (..), QueueProcessor (..), SupervisionStrategy (..), mkProcessor, runApp, waitApp) +import Shibuya.App (AppConfig (..), AppError (..), QueueProcessor (..), defaultAppConfig, mkProcessor, runApp, waitApp) import Shibuya.Core.Ack (AckDecision (..)) import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Error (PolicyError (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..)) +import Shibuya.Core.Error (ConfigError (..), PolicyError (..)) +import Shibuya.Core.Ingested (Ingested, Message (..), mkIngested) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..), mkEnvelope) import Shibuya.Handler (Handler) -import Shibuya.Policy (Concurrency (..), Ordering (..)) -import Shibuya.Runner.Metrics (ProcessorId (..)) +import Shibuya.Policy (Concurrency (..), OrderingPolicy (..)) import Shibuya.Telemetry.Effect (runTracingNoop) import Streamly.Data.Stream qualified as Stream import Test.Hspec -import Prelude hiding (Ordering) spec :: Spec spec = do describe "runApp" $ do + it "rejects inboxSize 0 before starting processors" $ do + result <- runEff $ runTracingNoop $ do + messages <- createTestMessages 1 + let processor = mkProcessor (testAdapter messages) alwaysAckOk + runApp defaultAppConfig {inboxSize = 0} [(ProcessorId "invalid-config", processor)] + + case result of + Left (AppConfigInvalid (InvalidInboxSize 0)) -> pure () + Left err -> expectationFailure $ "Expected AppConfigInvalid, got: " ++ show err + Right _ -> expectationFailure "Expected config validation to fail" + + it "rejects negative inboxSize before Natural conversion" $ do + result <- runEff $ runTracingNoop $ do + messages <- createTestMessages 1 + let processor = mkProcessor (testAdapter messages) alwaysAckOk + runApp defaultAppConfig {inboxSize = -5} [(ProcessorId "invalid-config", processor)] + + case result of + Left (AppConfigInvalid (InvalidInboxSize (-5))) -> pure () + Left err -> expectationFailure $ "Expected AppConfigInvalid, got: " ++ show err + Right _ -> expectationFailure "Expected config validation to fail" + it "processes messages from mock adapter" $ do result <- runEff $ runTracingNoop $ do -- Track processed messages @@ -42,8 +62,7 @@ -- Run the app res <- runApp - IgnoreFailures - 100 + defaultAppConfig [ (ProcessorId "test", processor) ] @@ -72,8 +91,7 @@ -- Run the app res <- runApp - IgnoreFailures - 100 + defaultAppConfig [ (ProcessorId "test", processor) ] @@ -105,8 +123,7 @@ res <- runApp - IgnoreFailures - 100 + defaultAppConfig [ (ProcessorId "proc1", proc1), (ProcessorId "proc2", proc2) ] @@ -128,7 +145,7 @@ -- Invalid combination: StrictInOrder requires Serial processor = QueueProcessor adapter handler StrictInOrder (Async 5) - runApp IgnoreFailures 100 [(ProcessorId "invalid", processor)] + runApp defaultAppConfig [(ProcessorId "invalid", processor)] case result of Left (AppPolicyError (InvalidPolicyCombo _)) -> pure () @@ -142,7 +159,7 @@ handler = alwaysAckOk processor = QueueProcessor adapter handler StrictInOrder (Ahead 5) - runApp IgnoreFailures 100 [(ProcessorId "invalid", processor)] + runApp defaultAppConfig [(ProcessorId "invalid", processor)] case result of Left (AppPolicyError (InvalidPolicyCombo _)) -> pure () @@ -160,8 +177,7 @@ res <- runApp - IgnoreFailures - 100 + defaultAppConfig [ (ProcessorId "async", proc1), (ProcessorId "ahead", proc2) ] @@ -178,14 +194,18 @@ messages <- runEff $ createTestMessages 1 let adapter = testAdapter messages handler = alwaysAckOk - QueueProcessor _ _ ordering _ = mkProcessor adapter handler + ordering = case mkProcessor adapter handler of + QueueProcessor {ordering = o} -> o + BatchingProcessor {ordering = o} -> o ordering `shouldBe` Unordered it "creates processor with Serial concurrency" $ do messages <- runEff $ createTestMessages 1 let adapter = testAdapter messages handler = alwaysAckOk - QueueProcessor _ _ _ concurrency = mkProcessor adapter handler + concurrency = case mkProcessor adapter handler of + QueueProcessor {concurrency = c} -> c + BatchingProcessor {concurrency = c} -> c concurrency `shouldBe` Serial -- Test helpers @@ -200,24 +220,12 @@ createMessage i = do let msgId = MessageId $ "msg-" <> (if i < 10 then "0" else "") <> Text.pack (show i) env = - Envelope - { messageId = msgId, - cursor = Just (CursorInt i), - partition = Nothing, - enqueuedAt = Just testTime, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = "message-" <> show i + (mkEnvelope msgId ("message-" <> show i)) + { cursor = Just (CursorInt i), + enqueuedAt = Just testTime } ackHandle = AckHandle $ \_ -> pure () -- No-op ack - pure $ - Ingested - { envelope = env, - ack = ackHandle, - lease = Nothing - } + pure $ mkIngested env ackHandle -- | Create N test messages with tracking acks createTrackedMessages :: (IOE :> es) => TrackingAck -> Int -> Eff es [Ingested es String] @@ -226,24 +234,12 @@ createMessage i = do let msgId = MessageId $ "msg-" <> (if i < 10 then "0" else "") <> Text.pack (show i) env = - Envelope - { messageId = msgId, - cursor = Just (CursorInt i), - partition = Nothing, - enqueuedAt = Just testTime, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = "message-" <> show i + (mkEnvelope msgId ("message-" <> show i)) + { cursor = Just (CursorInt i), + enqueuedAt = Just testTime } ackHandle = trackingAckHandle tracking msgId - pure $ - Ingested - { envelope = env, - ack = ackHandle, - lease = Nothing - } + pure $ mkIngested env ackHandle -- | Create a test adapter from a list of messages testAdapter :: [Ingested es String] -> Adapter es String diff --git a/test/Shibuya/Telemetry/SemanticSpec.hs b/test/Shibuya/Telemetry/SemanticSpec.hs --- a/test/Shibuya/Telemetry/SemanticSpec.hs +++ b/test/Shibuya/Telemetry/SemanticSpec.hs @@ -42,10 +42,10 @@ import Shibuya.Adapter.Mock (listAdapter) import Shibuya.Core.Ack (AckDecision (..)) import Shibuya.Core.AckHandle (AckHandle (..)) -import Shibuya.Core.Ingested (Ingested (..)) -import Shibuya.Core.Types (Envelope (..), MessageId (..)) -import Shibuya.Runner.Metrics (ProcessorId (..)) -import Shibuya.Runner.Supervised (runWithMetrics) +import Shibuya.Core.Ingested (mkIngested) +import Shibuya.Core.Metrics (ProcessorId (..)) +import Shibuya.Core.Types (Envelope (..), MessageId (..), mkEnvelope) +import Shibuya.Internal.Runner.Supervised (runWithMetrics) import Shibuya.Telemetry.Effect (runTracing) import Test.Hspec @@ -58,23 +58,8 @@ runEff $ runTracing tracer $ do let envelope = - Envelope - { messageId = MessageId "m-1", - cursor = Nothing, - partition = Nothing, - enqueuedAt = Nothing, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = HashMap.empty, - payload = ("hello" :: Text) - } - ingested = - Ingested - { envelope = envelope, - ack = AckHandle (\_ -> pure ()), - lease = Nothing - } + mkEnvelope (MessageId "m-1") ("hello" :: Text) + ingested = mkIngested envelope (AckHandle (\_ -> pure ())) adapter = listAdapter [ingested] handler _ = pure AckOk procId = ProcessorId "test-proc" @@ -113,15 +98,8 @@ runEff $ runTracing tracer $ do let envelope = - Envelope - { messageId = MessageId "orders-2-42", - cursor = Nothing, - partition = Nothing, - enqueuedAt = Nothing, - traceContext = Nothing, - headers = Nothing, - attempt = Nothing, - attributes = + (mkEnvelope (MessageId "orders-2-42") ("hello" :: Text)) + { attributes = HashMap.fromList [ ("messaging.system", toAttribute ("kafka" :: Text)), ( "messaging.kafka.destination.partition", @@ -130,15 +108,9 @@ ( "messaging.kafka.message.offset", toAttribute (42 :: Int64) ) - ], - payload = ("hello" :: Text) - } - ingested = - Ingested - { envelope = envelope, - ack = AckHandle (\_ -> pure ()), - lease = Nothing + ] } + ingested = mkIngested envelope (AckHandle (\_ -> pure ())) adapter = listAdapter [ingested] handler _ = pure AckOk procId = ProcessorId "orders-consumer"