Schema.hs

-- | The @keiro_timers@ table: storage and claim logic for durable timers.
--
-- Holds one row per scheduled timer with its 'TimerStatus' lifecycle.
-- 'scheduleTimerTx' inserts (or re-arms a still-@Scheduled@ timer with the same
-- id) inside the caller's transaction; 'claimDueTimer' atomically picks the
-- single earliest due timer with @FOR UPDATE SKIP LOCKED@ and moves it to
-- @Firing@, so competing workers never claim the same timer; 'markTimerFired'
-- records completion and the produced event id. Stale @Firing@ rows are requeued
-- by 'requeueStuckTimers' so a crashed worker does not strand a timer forever.
--
-- Callers normally use the re-exports from "Keiro.Timer" rather than this
-- module directly.
module Keiro.Timer.Schema
  ( -- * Rows and status
    TimerStatus (..),
    TimerRow (..),

    -- * Storage
    scheduleTimerTx,
    scheduleTimerOnceTx,
    claimDueTimer,
    markTimerFired,

    -- * Read-only counts
    countDueTimers,
    countStuckTimers,

    -- * Recovery
    StuckTimerFilter (..),
    anyStuckTimer,
    findStuckTimers,
    requeueStuckTimers,
    requeueStuckTimer,
    cancelTimer,
    deadLetterTimer,
  )
where

import Contravariant.Extras (contrazip2, contrazip6)
import Data.Time (NominalDiffTime, addUTCTime)
import Data.UUID (UUID)
import Effectful (Eff, (:>))
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Keiro.Prelude
import Keiro.Timer.Types (TimerId (..), TimerRequest (..))
import Kiroku.Store.Effect (Store)
import Kiroku.Store.Transaction (runTransaction)
import Kiroku.Store.Types (EventId (..))
import "hasql-transaction" Hasql.Transaction qualified as Tx

-- | A timer's lifecycle state.
--
-- * 'Scheduled' — waiting for its 'fireAt'; claimable.
-- * 'Firing' — claimed by a worker and being processed; stale rows become
--   claimable again when 'requeueStuckTimers' moves them back to 'Scheduled'.
-- * 'Fired' — successfully fired; terminal.
-- * 'Cancelled' — withdrawn before firing.
-- * 'Dead' — abandoned after exceeding the attempt ceiling; terminal; carries an
--   optional @last_error@ describing why it was given up on.
data TimerStatus
  = Scheduled
  | Firing
  | Fired
  | Cancelled
  | Dead
  deriving stock (Generic, Eq, Show)

-- | A timer row as stored: the original 'TimerRequest' fields plus the live
-- 'status', the 'attempts' count (incremented on each claim), and the
-- 'firedEventId' recorded once it fires.
data TimerRow = TimerRow
  { timerId :: !TimerId,
    processManagerName :: !Text,
    correlationId :: !Text,
    fireAt :: !UTCTime,
    payload :: !Value,
    status :: !TimerStatus,
    attempts :: !Int,
    firedEventId :: !(Maybe EventId)
  }
  deriving stock (Generic, Eq, Show)

-- | Criteria selecting timers stranded in 'Firing'. A row is "stuck" when its
-- 'status' is @firing@ and it matches every set bound: 'minAge' (it has been
-- firing at least this long, measured from @updated_at@) and 'minAttempts' (it
-- has been claimed at least this many times). Both unset selects every @firing@
-- row.
data StuckTimerFilter = StuckTimerFilter
  { minAge :: !(Maybe NominalDiffTime),
    minAttempts :: !(Maybe Int)
  }
  deriving stock (Generic, Eq, Show)

-- | Select every @firing@ row regardless of age or attempts.
anyStuckTimer :: StuckTimerFilter
anyStuckTimer = StuckTimerFilter Nothing Nothing

-- | Schedule a timer inside the caller's transaction (typically a process
-- manager's append). Upserts on 'timerId': a conflicting row is re-armed only
-- while it is still @Scheduled@, so a timer that has already fired or been
-- cancelled is not resurrected.
scheduleTimerTx :: TimerRequest -> Tx.Transaction ()
scheduleTimerTx request =
  Tx.statement
    ( timerIdToUuid (request ^. #timerId),
      request ^. #processManagerName,
      request ^. #correlationId,
      request ^. #fireAt,
      request ^. #payload,
      statusToText Scheduled
    )
    scheduleTimerStmt

-- | Schedule a timer only if no row with the same 'timerId' already exists.
--
-- This is for callers whose first arm must win, such as durable workflow sleeps:
-- every resume pass re-runs the sleep arm until the timer fires, and preserving
-- the original 'fireAt' keeps the sleep measured from the first arm. Process
-- managers that intentionally push a deadline back should keep using
-- 'scheduleTimerTx'. Returns 'True' when this call inserted the row and 'False'
-- when an existing timer won.
scheduleTimerOnceTx :: TimerRequest -> Tx.Transaction Bool
scheduleTimerOnceTx request =
  Tx.statement
    ( timerIdToUuid (request ^. #timerId),
      request ^. #processManagerName,
      request ^. #correlationId,
      request ^. #fireAt,
      request ^. #payload,
      statusToText Scheduled
    )
    scheduleTimerOnceStmt

-- | Atomically claim the single earliest timer due at @now@, moving it to
-- @Firing@ and bumping its attempt count. Uses @FOR UPDATE SKIP LOCKED@ so
-- concurrent workers each get a distinct timer. Returns 'Nothing' when none is
-- due.
claimDueTimer :: (Store :> es) => UTCTime -> Eff es (Maybe TimerRow)
claimDueTimer now =
  runTransaction $
    Tx.statement now claimDueTimerStmt

-- | Mark a claimed timer @Fired@, recording the id of the event its firing
-- produced. Returns 'False' when the row left @Firing@ while the fire action was
-- running (for example, it was requeued, cancelled, or dead-lettered).
markTimerFired :: (Store :> es) => TimerId -> EventId -> Eff es Bool
markTimerFired timerId eventId =
  runTransaction $
    Tx.statement (timerIdToUuid timerId, eventIdToUuid eventId) markTimerFiredStmt

-- | Count timers that are @scheduled@ and already due at @now@ — the timer
-- backlog. Read-only; mirrors 'claimDueTimer''s WHERE clause but counts rather
-- than locking, so it never claims or mutates a row.
countDueTimers :: (Store :> es) => UTCTime -> Eff es Int
countDueTimers now =
  runTransaction $
    Tx.statement now countDueTimersStmt

-- | Count timers stranded in @Firing@ that match the given 'StuckTimerFilter' —
-- the same "stuck" predicate 'findStuckTimers' lists, evaluated against @now@.
-- Read-only. 'anyStuckTimer' counts every @firing@ row.
countStuckTimers :: (Store :> es) => UTCTime -> StuckTimerFilter -> Eff es Int
countStuckTimers now stuckFilter =
  runTransaction $
    Tx.statement (cutoff, fmap fromIntegral (stuckFilter ^. #minAttempts)) countStuckTimersStmt
  where
    cutoff = fmap (\age -> addUTCTime (negate age) now) (stuckFilter ^. #minAge)

-- | List timers stranded in @Firing@ that match the given 'StuckTimerFilter'.
-- The @minAge@ bound is evaluated against @now@: a row qualifies when its
-- @updated_at@ is at least @minAge@ in the past (cutoff @now - minAge@). Results
-- are ordered oldest-first by @updated_at@.
findStuckTimers ::
  (Store :> es) => UTCTime -> StuckTimerFilter -> Eff es [TimerRow]
findStuckTimers now stuckFilter =
  runTransaction $
    Tx.statement (cutoff, fmap fromIntegral (stuckFilter ^. #minAttempts)) findStuckTimersStmt
  where
    cutoff = fmap (\age -> addUTCTime (negate age) now) (stuckFilter ^. #minAge)

-- | Move every timer stranded in @Firing@ for at least @olderThan@ back to
-- @Scheduled@. The statement preserves @fire_at@, so a due timer becomes
-- claimable on the same worker pass. Returns the number of rows requeued.
requeueStuckTimers :: (Store :> es) => NominalDiffTime -> UTCTime -> Eff es Int
requeueStuckTimers olderThan now =
  runTransaction $
    Tx.statement cutoff requeueStuckTimersStmt
  where
    cutoff = addUTCTime (negate olderThan) now

-- | Move a timer from @Firing@ back to @Scheduled@ so the ordinary claim loop
-- re-fires it. Leaves @fire_at@ unchanged, so a due timer becomes immediately
-- re-claimable. Idempotent: only @firing@ rows match, so re-running on an
-- already-requeued row affects nothing. Returns 'True' when a row changed.
requeueStuckTimer :: (Store :> es) => TimerId -> Eff es Bool
requeueStuckTimer timerId =
  runTransaction $
    Tx.statement (timerIdToUuid timerId) requeueStuckTimerStmt

-- | Move a timer from @Scheduled@ or @Firing@ to the terminal @Cancelled@
-- state so it never fires. Terminal rows (@fired@, @cancelled@, @dead@) are left
-- untouched. Idempotent. Returns 'True' when a row changed.
cancelTimer :: (Store :> es) => TimerId -> Eff es Bool
cancelTimer timerId =
  runTransaction $
    Tx.statement (timerIdToUuid timerId) cancelTimerStmt

-- | Move a timer from @Scheduled@ or @Firing@ to the terminal @Dead@ state,
-- recording @reason@ in @last_error@ so an operator can see why it was abandoned
-- (@SELECT * FROM keiro_timers WHERE status = 'dead'@). Terminal rows are left
-- untouched. Idempotent. Returns 'True' when a row changed.
deadLetterTimer :: (Store :> es) => TimerId -> Text -> Eff es Bool
deadLetterTimer timerId reason =
  runTransaction $
    Tx.statement (timerIdToUuid timerId, reason) deadLetterTimerStmt

scheduleTimerStmt :: Statement (UUID, Text, Text, UTCTime, Value, Text) ()
scheduleTimerStmt =
  preparable
    """
    INSERT INTO keiro.keiro_timers
      (timer_id, process_manager_name, correlation_id, fire_at, payload, status)
    VALUES
      ($1, $2, $3, $4, $5, $6)
    ON CONFLICT (timer_id) DO UPDATE
      SET process_manager_name = EXCLUDED.process_manager_name,
          correlation_id = EXCLUDED.correlation_id,
          fire_at = EXCLUDED.fire_at,
          payload = EXCLUDED.payload,
          status = EXCLUDED.status,
          updated_at = now()
      WHERE keiro_timers.status = 'scheduled'
    """
    ( contrazip6
        (E.param (E.nonNullable E.uuid))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
        (E.param (E.nonNullable E.jsonb))
        (E.param (E.nonNullable E.text))
    )
    D.noResult

scheduleTimerOnceStmt :: Statement (UUID, Text, Text, UTCTime, Value, Text) Bool
scheduleTimerOnceStmt =
  preparable
    """
    INSERT INTO keiro.keiro_timers
      (timer_id, process_manager_name, correlation_id, fire_at, payload, status)
    VALUES
      ($1, $2, $3, $4, $5, $6)
    ON CONFLICT (timer_id) DO NOTHING
    """
    ( contrazip6
        (E.param (E.nonNullable E.uuid))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable E.timestamptz))
        (E.param (E.nonNullable E.jsonb))
        (E.param (E.nonNullable E.text))
    )
    ((> 0) <$> D.rowsAffected)

claimDueTimerStmt :: Statement UTCTime (Maybe TimerRow)
claimDueTimerStmt =
  preparable
    """
    WITH due AS (
      SELECT timer_id
      FROM keiro.keiro_timers
      WHERE status = 'scheduled'
        AND fire_at <= $1
      ORDER BY fire_at, timer_id
      LIMIT 1
      FOR UPDATE SKIP LOCKED
    )
    UPDATE keiro.keiro_timers kt
    SET status = 'firing',
        attempts = kt.attempts + 1,
        updated_at = now()
    FROM due
    WHERE kt.timer_id = due.timer_id
    RETURNING kt.timer_id, kt.process_manager_name, kt.correlation_id, kt.fire_at,
      kt.payload, kt.status, kt.attempts, kt.fired_event_id
    """
    (E.param (E.nonNullable E.timestamptz))
    (D.rowMaybe timerRowDecoder)

markTimerFiredStmt :: Statement (UUID, UUID) Bool
markTimerFiredStmt =
  preparable
    """
    UPDATE keiro.keiro_timers
    SET status = 'fired',
        fired_event_id = $2,
        updated_at = now()
    WHERE timer_id = $1
      AND status = 'firing'
    """
    ( contrazip2
        (E.param (E.nonNullable E.uuid))
        (E.param (E.nonNullable E.uuid))
    )
    ((> 0) <$> D.rowsAffected)

countDueTimersStmt :: Statement UTCTime Int
countDueTimersStmt =
  preparable
    """
    SELECT count(*)
    FROM keiro.keiro_timers
    WHERE status = 'scheduled'
      AND fire_at <= $1
    """
    (E.param (E.nonNullable E.timestamptz))
    (D.singleRow (fromIntegral <$> D.column (D.nonNullable D.int8)))

countStuckTimersStmt :: Statement (Maybe UTCTime, Maybe Int64) Int
countStuckTimersStmt =
  preparable
    """
    SELECT count(*)
    FROM keiro.keiro_timers
    WHERE status = 'firing'
      AND ($1::timestamptz IS NULL OR updated_at <= $1)
      AND ($2::bigint IS NULL OR attempts >= $2)
    """
    ( contrazip2
        (E.param (E.nullable E.timestamptz))
        (E.param (E.nullable E.int8))
    )
    (D.singleRow (fromIntegral <$> D.column (D.nonNullable D.int8)))

findStuckTimersStmt :: Statement (Maybe UTCTime, Maybe Int64) [TimerRow]
findStuckTimersStmt =
  preparable
    """
    SELECT timer_id, process_manager_name, correlation_id, fire_at,
      payload, status, attempts, fired_event_id
    FROM keiro.keiro_timers
    WHERE status = 'firing'
      AND ($1::timestamptz IS NULL OR updated_at <= $1)
      AND ($2::bigint IS NULL OR attempts >= $2)
    ORDER BY updated_at, timer_id
    """
    ( contrazip2
        (E.param (E.nullable E.timestamptz))
        (E.param (E.nullable E.int8))
    )
    (D.rowList timerRowDecoder)

requeueStuckTimersStmt :: Statement UTCTime Int
requeueStuckTimersStmt =
  preparable
    """
    UPDATE keiro.keiro_timers
    SET status = 'scheduled',
        updated_at = now()
    WHERE status = 'firing'
      AND updated_at <= $1
    """
    (E.param (E.nonNullable E.timestamptz))
    (fromIntegral <$> D.rowsAffected)

requeueStuckTimerStmt :: Statement UUID Bool
requeueStuckTimerStmt =
  preparable
    """
    UPDATE keiro.keiro_timers
    SET status = 'scheduled',
        updated_at = now()
    WHERE timer_id = $1
      AND status = 'firing'
    """
    (E.param (E.nonNullable E.uuid))
    ((> 0) <$> D.rowsAffected)

cancelTimerStmt :: Statement UUID Bool
cancelTimerStmt =
  preparable
    """
    UPDATE keiro.keiro_timers
    SET status = 'cancelled',
        updated_at = now()
    WHERE timer_id = $1
      AND status IN ('scheduled', 'firing')
    """
    (E.param (E.nonNullable E.uuid))
    ((> 0) <$> D.rowsAffected)

deadLetterTimerStmt :: Statement (UUID, Text) Bool
deadLetterTimerStmt =
  preparable
    """
    UPDATE keiro.keiro_timers
    SET status = 'dead',
        last_error = $2,
        updated_at = now()
    WHERE timer_id = $1
      AND status IN ('scheduled', 'firing')
    """
    ( contrazip2
        (E.param (E.nonNullable E.uuid))
        (E.param (E.nonNullable E.text))
    )
    ((> 0) <$> D.rowsAffected)

timerRowDecoder :: D.Row TimerRow
timerRowDecoder =
  TimerRow
    <$> (TimerId <$> D.column (D.nonNullable D.uuid))
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.text)
    <*> D.column (D.nonNullable D.timestamptz)
    <*> D.column (D.nonNullable D.jsonb)
    <*> D.column (D.nonNullable (D.refine statusFromText D.text))
    <*> (fromIntegral <$> D.column (D.nonNullable D.int8))
    <*> (fmap EventId <$> D.column (D.nullable D.uuid))

statusToText :: TimerStatus -> Text
statusToText = \case
  Scheduled -> "scheduled"
  Firing -> "firing"
  Fired -> "fired"
  Cancelled -> "cancelled"
  Dead -> "dead"

statusFromText :: Text -> Either Text TimerStatus
statusFromText = \case
  "scheduled" -> Right Scheduled
  "firing" -> Right Firing
  "fired" -> Right Fired
  "cancelled" -> Right Cancelled
  "dead" -> Right Dead
  other -> Left ("unknown keiro_timers.status: " <> other)

timerIdToUuid :: TimerId -> UUID
timerIdToUuid (TimerId uuid) = uuid

eventIdToUuid :: EventId -> UUID
eventIdToUuid (EventId uuid) = uuid