Schema.hs

{- | The @keiro_subscription_shards@ table: storage and lease logic for
cooperative consumer-group ownership (EP-51).

A category subscription can be split across @N@ buckets — kiroku consumer-group
member indices in @[0, N)@ — so that @N@ cooperating workers each drain a
disjoint slice of the keyspace. This module owns the durable __assignment__
layer: one row per @(subscription_name, bucket)@ recording which worker holds it
right now, as a renewable __lease__ (an owner id plus an expiry timestamp). A
live worker renews its lease on a heartbeat; a dead worker stops renewing, its
lease expires, and another worker re-claims the bucket (failover). It does
__not__ store event positions — kiroku's per-member checkpoints
(@(subscription_name, consumer_group_member)@) do that, so a re-homed bucket
resumes where its previous owner left off.

The statements here are 'Hasql.Transaction.Transaction'-flavoured so callers can
compose several into one short transaction (e.g. renew-then-claim in
'Keiro.Subscription.Shard.acquireOwnedBuckets'); the typed 'Eff'-level wrappers
that run them through kiroku's pool live in "Keiro.Subscription.Shard". The
claim uses @FOR UPDATE SKIP LOCKED@ over the claimable rows, exactly as
'Keiro.Timer.Schema.claimDueTimer' does, so two workers racing the same bucket
can never both win — a stale "how many workers are live" estimate only changes
how aggressively a worker claims, never whether ownership stays disjoint.
-}
module Keiro.Subscription.Shard.Schema (
    -- * Worker identity
    WorkerId (..),

    -- * Lease statements (composable within a transaction)
    ensureShardRows,
    claimShardsTx,
    renewLeaseTx,
    releaseShardsTx,
    listShardOwnership,
    listShardCounts,
)
where

import Contravariant.Extras (contrazip2, contrazip3, contrazip4, contrazip5)
import Data.Int (Int32)
import Data.Time (NominalDiffTime, addUTCTime)
import Data.UUID (UUID)
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Keiro.Prelude
import Kiroku.Store.Subscription.Types (SubscriptionName (..))
import "hasql-transaction" Hasql.Transaction qualified as Tx

{- | A per-process unique id naming the owner in a lease row. A UUID minted once
at worker start ('Keiro.Subscription.Shard.freshWorkerId'); two restarts of the
same binary get two different ids, so a restarted process never inherits the
dead process's leases — it claims afresh once the old leases expire.
-}
newtype WorkerId = WorkerId UUID
    deriving stock (Eq, Ord, Show)

{- | Idempotently insert the @N@ rows @(name, bucket = 0..N-1, shard_count = N)@
with @owner_worker_id@ left @NULL@. @ON CONFLICT DO NOTHING@ makes calling it
on every worker startup safe; it converges the table to exactly @N@ rows.
-}
ensureShardRows :: SubscriptionName -> Int -> Tx.Transaction ()
ensureShardRows (SubscriptionName name) shardCount =
    Tx.statement (name, fromIntegral shardCount) ensureShardRowsStmt

{- | Claim up to @targetCount@ buckets that are currently unowned __or__ whose
lease has expired (@owner_worker_id IS NULL OR lease_expires_at < now@), in one
statement, returning the bucket numbers actually claimed. @FOR UPDATE SKIP
LOCKED@ over the claimable rows is the exclusion mechanism: two workers racing
the same bucket cannot both win. The TTL is added to @now@ here (in Haskell) to
form the new @lease_expires_at@.
-}
claimShardsTx ::
    SubscriptionName -> WorkerId -> Int -> UTCTime -> NominalDiffTime -> Tx.Transaction [Int]
claimShardsTx (SubscriptionName name) (WorkerId worker) targetCount now ttl =
    fmap (fmap fromIntegral) $
        Tx.statement
            (name, now, addUTCTime ttl now, worker, fromIntegral targetCount)
            claimShardsStmt

{- | Renew every lease this worker still holds: write a fresh @lease_expires_at =
now + ttl@ and @heartbeat_at = now@ for each row it owns, returning the buckets
still held. A bucket stolen after this worker's lease lapsed is owned by someone
else and so is __not__ in the result — that is how a worker learns it lost a
bucket and stops reading it.
-}
renewLeaseTx :: SubscriptionName -> WorkerId -> UTCTime -> NominalDiffTime -> Tx.Transaction [Int]
renewLeaseTx (SubscriptionName name) (WorkerId worker) now ttl =
    fmap (fmap fromIntegral) $
        Tx.statement (name, now, addUTCTime ttl now, worker) renewLeaseStmt

{- | Graceful relinquish: clear ownership of the given buckets this worker holds
so they become claimable immediately, without waiting for lease expiry. Called
on clean shutdown. Only rows still owned by @worker@ are affected, so a bucket
already stolen is left untouched.
-}
releaseShardsTx :: SubscriptionName -> WorkerId -> [Int] -> Tx.Transaction ()
releaseShardsTx (SubscriptionName name) (WorkerId worker) buckets =
    Tx.statement (name, worker, fmap fromIntegral buckets) releaseShardsStmt

{- | Observability/test read of @(bucket, owner, lease_expires_at)@ for one
subscription, ordered by bucket.
-}
listShardOwnership :: SubscriptionName -> Tx.Transaction [(Int, Maybe WorkerId, Maybe UTCTime)]
listShardOwnership (SubscriptionName name) =
    Tx.statement name listShardOwnershipStmt

-- | Read existing shard-count groups for a subscription as @(shard_count, rows)@.
listShardCounts :: SubscriptionName -> Tx.Transaction [(Int, Int)]
listShardCounts (SubscriptionName name) =
    Tx.statement name listShardCountsStmt

ensureShardRowsStmt :: Statement (Text, Int32) ()
ensureShardRowsStmt =
    preparable
        """
        INSERT INTO keiro.keiro_subscription_shards (subscription_name, bucket, shard_count)
        SELECT $1, g, $2
        FROM generate_series(0, $2 - 1) AS g
        ON CONFLICT (subscription_name, bucket) DO NOTHING
        """
        ( contrazip2
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.int4))
        )
        D.noResult

claimShardsStmt :: Statement (Text, UTCTime, UTCTime, UUID, Int32) [Int32]
claimShardsStmt =
    preparable
        """
        WITH claimable AS (
          SELECT bucket
          FROM keiro.keiro_subscription_shards
          WHERE subscription_name = $1
            AND (owner_worker_id IS NULL OR lease_expires_at < $2)
          ORDER BY bucket
          LIMIT $5
          FOR UPDATE SKIP LOCKED
        )
        UPDATE keiro.keiro_subscription_shards s
        SET owner_worker_id = $4,
            lease_expires_at = $3,
            heartbeat_at = $2,
            updated_at = $2
        FROM claimable c
        WHERE s.subscription_name = $1 AND s.bucket = c.bucket
        RETURNING s.bucket
        """
        ( contrazip5
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.timestamptz))
            (E.param (E.nonNullable E.timestamptz))
            (E.param (E.nonNullable E.uuid))
            (E.param (E.nonNullable E.int4))
        )
        (D.rowList (D.column (D.nonNullable D.int4)))

renewLeaseStmt :: Statement (Text, UTCTime, UTCTime, UUID) [Int32]
renewLeaseStmt =
    preparable
        """
        UPDATE keiro.keiro_subscription_shards
        SET lease_expires_at = $3,
            heartbeat_at = $2,
            updated_at = $2
        WHERE subscription_name = $1
          AND owner_worker_id = $4
        RETURNING bucket
        """
        ( contrazip4
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.timestamptz))
            (E.param (E.nonNullable E.timestamptz))
            (E.param (E.nonNullable E.uuid))
        )
        (D.rowList (D.column (D.nonNullable D.int4)))

releaseShardsStmt :: Statement (Text, UUID, [Int32]) ()
releaseShardsStmt =
    preparable
        """
        UPDATE keiro.keiro_subscription_shards
        SET owner_worker_id = NULL,
            lease_expires_at = NULL,
            updated_at = now()
        WHERE subscription_name = $1
          AND owner_worker_id = $2
          AND bucket = ANY($3)
        """
        ( contrazip3
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.uuid))
            (E.param (E.nonNullable (E.foldableArray (E.nonNullable E.int4))))
        )
        D.noResult

listShardOwnershipStmt :: Statement Text [(Int, Maybe WorkerId, Maybe UTCTime)]
listShardOwnershipStmt =
    preparable
        """
        SELECT bucket, owner_worker_id, lease_expires_at
        FROM keiro.keiro_subscription_shards
        WHERE subscription_name = $1
        ORDER BY bucket
        """
        (E.param (E.nonNullable E.text))
        (D.rowList ownershipRowDecoder)

listShardCountsStmt :: Statement Text [(Int, Int)]
listShardCountsStmt =
    preparable
        """
        SELECT shard_count, count(*)::int
        FROM keiro.keiro_subscription_shards
        WHERE subscription_name = $1
        GROUP BY shard_count
        ORDER BY shard_count
        """
        (E.param (E.nonNullable E.text))
        (D.rowList ((,) <$> (fromIntegral <$> D.column (D.nonNullable D.int4)) <*> (fromIntegral <$> D.column (D.nonNullable D.int4))))

ownershipRowDecoder :: D.Row (Int, Maybe WorkerId, Maybe UTCTime)
ownershipRowDecoder =
    (,,)
        <$> (fromIntegral <$> D.column (D.nonNullable D.int4))
        <*> (fmap WorkerId <$> D.column (D.nullable D.uuid))
        <*> D.column (D.nullable D.timestamptz)