Dlq.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Dead-letter queue inspection and redrive helpers.
--
-- The PGMQ adapter writes DLQ rows as a shibuya wrapper whose required keys are
-- @original_message@ and @dead_letter_reason@. Keiro's worker path always includes
-- metadata as well: @original_message_id@, @original_enqueued_at@, @last_read_at@,
-- @read_count@, and @original_headers@. These helpers parse the required keys and
-- treat metadata as optional so operators can still inspect legacy or hand-written
-- rows.
--
-- PGMQ does not expire DLQ rows by itself. The retention model is "archive-then-
-- purge": 'archiveDlq' retains dead letters by moving them into the archive table
-- @pgmq.a_<dlq>@ (preserving @enqueued_at@ / @read_ct@ and stamping @archived_at@)
-- for audit, while 'purgeDlq' deletes them permanently. An operator who needs an
-- audit trail runs 'archiveDlq' (retain) and may then 'purgeDlq' (clear the active
-- table); an operator who does not keeps using 'purgeDlq' alone. Either way, alert
-- on the DLQ's depth via 'Keiro.PGMQ.Metrics.jobDlqMetrics'.
--
-- Redrive is at-least-once: a crash after sending the original payload back to the
-- main queue but before deleting the DLQ row leaves the payload in both places.
-- Handlers must therefore be idempotent.
module Keiro.PGMQ.Dlq
  ( DlqEntry (..),
    readDlq,
    redriveDlq,
    purgeDlq,
    archiveDlq,
    archiveDlqEntry,
  )
where

import Keiro.PGMQ.Codec (JobDecodeError (..), decodeJob)
import Keiro.PGMQ.Job (Job (..))
import Keiro.PGMQ.Runtime (QueueRef (..))
import "aeson" Data.Aeson (Value (..), withObject, (.:), (.:?))
import "aeson" Data.Aeson.Types (Parser, parseEither)
import "base" Control.Monad (foldM, void)
import "base" Data.Foldable (toList)
import "base" Data.Int (Int32, Int64)
import "effectful-core" Effectful (Eff, IOE, (:>))
import "pgmq-effectful" Pgmq.Effectful
  ( Message (..),
    MessageBody (..),
    MessageId,
    MessageQuery (..),
    Pgmq,
    ReadMessage (..),
    SendMessage (..),
  )
import "pgmq-effectful" Pgmq.Effectful qualified as Pgmq
import "text" Data.Text (Text)
import "text" Data.Text qualified as Text
import "time" Data.Time (UTCTime)

-- | A decoded dead-letter entry: shibuya's DLQ wrapper, unwrapped.
data DlqEntry p = DlqEntry
  { -- | The DLQ row's own PGMQ message id.
    dlqMessageId :: !MessageId,
    -- | @poison_pill: ...@, @invalid_payload: ...@, or @max_retries_exceeded@.
    reason :: !Text,
    -- | The preserved @original_message@ decoded with the job's codec.
    originalPayload :: !(Either JobDecodeError p),
    originalMessageId :: !(Maybe Int64),
    originalEnqueuedAt :: !(Maybe UTCTime),
    readCount :: !(Maybe Int64),
    -- | Full DLQ wrapper for forensics.
    rawBody :: !Value
  }
  deriving stock (Show)

data DlqEnvelope = DlqEnvelope
  { originalMessage :: !Value,
    deadLetterReason :: !Text,
    envelopeOriginalMessageId :: !(Maybe Int64),
    envelopeOriginalEnqueuedAt :: !(Maybe UTCTime),
    envelopeReadCount :: !(Maybe Int64)
  }

parseDlqEnvelope :: Value -> Either Text DlqEnvelope
parseDlqEnvelope =
  firstText . parseEither parser
  where
    parser :: Value -> Parser DlqEnvelope
    parser =
      withObject "DLQ payload" \obj -> do
        originalMessage <- obj .: "original_message"
        deadLetterReason <- obj .: "dead_letter_reason"
        envelopeOriginalMessageId <- obj .:? "original_message_id"
        envelopeOriginalEnqueuedAt <- obj .:? "original_enqueued_at"
        envelopeReadCount <- obj .:? "read_count"
        pure
          DlqEnvelope
            { originalMessage,
              deadLetterReason,
              envelopeOriginalMessageId,
              envelopeOriginalEnqueuedAt,
              envelopeReadCount
            }

    firstText = \case
      Left err -> Left (Text.pack err)
      Right value -> Right value

-- | Read and decode up to @n@ DLQ entries. The read uses a 30 second visibility
-- timeout so concurrent inspections do not immediately see the same rows.
readDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int32 -> Eff es [DlqEntry p]
readDlq job n
  | n <= 0 = pure []
  | otherwise = do
      messages <-
        Pgmq.readMessage
          ReadMessage
            { queueName = job.jobQueue.dlqName,
              delay = 30,
              batchSize = Just n,
              conditional = Nothing
            }
      pure (fmap (toEntry job) (toList messages))

toEntry :: Job p -> Message -> DlqEntry p
toEntry job message =
  let body = unMessageBody message.body
   in case parseDlqEnvelope body of
        Left err ->
          DlqEntry
            { dlqMessageId = message.messageId,
              reason = "malformed_dlq_payload: " <> err,
              originalPayload = Left (JobPayloadMalformed err),
              originalMessageId = Nothing,
              originalEnqueuedAt = Nothing,
              readCount = Nothing,
              rawBody = body
            }
        Right envelope ->
          DlqEntry
            { dlqMessageId = message.messageId,
              reason = envelope.deadLetterReason,
              originalPayload = decodeJob job.jobCodec envelope.originalMessage,
              originalMessageId = envelope.envelopeOriginalMessageId,
              originalEnqueuedAt = envelope.envelopeOriginalEnqueuedAt,
              readCount = envelope.envelopeReadCount,
              rawBody = body
            }

-- | Move up to @n@ DLQ rows back to the main queue. Redriven messages start a
-- fresh PGMQ @read_ct@ on the main queue. Malformed DLQ wrappers are left in the
-- DLQ for inspection.
redriveDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int -> Eff es Int
redriveDlq job n
  | n <= 0 = pure 0
  | otherwise = loop 0
  where
    loop moved
      | moved >= n = pure moved
      | otherwise = do
          messages <-
            Pgmq.readMessage
              ReadMessage
                { queueName = job.jobQueue.dlqName,
                  delay = 30,
                  batchSize = Just (fromIntegral (min 100 (n - moved))),
                  conditional = Nothing
                }
          if null messages
            then pure moved
            else do
              movedInBatch <- foldM redriveOne 0 messages
              if movedInBatch == 0
                then pure moved
                else loop (moved + movedInBatch)

    redriveOne count message =
      case parseDlqEnvelope (unMessageBody message.body) of
        Left _err ->
          pure count
        Right envelope -> do
          _ <-
            Pgmq.sendMessage
              SendMessage
                { queueName = job.jobQueue.physicalName,
                  messageBody = MessageBody envelope.originalMessage,
                  delay = Nothing
                }
          void $
            Pgmq.deleteMessage
              MessageQuery
                { queueName = job.jobQueue.dlqName,
                  messageId = message.messageId
                }
          pure (count + 1)

-- | Delete all rows currently in the DLQ.
purgeDlq :: (Pgmq :> es, IOE :> es) => Job p -> Eff es ()
purgeDlq job =
  void (Pgmq.deleteAllMessagesFromQueue job.jobQueue.dlqName)

-- | Archive (retain) up to @n@ DLQ rows: move each out of the active DLQ table
-- @pgmq.q_<dlq>@ into the archive table @pgmq.a_<dlq>@, preserving @enqueued_at@ /
-- @read_ct@ and stamping @archived_at@. Returns the number archived. This is the
-- audit-retention counterpart to the delete-only 'purgeDlq'. At-most-once per row
-- per call; a crash before archiving leaves the row in the active DLQ for a re-run.
archiveDlq :: (Pgmq :> es, IOE :> es) => Job p -> Int -> Eff es Int
archiveDlq job n
  | n <= 0 = pure 0
  | otherwise = loop 0
  where
    loop archived
      | archived >= n = pure archived
      | otherwise = do
          messages <-
            Pgmq.readMessage
              ReadMessage
                { queueName = job.jobQueue.dlqName,
                  delay = 30,
                  batchSize = Just (fromIntegral (min 100 (n - archived))),
                  conditional = Nothing
                }
          if null messages
            then pure archived
            else do
              archivedInBatch <- foldM archiveOne 0 messages
              if archivedInBatch == 0
                then pure archived
                else loop (archived + archivedInBatch)

    archiveOne count message = do
      moved <- archiveDlqEntry job message.messageId
      pure (if moved then count + 1 else count)

-- | Archive one specific DLQ row by message id. 'True' if a row was moved.
archiveDlqEntry :: (Pgmq :> es, IOE :> es) => Job p -> MessageId -> Eff es Bool
archiveDlqEntry job msgId =
  Pgmq.archiveMessage
    MessageQuery
      { queueName = job.jobQueue.dlqName,
        messageId = msgId
      }