EventPublisher.hs

-- | The centralized @$all@ broadcaster shared by queue-consuming live clients.
--
-- A single 'EventPublisher' reads new events from the global stream once when at
-- least one 'Subscriber' is registered and fans them out through each
-- subscriber's bounded 'Control.Concurrent.STM.TBQueue'. When the registry is
-- empty, it does not fetch full event rows; it advances 'lastPublished' from the
-- @$all@ tail with a single-row query so DB-driven subscriptions can still use the
-- position as their live gate. 'startPublisher' launches the loop,
-- 'subscribePublisher' registers a subscriber (returning its queue, a status
-- 'TVar', and an unsubscribe action), 'publisherPosition' reports the last
-- published cursor, and 'stopPublisher' tears the loop down.
--
-- When a subscriber's queue fills, the publisher applies that subscriber's
-- 'Kiroku.Store.Subscription.Types.OverflowPolicy' by setting its
-- 'SubscriberStatus': under the default @PauseAndResume@ it marks the subscriber
-- 'Paused' and stops pushing (the worker drains and re-catches-up losslessly),
-- under @DropSubscription@ it marks it 'Overflowed', and under @DropOldest@ it
-- evicts the oldest batch and counts the drop in 'subDropped'. The publisher itself never blocks on a slow consumer.
module Kiroku.Store.Subscription.EventPublisher
  ( EventPublisher (..),
    Subscriber (..),
    SubscriberStatus (..),
    startPublisher,
    stopPublisher,
    subscribePublisher,
    PublisherSubscription (..),
    subscribePublisherWith,
    publisherPosition,
  )
where

import Control.Concurrent.Async (Async)
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM
  ( STM,
    TBQueue,
    TChan,
    TVar,
    atomically,
    dupTChan,
    isFullTBQueue,
    modifyTVar',
    newTBQueue,
    newTVar,
    newTVarIO,
    readTChan,
    readTVar,
    readTVarIO,
    registerDelay,
    tryReadTBQueue,
    tryReadTChan,
    writeTBQueue,
    writeTVar,
  )
import Control.Concurrent.STM qualified as STM
import Control.Exception (SomeAsyncException, SomeException, asyncExceptionFromException, catch, throwIO)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Foldable (for_)
import Data.Int (Int32)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Vector qualified as V
import Data.Word (Word64)
import Hasql.Pool (Pool)
import Hasql.Pool qualified as Pool
import Hasql.Session qualified as Session
import Kiroku.Store.Observability (KirokuEvent (..), emitOrDrop)
import Kiroku.Store.SQL qualified as SQL
import Kiroku.Store.Settings (DecodedBatch (..), DecodedEvent (..), StoreSettings, decodeEvents)
import Kiroku.Store.Subscription.Types (OverflowPolicy (..))
import Kiroku.Store.Types (GlobalPosition (..), RecordedEvent (..))
import Numeric.Natural (Natural)

-- | The centralized EventPublisher. With registered subscribers, reads events
-- from the database once per notification and fans them out to a registry of
-- bounded per-subscriber queues. With an empty registry, fetches only the current
-- @$all@ tail to keep 'lastPublished' moving without decoding event rows.
data EventPublisher = EventPublisher
  { -- | Active subscriber registry, keyed by an internal id.
    subscribers :: !(TVar (IntMap Subscriber)),
    nextSubscriberId :: !(TVar Int),
    publisherThread :: !(Async ()),
    -- | Last-published position; workers read this to know when catch-up is done
    lastPublished :: !(TVar GlobalPosition)
  }

-- | A registered subscriber. The publisher owns one of these per
-- 'Kiroku.Store.Subscription.subscribe' call. The worker reads from
-- 'subQueue' under normal operation; on overflow under 'DropSubscription'
-- the publisher flips 'subStatus' to 'Overflowed' and the worker
-- terminates the subscription with 'SubscriptionOverflowed'.
data Subscriber = Subscriber
  { subQueue :: !(TBQueue DecodedBatch),
    subStatus :: !(TVar SubscriberStatus),
    subPolicy :: !OverflowPolicy,
    -- | Dropped batches under DropOldest; compare successive readings (modulo Word64).
    subDropped :: !(TVar Word64)
  }

-- | A bounded publisher subscription with its exact dropped-batch counter.
data PublisherSubscription = PublisherSubscription
  { subscriptionQueue :: !(TBQueue DecodedBatch),
    subscriptionStatus :: !(TVar SubscriberStatus),
    subscriptionDropped :: !(TVar Word64),
    -- | Idempotent deregistration; invoke on every exit.
    unsubscribe :: !(IO ())
  }

-- | Subscriber lifecycle status as observed by its worker.
data SubscriberStatus
  = -- | Healthy; worker reads from the queue normally.
    Active
  | -- | Recoverable backpressure under 'PauseAndResume': the queue was full
    --       so the publisher stopped pushing and set this flag. The worker observes
    --       it, drains the stale queue, clears the flag back to 'Active', and
    --       re-catches-up from its checkpoint to recover any skipped events.
    Paused
  | -- | Publisher signalled overflow under 'DropSubscription'; worker should
    --       surface a structured error.
    Overflowed
  deriving stock (Eq, Show)

-- | Publisher batch size — larger than subscriber batch size because the
-- Publisher serves all subscribers.
publisherBatchSize :: Int32
publisherBatchSize = 1000

-- | Safety poll interval in microseconds (30 seconds).
safetyPollMicros :: Int
safetyPollMicros = 30_000_000

-- | Start the EventPublisher. Spawns a thread that waits for ticks from the
-- Notifier (or a 30-second safety poll timeout). With at least one registered
-- subscriber, the thread queries full event rows and delivers them to every active
-- subscriber; with an empty registry, it advances 'lastPublished' from the @$all@
-- tail using a single-row query and fetches no payload rows. Per-subscriber
-- overflow handling is determined by each subscription's 'OverflowPolicy'.
--
-- If an 'eventHandler' callback is supplied, the publisher emits
-- 'Kiroku.Store.Observability.KirokuEventPublisherPoolError' when its read
-- query returns a 'Pool.UsageError'. The publisher continues to run on the
-- 30-second safety poll regardless of the event firing; the event is the
-- operator's only structured signal that pool exhaustion or a server
-- error stalled the broadcast.
startPublisher ::
  (MonadIO m) =>
  Pool ->
  TChan () ->
  -- | optional event handler for publisher-side observability
  Maybe (KirokuEvent -> IO ()) ->
  -- | interpreter-level event hooks; 'Kiroku.Store.Settings.decodeHook'
  --     runs once per fetched batch before fan-out to subscribers.
  StoreSettings ->
  m EventPublisher
startPublisher pool notifierChan mHandler stSettings = liftIO $ do
  subsVar <- newTVarIO IntMap.empty
  nextIdVar <- newTVarIO 0
  -- Duplicate the notifier channel before reading the tail so ticks arriving
  -- during startup are redundant rather than lost until the safety poll.
  tickChan <- atomically (dupTChan notifierChan)
  tailResult <- Pool.use pool (Session.statement () SQL.currentGlobalPositionStmt)
  tailPos <- either throwIO pure tailResult
  pos <- newTVarIO (GlobalPosition tailPos)
  thread <- Async.async (publisherLoop pool tickChan subsVar pos mHandler stSettings)
  pure
    EventPublisher
      { subscribers = subsVar,
        nextSubscriberId = nextIdVar,
        publisherThread = thread,
        lastPublished = pos
      }

-- | Stop the EventPublisher thread.
stopPublisher :: (MonadIO m) => EventPublisher -> m ()
stopPublisher pub = liftIO $ do
  Async.cancel (publisherThread pub)
  () <$ Async.waitCatch (publisherThread pub)

-- | Register a new subscriber with the given queue capacity and overflow
-- policy. Returns the bounded queue, the subscriber's status TVar, and an
-- @unsubscribe@ action that the caller must invoke when the subscription
-- ends (worker exits, cancellation). Failing to unsubscribe causes the
-- publisher to keep delivering to a queue with no reader, which will fill
-- and trigger the configured policy unnecessarily.
subscribePublisher ::
  EventPublisher ->
  -- | Queue capacity (number of batches)
  Natural ->
  OverflowPolicy ->
  STM (TBQueue DecodedBatch, TVar SubscriberStatus, IO ())
subscribePublisher pub cap policy = do
  sub <- subscribePublisherWith pub cap policy
  pure (subscriptionQueue sub, subscriptionStatus sub, unsubscribe sub)

-- | Register a subscriber, also exposing the DropOldest counter (initially zero).
subscribePublisherWith :: EventPublisher -> Natural -> OverflowPolicy -> STM PublisherSubscription
subscribePublisherWith pub cap policy = do
  queue <- newTBQueue cap
  status <- newTVar Active
  dropped <- newTVar 0
  sid <- readTVar (nextSubscriberId pub)
  writeTVar (nextSubscriberId pub) (sid + 1)
  let sub = Subscriber {subQueue = queue, subStatus = status, subPolicy = policy, subDropped = dropped}
  modifyTVar' (subscribers pub) (IntMap.insert sid sub)
  let deregister =
        atomically $
          modifyTVar' (subscribers pub) (IntMap.delete sid)
  pure (PublisherSubscription queue status dropped deregister)

-- | Read the last-published global position.
publisherPosition :: EventPublisher -> STM GlobalPosition
publisherPosition pub = readTVar (lastPublished pub)

-- Internal: the publisher loop.
publisherLoop ::
  Pool ->
  TChan () ->
  TVar (IntMap Subscriber) ->
  TVar GlobalPosition ->
  Maybe (KirokuEvent -> IO ()) ->
  StoreSettings ->
  IO ()
publisherLoop pool tickChan subsVar posVar mHandler stSettings = loop
  where
    loop = do
      -- Wait for a tick OR safety poll timeout
      waitForWakeup tickChan
      -- Drain all pending ticks (debouncing)
      drainTicks tickChan
      -- Fetch and broadcast. If a synchronous callback failure occurs after
      -- partial delivery, lastPublished is not advanced, so the next tick
      -- re-fetches under Kiroku's at-least-once delivery contract.
      fetchAndBroadcast `catch` \(e :: SomeException) ->
        case asyncExceptionFromException e of
          Just (ae :: SomeAsyncException) -> throwIO ae
          Nothing -> emitOrDrop mHandler (KirokuEventPublisherLoopError e)
      loop

    fetchAndBroadcast = do
      subs <- readTVarIO subsVar
      if IntMap.null subs
        then cheapAdvance
        else fullFetch

    cheapAdvance = do
      result <- Pool.use pool (Session.statement () SQL.currentGlobalPositionStmt)
      case result of
        Left err -> do
          -- Surface the pool error so operators see why the publisher
          -- position has stalled; the 30-second safety poll will retry.
          emitOrDrop mHandler (KirokuEventPublisherPoolError err)
        Right tailPos -> do
          -- Re-check registry emptiness in the same STM transaction as
          -- the position write. If a queue subscriber registered after
          -- our snapshot, fall through to a full fetch so no event skips
          -- a queue that now exists.
          raced <- atomically $ do
            subs' <- readTVar subsVar
            if IntMap.null subs'
              then do
                GlobalPosition cur <- readTVar posVar
                let nextPos = max cur tailPos
                -- Do not store a chain of lazy max expressions that
                -- retains earlier tail-query results.
                nextPos `seq` writeTVar posVar (GlobalPosition nextPos)
                pure False
              else pure True
          when raced fullFetch

    fullFetch = do
      GlobalPosition pos <- readTVarIO posVar
      result <- Pool.use pool (Session.statement (pos, publisherBatchSize) SQL.readAllForwardStmt)
      case result of
        Left err -> do
          -- Surface the pool error so operators see why broadcast
          -- has stalled; the 30-second safety poll will retry.
          emitOrDrop mHandler (KirokuEventPublisherPoolError err)
        Right rawEvents
          | V.null rawEvents -> pure ()
          | otherwise -> do
              -- Apply the decodeHook once per batch — every subscriber
              -- observes the same transformed view, and the cost is
              -- paid in one place instead of per-subscriber.
              events <- decodeEvents stSettings rawEvents
              case events of
                UnchangedBatch _ -> pure ()
                TransformedBatch decoded -> for_ decoded $ \case
                  Decoded _ -> pure ()
                  Undecodable event failure -> emitOrDrop mHandler (KirokuEventPublisherDecodeFailed (globalPosition event) (eventId event) failure)
              let lastEvent = V.last rawEvents
                  newPos = globalPosition lastEvent
              -- Snapshot the current subscriber set, then deliver outside
              -- the snapshot's STM transaction. Each delivery is its own
              -- atomic step so a slow subscriber under DropSubscription
              -- cannot rollback another's enqueue.
              subs <- readTVarIO subsVar
              for_ (IntMap.elems subs) (deliverBatch events)
              -- Advance posVar only after attempting delivery to the
              -- snapshot cohort. In the same transaction, re-read the
              -- registry and offer the in-flight batch to subscribers that
              -- registered after the snapshot. This closes the attach race:
              -- every subscriber either received the batch before the
              -- advance or registered after the advanced position was
              -- visible and will cover the range through SQL catch-up.
              atomically $ do
                subs' <- readTVar subsVar
                for_ (IntMap.elems (subs' `IntMap.difference` subs)) (deliverBatchSTM events)
                writeTVar posVar newPos
              -- If we got a full batch, there may be more — loop immediately
              if V.length rawEvents >= fromIntegral publisherBatchSize
                then fullFetch
                else pure ()

    deliverBatch events sub = atomically $ do
      deliverBatchSTM events sub

    deliverBatchSTM events sub = do
      full <- isFullTBQueue (subQueue sub)
      if not full
        then do
          -- There is space again: clear a prior 'PauseAndResume' pause so
          -- the worker (which also clears it on resume) is not left waiting.
          -- Never clear a terminal 'DropSubscription' overflow.
          status <- readTVar (subStatus sub)
          case status of
            Paused -> writeTVar (subStatus sub) Active
            Active -> pure ()
            Overflowed -> pure ()
          writeTBQueue (subQueue sub) events
        else case subPolicy sub of
          -- Recoverable: signal the pause and stop pushing. Do not drop;
          -- the worker re-reads the skipped events from its checkpoint.
          PauseAndResume -> writeTVar (subStatus sub) Paused
          DropSubscription -> writeTVar (subStatus sub) Overflowed
          DropOldest -> do
            _ <- tryReadTBQueue (subQueue sub)
            modifyTVar' (subDropped sub) (+ 1)
            writeTBQueue (subQueue sub) events

-- Wait for either a tick or a 30-second timeout (safety poll).
waitForWakeup :: TChan () -> IO ()
waitForWakeup tickChan = do
  timerVar <- registerDelay safetyPollMicros
  atomically $
    (readTChan tickChan)
      `STM.orElse` (readTVar timerVar >>= STM.check)

-- Drain all pending ticks from the channel (debouncing).
drainTicks :: TChan () -> IO ()
drainTicks tickChan = atomically go
  where
    go = do
      mTick <- tryReadTChan tickChan
      case mTick of
        Nothing -> pure ()
        Just () -> go