Subscription.hs

module Kiroku.Store.Subscription (
    -- * Subscribe
    subscribe,
    withSubscription,

    -- * Observability
    subscriptionStates,
    SubscriptionStateView (..),

    -- * Types
    module Kiroku.Store.Subscription.Types,
) where

import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM (atomically, modifyTVar', newTVarIO, readTVarIO)
import Control.Exception (bracket, bracketOnError, finally, mask, throwIO)
import Control.Lens ((^.))
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.IO.Unlift (MonadUnliftIO, withRunInIO)
import Data.Foldable (for_)
import Data.Generics.Labels ()
import Data.Int (Int32)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Unique (newUnique)
import GHC.Generics (Generic)
import Kiroku.Store.Connection (KirokuStore (..))
import Kiroku.Store.Notification qualified as Notifier
import Kiroku.Store.Subscription.EventPublisher qualified as Pub
import Kiroku.Store.Subscription.Fsm (SubscriptionState (..), stateCursor, stateName)
import Kiroku.Store.Subscription.Types
import Kiroku.Store.Subscription.Worker (LiveSource (..), configMember, runWorker)
import Kiroku.Store.Types (CategoryName (..), GlobalPosition (..))

{- | Start a subscription. Returns a handle for cancellation and waiting.

The subscription spawns a worker thread that:

1. Reads the checkpoint from the database (or starts from global position 0
   for a fresh subscription name). A database error while loading the checkpoint
   fails the worker loudly through 'wait'; it does not fall back to 0.
2. Catches up by querying the database directly until it reaches the
   'Kiroku.Store.Subscription.EventPublisher.lastPublished' cursor.
3. Switches to live mode. For 'Kiroku.Store.Subscription.Types.AllStreams'
   subscriptions outside a consumer group, the worker reads pre-broadcast
   events from the publisher's bounded per-subscriber queue. For
   'Kiroku.Store.Subscription.Types.Category' subscriptions and all
   consumer-group subscriptions, no publisher queue is created; the worker
   re-queries the database in live mode.

The returned 'SubscriptionHandle' carries an implicit lifecycle contract:
the worker thread runs until 'cancel' is called or the handler returns 'Stop'.
Forgetting to cancel leaks the thread. Prefer 'withSubscription' for any
non-trivial code path; use the bare 'subscribe' only when the caller already
has a structured lifecycle (e.g., the Streamly 'subscriptionStream' bridge).

=== Delivery semantics

Events are delivered __at least once__. The store's checkpoint is advanced
__per batch__, not per event: when the handler returns 'Continue' for every
event in a batch the checkpoint is saved at the batch tail; when the handler
returns 'Stop' for some event the checkpoint is saved at that event. The
following events on the boundary therefore __replay__ on the next
subscription with the same 'Kiroku.Store.Subscription.Types.SubscriptionName':

* The handler returned 'Continue' but the worker was cancelled or the
  process crashed before 'saveCheckpoint' completed. The events processed
  by the handler are re-delivered.
* The handler was interrupted between events (cancellation, crash) inside
  one batch. The events already processed are re-delivered along with the
  not-yet-processed ones.
* The publisher could not reach the database for one cycle (transient
  pool error). The next cycle re-fetches and re-broadcasts; subscribers
  that had already exited catch-up may see those events delivered late
  with no checkpoint advance in between, so a subsequent restart replays
  them.

Handlers must therefore be idempotent — process a duplicate event without
producing a wrong-on-replay result — or be tolerant of duplicates by some
domain-specific check (e.g., a unique key on the projection table).

=== Failure modes

The 'Kiroku.Store.Subscription.Types.SubscriptionHandleM.wait' on the
returned handle resolves with one of:

* @Right ()@ — the handler returned 'Stop' for some event and the worker
  exited cleanly. The checkpoint is saved at that event.
* @Left e@ where @e@ is 'Control.Concurrent.Async.AsyncCancelled' — the
  caller invoked 'cancel'. No checkpoint advance is guaranteed; events
  in flight at cancellation time will replay.
* @Left e@ where @e@ is
  'Kiroku.Store.Subscription.Types.SubscriptionOverflowed' — the publisher
  marked this non-group 'Kiroku.Store.Subscription.Types.AllStreams'
  subscription overflowed under
  'Kiroku.Store.Subscription.Types.DropSubscription'. Category and
  consumer-group subscriptions do not own publisher queues, so
  'queueCapacity' and 'overflowPolicy' have no effect for them.
  Investigate the slow handler and either fix the slowness, raise
  'queueCapacity', or switch to 'Kiroku.Store.Subscription.Types.DropOldest'
  if the consumer can tolerate event loss.
* @Left e@ where @e@ is a 'Hasql.Pool.UsageError' from checkpoint load —
  startup could not read the saved checkpoint. The worker stops rather than
  silently replaying from global position 0.
* @Left e@ for any exception thrown by the handler — handler exceptions
  are not caught; the worker thread dies and the original exception
  propagates to the consumer. This is intentional: a handler that
  throws is signalling that the subscription cannot proceed safely.
-}
subscribe :: (MonadIO m) => KirokuStore -> SubscriptionConfig -> m SubscriptionHandle
subscribe store config = liftIO $ do
    -- Fail fast on a misconfigured group, before any thread is spawned. A bad
    -- (member, size) is a programmer error; throwing here keeps subscribe's
    -- non-Either signature intact for every existing caller (see EP-2 Decision Log).
    for_ (consumerGroup config) $ \(ConsumerGroup m n) ->
        when (n < 1 || m < 0 || m >= n) $
            throwIO (InvalidConsumerGroup m n)
    mask $ \_restore ->
        bracketOnError
            ( case (consumerGroup config, target config) of
                (Nothing, AllStreams) -> do
                    (queue, statusVar, unsubscribe) <-
                        atomically $
                            Pub.subscribePublisher
                                (store ^. #publisher)
                                (queueCapacity config)
                                (overflowPolicy config)
                    pure (LiveFromPublisherQueue queue statusVar, unsubscribe)
                (Nothing, Category (CategoryName cat)) ->
                    pure (LiveFromCategoryNotify cat, pure ())
                (Just _, _) ->
                    pure (LiveFromGroupPolling, pure ())
            )
            (\(_, unsubscribe) -> unsubscribe)
            $ \(liveSource, unsubscribe) -> do
                -- The worker writes its current FSM state here on every transition; the
                -- handle's 'currentState' reads it. Seeded with the catch-up entry state so
                -- a read before the worker's first transition is sensible.
                stateVar <- newTVarIO (CatchingUp (GlobalPosition 0) 0)
                let pubPosVar = Pub.lastPublished (store ^. #publisher)
                    catGenVar = Notifier.categoryGenerations (store ^. #notifier)
                -- Register this worker's state cell into the store's central registry under
                -- (name, member). A fresh token identifies *this* worker so that an older
                -- worker's cleanup cannot delete a newer worker's replacement entry under
                -- the same key, and a held handle reads only the cell it registered.
                token <- newUnique
                let reg = store ^. #subscriptionRegistry
                    key = (name config, configMember config)
                    deregister =
                        atomically
                            ( modifyTVar' reg $
                                Map.update
                                    ( \(tok', cell) ->
                                        if tok' == token
                                            then Nothing
                                            else Just (tok', cell)
                                    )
                                    key
                            )
                    -- `finally cleanup` removes this subscription from the publisher's
                    -- registry whenever the worker exits — gracefully on Stop, by
                    -- cancellation, or on any exception (including SubscriptionOverflowed).
                    -- Forgetting to unsubscribe would leave a registry entry with no
                    -- reader, and the publisher would needlessly trigger this subscriber's
                    -- overflow policy on the next batch. We extend the same `finally` to
                    -- also deregister from the subscription-state registry on ANY exit, so
                    -- the registry never leaks a stale entry. The delete is
                    -- token-conditional so stale cleanup from an older duplicate-key worker
                    -- cannot remove a newer worker's live entry.
                    cleanup = unsubscribe >> deregister
                -- Insert before forking so a caller that reads a snapshot immediately after
                -- `subscribe` returns already sees this subscription.
                bracketOnError
                    (atomically $ modifyTVar' reg (Map.insert key (token, stateVar)))
                    (const deregister)
                    $ \() -> do
                        -- Every acquisition above is paired with a release in the
                        -- pre-fork window; once asyncWithUnmask returns, ownership of
                        -- both releases has transferred to the worker thread's cleanup.
                        thread <-
                            Async.asyncWithUnmask $ \unmask ->
                                unmask
                                    (runWorker (store ^. #pool) liveSource stateVar pubPosVar catGenVar config (store ^. #eventHandler) (store ^. #storeSettings))
                                    `finally` cleanup
                        pure
                            SubscriptionHandle
                                { cancel = Async.cancel thread
                                , wait = Async.waitCatch thread
                                , currentState = do
                                    m <- readTVarIO reg
                                    case Map.lookup key m of
                                        Just (tok, cell) | tok == token -> Just <$> readTVarIO cell
                                        _ -> pure Nothing
                                }

{- | Bracket-style subscription lifecycle.

Starts a subscription, runs the body, and ensures the underlying worker
thread is cancelled on either normal exit or an exception thrown inside
the body. This is the recommended way to use subscriptions: forgetting
to call 'cancel' on a 'SubscriptionHandle' leaks the worker thread for
the lifetime of the process.

Equivalent to:

    bracket (subscribe store config) cancel action

but lifted to any 'MonadUnliftIO' so it composes with effectful stacks
that already have an unlift in scope.
-}
withSubscription ::
    (MonadUnliftIO m) =>
    KirokuStore ->
    SubscriptionConfig ->
    (SubscriptionHandle -> m a) ->
    m a
withSubscription store config action = withRunInIO $ \runInIO ->
    bracket (subscribe store config) cancel (runInIO . action)

{- | A public, point-in-time view of one live subscription's state, as returned
by 'subscriptionStates'. This is the committed observability surface external
consumers (a future Prometheus exporter, an admin tool) read directly; it is
intentionally a flat view, not the internal 'SubscriptionState' cell type.

* 'subscriptionName' / 'member' identify the subscription (member 0 for a
  non-group subscription), matching the registry/checkpoint key.
* 'state' is the live 'SubscriptionState' read from the worker's registered cell.
* 'statePhase' is a stable, low-cardinality label (via
  'Kiroku.Store.Subscription.Fsm.stateName') suitable as a metric label value or
  admin column; it does not drift if a constructor's fields change.
* 'cursor' is the worker FSM cursor (via 'Kiroku.Store.Subscription.Fsm.stateCursor').
  It is the cheap live progress signal for observability, not a guaranteed
  durable checkpoint row.

Read fields via the codebase's @generic-lens@ convention, e.g. @view ^. #state@.
-}
data SubscriptionStateView = SubscriptionStateView
    { subscriptionName :: !SubscriptionName
    , member :: !Int32
    , state :: !SubscriptionState
    , statePhase :: !Text
    , cursor :: !GlobalPosition
    }
    deriving stock (Show, Generic)

{- | A near-instant snapshot of every live subscription as a public
'SubscriptionStateView', keyed by (subscription name, consumer-group member;
0 for non-group).

Snapshots the registry's outer map with 'readTVarIO', then reads each registered
state cell with 'readTVarIO' __outside__ STM. This is deliberately not one STM
transaction over all cells: a single transaction would put every subscription's
state cell in the reader's STM read set, and since each worker writes its cell
~once per batch, any such write would force the whole scan to re-run — a
reader-side cost that scales with subscription count. The named consumers (a
Prometheus scrape, an admin listing) read independent values and do not need a
globally atomic snapshot, so each entry is read as its own freshest value.

A subscription appears here from the moment 'subscribe' returns its handle until
its worker exits (stop, cancel, or crash), at which point its key is removed; a
stopped/cancelled/crashed subscription is represented by __absence__, never by a
@"stopped"@ phase (the FSM never writes 'Stopped' into the cell). This map of
view records is the committed surface the future Prometheus exporter and admin
tool consume.
-}
subscriptionStates :: KirokuStore -> IO (Map (SubscriptionName, Int32) SubscriptionStateView)
subscriptionStates store = do
    cells <- readTVarIO (store ^. #subscriptionRegistry)
    Map.traverseWithKey
        ( \(nm, mbr) (_tok, cell) -> do
            st <- readTVarIO cell
            pure
                SubscriptionStateView
                    { subscriptionName = nm
                    , member = mbr
                    , state = st
                    , statePhase = stateName st
                    , cursor = stateCursor st
                    }
        )
        cells