Checkpoint.hs

-- | Explicit mutation operations for durable subscription checkpoints.
--
-- Ordinary subscription checkpoint saves are monotonic. This module owns the
-- separate reset, resize and rebind operations for callers that need to move
-- persisted progress backward or forward as part of a larger transaction.
module Kiroku.Store.Subscription.Checkpoint
  ( SubscriptionCheckpointResetReport (..),
    resetSubscriptionCheckpointsTx,
    ConsumerGroupResizeReport (..),
    resizeConsumerGroupTx,
    SubscriptionTargetRebindReport (..),
    rebindSubscriptionTargetTx,
  )
where

import Data.Int (Int32)
import Data.List (nub, sort)
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import GHC.Generics (Generic)
import Hasql.Transaction qualified as Tx
import Kiroku.Store.Subscription.Checkpoint.SQL qualified as SQL
import Kiroku.Store.Subscription.Types
  ( ConsumerGroupSize,
    SubscriptionCheckpointKey (..),
    SubscriptionName (..),
    SubscriptionTarget,
    consumerGroupSizeValue,
  )
import Kiroku.Store.Types (GlobalPosition (..))

-- | Exact result of resetting a non-empty set of subscription names.
--
-- 'resetCheckpointKeys' contains every persisted @(name, member)@ row that was
-- updated. 'missingSubscriptionNames' contains requested names for which no row
-- existed. Both vectors are sorted by subscription name (and then member for
-- keys); duplicate requested names appear only once in the report.
data SubscriptionCheckpointResetReport = SubscriptionCheckpointResetReport
  { resetCheckpointKeys :: !(Vector SubscriptionCheckpointKey),
    missingSubscriptionNames :: !(Vector SubscriptionName)
  }
  deriving stock (Eq, Show, Generic)

-- | Set every existing checkpoint member for the requested subscription names
-- to the exact target position and return complete deterministic evidence.
--
-- The operation treats duplicate requested names as one name, updates all
-- persisted members for each name, and never creates checkpoint rows for missing
-- names. Unlike ordinary worker saves, this operation can move a checkpoint
-- backward. It is a 'Tx.Transaction' combinator so a caller can atomically compose
-- the reset with its own projection fence and target preparation; condemning that
-- surrounding transaction rolls back all of those writes together.
resetSubscriptionCheckpointsTx ::
  NonEmpty SubscriptionName ->
  GlobalPosition ->
  Tx.Transaction SubscriptionCheckpointResetReport
resetSubscriptionCheckpointsTx names (GlobalPosition position) = do
  rows <-
    Tx.statement
      ( Vector.fromList
          [name | SubscriptionName name <- NonEmpty.toList names],
        position
      )
      SQL.resetSubscriptionCheckpointsStmt
  pure
    SubscriptionCheckpointResetReport
      { resetCheckpointKeys = Vector.mapMaybe resetKey rows,
        missingSubscriptionNames = Vector.mapMaybe missingName rows
      }
  where
    resetKey (name, Just member) =
      Just (SubscriptionCheckpointKey (SubscriptionName name) member)
    resetKey (_, Nothing) = Nothing
    missingName (name, Nothing) = Just (SubscriptionName name)
    missingName (_, Just _) = Nothing

-- | Evidence of explicit topology equalization. Previous sizes are sorted and
-- distinct; all new members resume from 'resumePosition'. A missing group starts
-- at zero. A repeat reports the now-equalized topology without moving progress.
data ConsumerGroupResizeReport = ConsumerGroupResizeReport
  { previousSizes :: !(Vector Int32),
    previousMemberCount :: !Int,
    newSize :: !ConsumerGroupSize,
    resumePosition :: !GlobalPosition
  }
  deriving stock (Eq, Show, Generic)

-- | Stop every worker for the name before calling this operation. Lock its
-- checkpoint set, rewind every new member to the old minimum, and remove obsolete
-- members atomically. Existing member identities are retained. A hash-assignment
-- change also requires this equalization even when the group size is unchanged.
-- The caller can compose or roll back the resize with application-owned SQL.
resizeConsumerGroupTx ::
  SubscriptionName -> ConsumerGroupSize -> Tx.Transaction ConsumerGroupResizeReport
resizeConsumerGroupTx (SubscriptionName name) newSize = do
  Tx.statement name SQL.lockCheckpointNameStmt
  rows <- Tx.statement name SQL.lockCheckpointRowsStmt
  let position = if Vector.null rows then 0 else Vector.minimum (Vector.map (\(_, _, p, _) -> p) rows)
      previousSizes = Vector.fromList . sort . nub $ [n | (_, n, _, _) <- Vector.toList rows]
  let binding = if Vector.null rows then Nothing else let (_, _, _, b) = Vector.head rows in b
      (kind, category) = SQL.targetColumns binding
  Tx.statement (name, consumerGroupSizeValue newSize, position, kind, category) SQL.resizeCheckpointMembersStmt
  pure
    ConsumerGroupResizeReport
      { previousSizes = previousSizes,
        previousMemberCount = Vector.length rows,
        newSize = newSize,
        resumePosition = GlobalPosition position
      }

-- | Rebind evidence includes every distinct prior binding; Nothing denotes legacy rows.
data SubscriptionTargetRebindReport = SubscriptionTargetRebindReport
  { reboundMemberCount :: !Int,
    previousBindings :: !(Vector (Maybe SubscriptionTarget)),
    reboundTarget :: !SubscriptionTarget,
    reboundPosition :: !GlobalPosition
  }
  deriving stock (Eq, Show)

-- | Stop all workers first. Bind every existing member and explicitly reset its
-- position in the caller's transaction. A missing name fails the transaction;
-- repeating the operation leaves target and progress unchanged.
rebindSubscriptionTargetTx ::
  SubscriptionName -> SubscriptionTarget -> GlobalPosition -> Tx.Transaction SubscriptionTargetRebindReport
rebindSubscriptionTargetTx (SubscriptionName name) target position@(GlobalPosition pos) = do
  Tx.statement name SQL.lockCheckpointNameStmt
  rows <- Tx.statement name SQL.lockCheckpointRowsStmt
  let (kind, category) = SQL.targetColumns (Just target)
  count <- Tx.statement (name, kind, category, pos) SQL.rebindCheckpointTargetStmt
  pure
    SubscriptionTargetRebindReport
      { reboundMemberCount = fromIntegral count,
        previousBindings = Vector.fromList . nub $ [binding | (_, _, _, binding) <- Vector.toList rows],
        reboundTarget = target,
        reboundPosition = position
      }