{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Lifecycle SQL templates.
module Arbiter.Core.Sql.Lifecycle
  ( smartAckJobSQL
  , smartAckJobsBatchSQL
  , setVisibilityTimeoutSQL
  , setVisibilityTimeoutBatchSQL
  , updateJobForRetrySQL
  , nackJobSQL
  , nackJobsBatchSQL
  , promoteJobSQL
  , rescheduleJobSQL
  ) where

import Data.Int (Int32, Int64)
import Data.Text (Text)
import Data.Time (UTCTime)

import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueGroupsTable, jobQueueTable)
import Arbiter.Core.Job.Types (JobRead, claimSeq, claimedBy, primaryKey)
import Arbiter.Core.Sql.Archive (archiveAckCte)
import Arbiter.Core.Sql.Groups (settleGroupLocksCte)
import Arbiter.Core.Sql.QQ (sql, stmt)
import Arbiter.Core.Sql.Query (Query, mwhen, sepBy)
import Arbiter.Core.Sql.Tree (lockedByIdsCte, wakeParentCte)
import Arbiter.Core.SqlLiterals (attemptsLeftSQL)

-- | Parent-aware ack. Deletes a childless job, suspends one whose children are still
-- running, and wakes a suspended parent whose last child left the queue. Returns 1,
-- or 0 for a job gone, reclaimed or cancelled. When @archiveEnabled@, the deleted row
-- is teed into the archive per-row on @archive_for@.
smartAckJobSQL :: Bool -> SchemaName -> TableName -> Int64 -> Int64 -> Query Int64
smartAckJobSQL :: Bool -> Text -> Text -> Int64 -> Int64 -> Query Int64
smartAckJobSQL Bool
archiveEnabled Text
schemaName Text
tableName =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      returning :: Text
returning = if Bool
archiveEnabled then Text
"*" else Text
"id, parent_id, group_key" :: Text
      archived :: Text
archived = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
archiveEnabled (Text -> Text -> Text -> Text
archiveAckCte Text
schemaName Text
tableName Text
"ack")
      wakeParent :: Query ()
wakeParent = Text -> Text -> Query ()
wakeParentCte Text
tbl Text
"ack"
      groupLocks :: Text
groupLocks = Text -> Text -> Text -> [Text] -> Text
settleGroupLocksCte Text
schemaName Text
tableName Text
"ack" [Text
"wake_parent"]
   in Int64 -> Int64 -> Query Int64
[stmt|
        WITH ack AS (
          DELETE FROM ${tbl}
          WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8}
            AND NOT EXISTS (SELECT 1 FROM ${tbl} WHERE parent_id = #{jobId :: CInt8})
          RETURNING ${returning}
        ),
        ${archived}
        suspend AS (
          UPDATE ${tbl}
          SET suspended = TRUE, not_visible_until = NULL, claimed_by = NULL,
              attempts = GREATEST(attempts - 1, 0), updated_at = NOW()
          WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8}
            AND NOT EXISTS (SELECT 1 FROM ack)
            AND EXISTS (SELECT 1 FROM ${tbl} WHERE parent_id = #{jobId :: CInt8})
          RETURNING id
        ),
        ${wakeParent},
        ${groupLocks}
        SELECT
          (SELECT count(*) FROM ack) + (SELECT count(*) FROM suspend) AS @{result :: CInt8}
        FROM (SELECT count(*) FROM group_locks) held
      |]

-- | Set-based smart ack over @unnest@ed @(id, claim_seq)@ arrays. Deletes leaves,
-- suspends finalizers that still have children, and wakes parents whose last
-- child completed. The wake check excludes acked children explicitly. Returns the
-- ids acked or suspended. Gone or reclaimed jobs are absent. Locks rows in descending id order.
-- The caller holds the parent locks.
smartAckJobsBatchSQL :: Bool -> SchemaName -> TableName -> [Int64] -> [Int64] -> Query Int64
smartAckJobsBatchSQL :: Bool -> Text -> Text -> [Int64] -> [Int64] -> Query Int64
smartAckJobsBatchSQL Bool
archiveEnabled Text
schemaName Text
tableName =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      returning :: Text
returning = if Bool
archiveEnabled then Text
"job.*" else Text
"job.id, job.parent_id, job.group_key" :: Text
      archived :: Text
archived = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
archiveEnabled (Text -> Text -> Text -> Text
archiveAckCte Text
schemaName Text
tableName Text
"ack")
      groupLocks :: Text
groupLocks = Text -> Text -> Text -> [Text] -> Text
settleGroupLocksCte Text
schemaName Text
tableName Text
"ack" [Text
"suspend", Text
"wake_parent"]
   in [Int64] -> [Int64] -> Query Int64
[stmt|
        WITH input AS (
          SELECT unnest(#{ids :: [CInt8]}::bigint[]) AS id, unnest(#{cseqs :: [CInt8]}::bigint[]) AS cseq
        ),
        locked AS (
          SELECT id FROM ${tbl} WHERE id = ANY(#{ids :: [CInt8]})
          ORDER BY id DESC
          FOR UPDATE
        ),
        ack AS (
          DELETE FROM ${tbl} job
          USING input input_row
          WHERE job.id = input_row.id AND job.claim_seq = input_row.cseq
            AND job.id IN (SELECT id FROM locked)
            AND NOT EXISTS (SELECT 1 FROM ${tbl} child WHERE child.parent_id = job.id)
          RETURNING ${returning}
        ),
        ${archived}
        suspend AS (
          UPDATE ${tbl} job
          SET suspended = TRUE, not_visible_until = NULL, claimed_by = NULL,
              attempts = GREATEST(job.attempts - 1, 0), updated_at = NOW()
          FROM input input_row
          WHERE job.id = input_row.id AND job.claim_seq = input_row.cseq
            AND job.id IN (SELECT id FROM locked)
            AND NOT EXISTS (SELECT 1 FROM ack acked WHERE acked.id = job.id)
            AND EXISTS (SELECT 1 FROM ${tbl} child WHERE child.parent_id = job.id)
          RETURNING job.id, job.group_key
        ),
        wake_parent AS (
          UPDATE ${tbl} parent
          SET suspended = FALSE, updated_at = NOW()
          WHERE parent.id IN (SELECT DISTINCT parent_id FROM ack WHERE parent_id IS NOT NULL)
            AND parent.suspended = TRUE
            AND NOT EXISTS (
              SELECT 1 FROM ${tbl} child
              WHERE child.parent_id = parent.id
                AND NOT EXISTS (SELECT 1 FROM ack acked WHERE acked.id = child.id)
            )
          RETURNING parent.id, parent.group_key
        ),
        ${groupLocks}
        SELECT @{id :: CInt8} FROM (SELECT id FROM ack UNION SELECT id FROM suspend) settled
        CROSS JOIN (SELECT count(*) FROM group_locks) held
      |]

-- | Extend a job's visibility timeout. Matches on the claim token and does not check
-- the holder. Suspended rows hold no lease. @secs@ at or below 0 clears the lease.
setVisibilityTimeoutSQL :: SchemaName -> TableName -> Double -> Int64 -> Int64 -> Query ()
setVisibilityTimeoutSQL :: Text -> Text -> Double -> Int64 -> Int64 -> Query ()
setVisibilityTimeoutSQL Text
schemaName Text
tableName Double
secs Int64
jobId Int64
cseq =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      leaseUntil :: Query ()
leaseUntil = Double -> Query ()
leaseDeadline Double
secs
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = ${leaseUntil},
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
      |]

-- | 'setVisibilityTimeoutSQL' over a batch, for the heartbeat. Extends every job still
-- under this claim, held by the same worker and unsuspended, and reports per row whether
-- the update landed alongside its claim token, holder, cancel flag and suspension.
-- @secs@ at or below 0 clears the lease. The statement never waits on a lock. It takes
-- the group summaries first, then the rows, both with SKIP LOCKED, and a row whose
-- summary or row is busy reads back unchanged. The caller attaches the row decoder.
setVisibilityTimeoutBatchSQL :: SchemaName -> TableName -> [JobRead payload] -> Double -> Query ()
setVisibilityTimeoutBatchSQL :: forall payload.
Text -> Text -> [JobRead payload] -> Double -> Query ()
setVisibilityTimeoutBatchSQL Text
schemaName Text
tableName [JobRead payload]
jobs Double
secs =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      leaseUntil :: Query ()
leaseUntil = Double -> Query ()
leaseDeadline Double
secs
      ids :: [Int64]
ids = (JobRead payload -> Int64) -> [JobRead payload] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey [JobRead payload]
jobs
      valuesFrag :: Query ()
valuesFrag = Text -> [Query ()] -> Query ()
sepBy Text
"," ((JobRead payload -> Query ()) -> [JobRead payload] -> [Query ()]
forall a b. (a -> b) -> [a] -> [b]
map JobRead payload -> Query ()
forall {payload} {q} {insertedAt} {adm}.
Job payload Int64 q insertedAt adm -> Query ()
inputRow [JobRead payload]
jobs)
      inputRow :: Job payload Int64 q insertedAt adm -> Query ()
inputRow Job payload Int64 q insertedAt adm
job =
        let jid :: Int64
jid = Job payload Int64 q insertedAt adm -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey Job payload Int64 q insertedAt adm
job
            cseq :: Int64
cseq = Job payload Int64 q insertedAt adm -> Int64
forall payload q insertedAt adm.
Job payload Int64 q insertedAt adm -> Int64
claimSeq Job payload Int64 q insertedAt adm
job
            holder :: Maybe UUID
holder = Job payload Int64 q insertedAt adm -> Maybe UUID
forall payload q insertedAt adm.
Job payload Int64 q insertedAt adm -> Maybe UUID
claimedBy Job payload Int64 q insertedAt adm
job
         in [sql|(#{jid :: CInt8}, #{cseq :: CInt8}, #{holder :: Maybe CUuid})|]
   in [sql|
        WITH input_jobs AS (
          SELECT input_values.id::bigint AS id,
                 input_values.expected_claim_seq::bigint AS expected_claim_seq,
                 input_values.expected_claimed_by::uuid AS expected_claimed_by
          FROM (VALUES ${valuesFrag}) AS input_values(id, expected_claim_seq, expected_claimed_by)
        ),
        wanted_groups AS MATERIALIZED (
          SELECT DISTINCT group_key FROM ${tbl}
          WHERE id = ANY(#{ids :: [CInt8]}) AND group_key IS NOT NULL
        ),
        held_groups AS (
          SELECT group_key FROM ${groupsTbl}
          WHERE group_key IN (SELECT group_key FROM wanted_groups)
          ORDER BY group_key
          FOR UPDATE SKIP LOCKED
        ),
        busy_groups AS MATERIALIZED (
          SELECT group_key FROM ${groupsTbl}
          WHERE group_key IN (SELECT group_key FROM wanted_groups)
          EXCEPT
          SELECT group_key FROM held_groups
        ),
        locked AS (
          SELECT id FROM ${tbl} job
          WHERE job.id = ANY(#{ids :: [CInt8]})
            AND (job.group_key IS NULL OR NOT EXISTS (SELECT 1 FROM busy_groups busy WHERE busy.group_key = job.group_key))
          ORDER BY id DESC
          FOR UPDATE SKIP LOCKED
        ),
        updated AS (
          UPDATE ${tbl} job
          SET not_visible_until = ${leaseUntil},
              updated_at = NOW()
          FROM input_jobs input_job
          WHERE job.id = input_job.id AND job.id IN (SELECT id FROM locked)
            AND job.claim_seq = input_job.expected_claim_seq AND NOT job.suspended
            AND job.claimed_by IS NOT DISTINCT FROM input_job.expected_claimed_by
          RETURNING job.id
        )
        SELECT
          input_job.id,
          (heartbeated.id IS NOT NULL) as was_heartbeated,
          job.claim_seq as current_db_claim_seq,
          (job.cancel_requested_at IS NOT NULL) as cancel_requested,
          (job.suspended IS TRUE) as suspended,
          job.claimed_by as claimed_by
        FROM input_jobs input_job
        LEFT JOIN updated heartbeated ON heartbeated.id = input_job.id
        LEFT JOIN ${tbl} job ON job.id = input_job.id
      |]

-- | Park a failed job for its retry backoff. Matches on the claim token.
updateJobForRetrySQL :: SchemaName -> TableName -> Int64 -> Text -> Int64 -> Int64 -> Query ()
updateJobForRetrySQL :: Text -> Text -> Int64 -> Text -> Int64 -> Int64 -> Query ()
updateJobForRetrySQL Text
schemaName Text
tableName Int64
backoff Text
errorMsg Int64
jobId Int64
cseq =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = NOW() + (#{backoff :: CInt8} * interval '1 second'),
            last_error = #{errorMsg :: CText},
            updated_at = NOW(),
            claimed_by = NULL
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
      |]

-- | Soft nack. Releases the claim and hands back the attempt it consumed, recording no
-- failure. A delay in seconds parks the job for that long. Without one,
-- @not_visible_until@ stands.
nackJobSQL :: SchemaName -> TableName -> Maybe Int64 -> Int64 -> Int64 -> Int32 -> Query ()
nackJobSQL :: Text -> Text -> Maybe Int64 -> Int64 -> Int64 -> Int32 -> Query ()
nackJobSQL Text
schemaName Text
tableName Maybe Int64
delay Int64
jobId Int64
cseq Int32
att =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET attempts = LEAST(GREATEST(#{att :: CInt4} - 1, attempts - 1, 0), attempts),
            not_visible_until = COALESCE(NOW() + (#{delay :: Maybe CInt8}::bigint * interval '1 second'), not_visible_until),
            claimed_by = NULL,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8} AND claim_seq = #{cseq :: CInt8} AND NOT suspended
          AND claimed_by IS NOT NULL
      |]

-- | 'nackJobSQL' over @unnest@ed @(id, claim_seq, attempts)@ arrays. Locks rows in
-- descending id order. Returns the ids nacked.
nackJobsBatchSQL :: SchemaName -> TableName -> Maybe Int64 -> [Int64] -> [Int64] -> [Int32] -> Query Int64
nackJobsBatchSQL :: Text
-> Text
-> Maybe Int64
-> [Int64]
-> [Int64]
-> [Int32]
-> Query Int64
nackJobsBatchSQL Text
schemaName Text
tableName Maybe Int64
delay [Int64]
ids [Int64]
cseqs [Int32]
atts =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      locked :: Query ()
locked = Text -> [Int64] -> Query ()
lockedByIdsCte Text
tbl [Int64]
ids
   in [sql|
        WITH input AS (
          SELECT unnest(#{ids :: [CInt8]}::bigint[]) AS in_id,
                 unnest(#{cseqs :: [CInt8]}::bigint[]) AS cseq,
                 unnest(#{atts :: [CInt4]}::int[]) AS att
        ),
        ${locked}
        UPDATE ${tbl} job
        SET attempts = LEAST(GREATEST(input_row.att - 1, job.attempts - 1, 0), job.attempts),
            not_visible_until = COALESCE(NOW() + (#{delay :: Maybe CInt8}::bigint * interval '1 second'), job.not_visible_until),
            claimed_by = NULL,
            updated_at = NOW()
        FROM input input_row
        WHERE job.id = input_row.in_id AND job.id IN (SELECT id FROM locked)
          AND job.claim_seq = input_row.cseq AND NOT job.suspended
          AND job.claimed_by IS NOT NULL
        RETURNING @{id :: CInt8}
      |]

-- | Make a delayed or retrying job immediately visible. Refuses an in-flight or suspended job.
promoteJobSQL :: SchemaName -> TableName -> Int64 -> Query ()
promoteJobSQL :: Text -> Text -> Int64 -> Query ()
promoteJobSQL Text
schemaName Text
tableName Int64
jobId =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = NULL,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8}
          AND NOT suspended
          AND not_visible_until IS NOT NULL
          AND not_visible_until > NOW()
          AND claimed_by IS NULL
      |]

-- | Set when a job next becomes visible, clear its throttle marker and void a lapsed
-- claim. Refuses an in-flight, suspended, force-cancel-flagged or exhausted job.
rescheduleJobSQL :: SchemaName -> TableName -> Int64 -> UTCTime -> Query ()
rescheduleJobSQL :: Text -> Text -> Int64 -> UTCTime -> Query ()
rescheduleJobSQL Text
schemaName Text
tableName Int64
jobId UTCTime
runAt =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      attemptsLeft :: Text
attemptsLeft = Text -> Text
attemptsLeftSQL Text
""
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = CASE WHEN #{runAt :: CTimestamptz}::timestamptz > NOW() THEN #{runAt :: CTimestamptz}::timestamptz END,
            throttled_until = NULL,
            claimed_by = NULL,
            claim_seq = claim_seq + 1,
            updated_at = NOW()
        WHERE id = #{jobId :: CInt8}
          AND NOT suspended
          AND (claimed_by IS NULL OR not_visible_until <= NOW())
          AND cancel_requested_at IS NULL
          AND ${attemptsLeft}
      |]

-- | The lease deadline @secs@ from now, or NULL when @secs@ is at or below 0.
leaseDeadline :: Double -> Query ()
leaseDeadline :: Double -> Query ()
leaseDeadline Double
secs =
  [sql|CASE WHEN #{secs :: CFloat8}::double precision <= 0 THEN NULL ELSE NOW() + (#{secs :: CFloat8}::double precision * interval '1 second') END|]