ReplayAudit.hs

{-# LANGUAGE GADTs #-}

{- | Read-only, differential replay audits for aggregate event streams.

Routine deploys should first consume @keiro-dsl diff@'s replay-impact verdict:
a replay-neutral deploy touches no data, while an affected verdict supplies the
event types for 'AuditTargeted'. 'AuditFull' is intentionally reserved for
one-time runtime cutovers and forensics.

Selection is read-only and server-side. It uses Kiroku's indexed category,
event-type, and global-position schema to discover only candidate stream names;
each candidate is then hydrated through the public Store effect. The audit
never appends events, calls @verifyAndSnapshot@, or writes snapshots.

Hand-written services have no spec from which to derive an affected set. They
must supply a conservative set explicitly or choose 'AuditFull'.
-}
module Keiro.ReplayAudit (
    AuditMode (..),
    AffectedSet (..),
    AuditBudget (..),
    defaultAuditBudget,
    AuditTarget (..),
    SomeAuditTarget (..),
    streamInCategory,
    AuditOutcome (..),
    StreamAuditResult (..),
    AuditReport (..),
    auditStream,
    auditStreams,
    auditTargets,
    renderAuditReport,
    auditExitCode,
) where

import Contravariant.Extras (contrazip3, contrazip5)
import Data.Int (Int32)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text qualified as Text
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Effectful (Eff, IOE, (:>))
import Effectful.Concurrent (runConcurrent)
import Effectful.Concurrent.Async qualified as Async
import Hasql.Decoders qualified as D
import Hasql.Encoders qualified as E
import Hasql.Statement (Statement, preparable)
import Keiki.Core (BoolAlg, RegFile)
import Keiro.Command (
    CommandError,
    Hydrated (..),
    RunCommandOptions,
    defaultRunCommandOptions,
    hydrateFull,
    hydrateSeeded,
 )
import Keiro.EventStream (EventStream, StateCodec)
import Keiro.EventStream.Validate (ValidatedEventStream, unvalidated)
import Keiro.Prelude
import Keiro.ReplayDigest (canonicalJsonBytes, replayDigest)
import Keiro.Snapshot (SnapshotLookup (..), lookupSnapshotSeed)
import Keiro.Stream (Stream (..))
import Kiroku.Store.Effect (Store)
import Kiroku.Store.Read (lookupStreamNames)
import Kiroku.Store.Transaction (runTransaction)
import Kiroku.Store.Types (
    CategoryName (..),
    EventType (..),
    GlobalPosition (..),
    StreamId (..),
    StreamName (..),
    StreamVersion,
 )
import Kiroku.Store.Types qualified as StoreTypes
import "hasql-transaction" Hasql.Transaction qualified as Tx
import Prelude qualified

-- | Which streams to inspect.
data AuditMode
    = AuditFull
    | AuditTargeted !AffectedSet
    deriving stock (Eq, Show, Generic)

-- | Conservative stored-data surface emitted by the DSL replay-impact diff.
data AffectedSet = AffectedSet
    { affectedEventTypes :: !(Set EventType)
    , includeSnapshotStreams :: !Bool
    }
    deriving stock (Eq, Show, Generic)

{- | Per-run cost and resume controls.

The checkpoint is the maximum global position of the last audited stream at
the selection snapshot. Re-running with that checkpoint selects only streams
whose latest event lies after it. A category that receives new events during
or after a run may intentionally cause a previously checked stream to be
selected again: its durable history changed and warrants another verdict.
-}
data AuditBudget = AuditBudget
    { maxStreams :: !(Maybe Int)
    , parallelism :: !Int
    , resumeFrom :: !(Maybe GlobalPosition)
    }
    deriving stock (Eq, Show, Generic)

defaultAuditBudget :: AuditBudget
defaultAuditBudget =
    AuditBudget
        { maxStreams = Nothing
        , parallelism = 4
        , resumeFrom = Nothing
        }

-- | One typed aggregate/category assembly.
data AuditTarget phi rs s ci co = AuditTarget
    { eventStream :: !(ValidatedEventStream phi rs s ci co)
    , category :: !Text
    , mkStream :: !(StreamName -> Maybe (Stream (EventStream phi rs s ci co)))
    }
    deriving stock (Generic)

-- | Existential packaging for services with multiple aggregate types.
data SomeAuditTarget where
    SomeAuditTarget ::
        (BoolAlg phi (RegFile rs, ci), Eq co) =>
        AuditTarget phi rs s ci co ->
        SomeAuditTarget

-- | Accept a raw store name only when it belongs to the expected category.
streamInCategory :: Text -> StreamName -> Maybe (Stream eventStream)
streamInCategory expected streamName =
    case StoreTypes.categoryName streamName of
        CategoryName actual
            | actual == expected -> Just (Stream streamName)
        _ -> Nothing

-- | Replay result for one accepted stream name.
data AuditOutcome
    = ReplayOk
        { streamVersion :: !StreamVersion
        , digest :: !(Maybe Text)
        }
    | ReplayFailed
        { commandError :: !CommandError
        }
    | SeedDivergence
        { seedVersion :: !StreamVersion
        , seededDigest :: !Text
        , fullDigest :: !Text
        }
    deriving stock (Eq, Show, Generic)

data StreamAuditResult = StreamAuditResult
    { streamName :: !StreamName
    , outcome :: !AuditOutcome
    }
    deriving stock (Eq, Show, Generic)

data AuditReport = AuditReport
    { targetCategory :: !Text
    , mode :: !Text
    , results :: ![StreamAuditResult]
    , rejectedStreams :: ![StreamName]
    , streamsSelected :: !Int
    , streamsSkipped :: !Int
    , failures :: !Int
    , divergences :: !Int
    , checkpoint :: !(Maybe GlobalPosition)
    }
    deriving stock (Eq, Show, Generic)

auditStream ::
    forall phi rs s ci co es.
    (Store :> es, BoolAlg phi (RegFile rs, ci), Eq co) =>
    AuditTarget phi rs s ci co ->
    Stream (EventStream phi rs s ci co) ->
    Eff es AuditOutcome
auditStream target stream = do
    let raw = unvalidated (target ^. #eventStream)
        options = defaultRunCommandOptions
    hydrateFull options raw stream >>= \case
        Left err -> pure (ReplayFailed err)
        Right full ->
            case raw ^. #stateCodec of
                Nothing ->
                    pure
                        ReplayOk
                            { streamVersion = full ^. #streamVersion
                            , digest = Nothing
                            }
                Just codec -> auditSeed options raw codec full
  where
    auditSeed ::
        RunCommandOptions ->
        EventStream phi rs s ci co ->
        StateCodec (s, RegFile rs) ->
        Hydrated rs s ->
        Eff es AuditOutcome
    auditSeed options raw codec full = do
        let name = (raw ^. #resolveStreamName) stream
            fullValue = (codec ^. #encode) (full ^. #state, full ^. #registers)
            fullBytes = canonicalJsonBytes fullValue
            fullHash = replayDigest fullValue
        lookupSnapshotSeed name codec >>= \case
            SnapshotUnavailable _ ->
                pure
                    ReplayOk
                        { streamVersion = full ^. #streamVersion
                        , digest = Just fullHash
                        }
            SnapshotHit seed ->
                hydrateSeeded
                    options
                    raw
                    stream
                    (seed ^. #state)
                    (seed ^. #registers)
                    (seed ^. #streamVersion)
                    >>= \case
                        Left err ->
                            pure
                                SeedDivergence
                                    { seedVersion = seed ^. #streamVersion
                                    , seededDigest = "replay-failed:" <> Text.pack (show err)
                                    , fullDigest = fullHash
                                    }
                        Right seeded ->
                            let seededValue =
                                    (codec ^. #encode)
                                        (seeded ^. #state, seeded ^. #registers)
                                seededHash = replayDigest seededValue
                             in if canonicalJsonBytes seededValue == fullBytes
                                    then
                                        pure
                                            ReplayOk
                                                { streamVersion = full ^. #streamVersion
                                                , digest = Just fullHash
                                                }
                                    else
                                        pure
                                            SeedDivergence
                                                { seedVersion = seed ^. #streamVersion
                                                , seededDigest = seededHash
                                                , fullDigest = fullHash
                                                }

auditStreams ::
    forall phi rs s ci co es.
    (IOE :> es, Store :> es, BoolAlg phi (RegFile rs, ci), Eq co) =>
    AuditMode ->
    AuditBudget ->
    AuditTarget phi rs s ci co ->
    Eff es AuditReport
auditStreams auditMode budget target = do
    totalStreams <- categoryStreamCount (target ^. #category)
    selectedTotal <- selectedStreamCount auditMode (target ^. #category)
    let skipped =
            case auditMode of
                AuditFull -> 0
                AuditTargeted _ -> Prelude.max 0 (totalStreams Prelude.- selectedTotal)
        workerCount = Prelude.max 1 (budget ^. #parallelism)
        remaining0 = fmap (Prelude.max 0) (budget ^. #maxStreams)
        cursor0 = fromMaybe (GlobalPosition 0) (budget ^. #resumeFrom)
    (allResults, rejected, selected, finalCursor) <-
        go workerCount remaining0 cursor0 [] [] 0
    let failureCount =
            Prelude.length rejected
                Prelude.+ Prelude.length
                    [ ()
                    | StreamAuditResult{outcome = ReplayFailed{}} <- allResults
                    ]
        divergenceCount =
            Prelude.length
                [ ()
                | StreamAuditResult{outcome = SeedDivergence{}} <- allResults
                ]
    pure
        AuditReport
            { targetCategory = target ^. #category
            , mode = case auditMode of AuditFull -> "full"; AuditTargeted _ -> "targeted"
            , results = allResults
            , rejectedStreams = rejected
            , streamsSelected = selected
            , streamsSkipped = skipped
            , failures = failureCount
            , divergences = divergenceCount
            , checkpoint =
                if selected == 0
                    then budget ^. #resumeFrom
                    else Just finalCursor
            }
  where
    pageSize = 128

    go ::
        Int ->
        Maybe Int ->
        GlobalPosition ->
        [StreamAuditResult] ->
        [StreamName] ->
        Int ->
        Eff es ([StreamAuditResult], [StreamName], Int, GlobalPosition)
    go workerCount remaining cursor resultAcc rejectedAcc selectedAcc
        | Just 0 <- remaining =
            pure
                ( Prelude.reverse resultAcc
                , Prelude.reverse rejectedAcc
                , selectedAcc
                , cursor
                )
        | otherwise = do
            let requestSize =
                    maybe pageSize (Prelude.min pageSize) remaining
            page <- selectStreamPage auditMode (target ^. #category) cursor requestSize
            if Vector.null page
                then
                    pure
                        ( Prelude.reverse resultAcc
                        , Prelude.reverse rejectedAcc
                        , selectedAcc
                        , cursor
                        )
                else do
                    names <- lookupStreamNames (Prelude.fst <$> Vector.toList page)
                    let resolved =
                            [ (streamName, streamValue)
                            | (streamId, _watermark) <- Vector.toList page
                            , Just streamName <- [Map.lookup streamId names]
                            , Just streamValue <- [(target ^. #mkStream) streamName]
                            ]
                        acceptedNames = Set.fromList (Prelude.fst <$> resolved)
                        pageNames =
                            [ streamName
                            | (streamId, _watermark) <- Vector.toList page
                            , Just streamName <- [Map.lookup streamId names]
                            ]
                        rejected =
                            Prelude.filter (`Set.notMember` acceptedNames) pageNames
                    audited <-
                        runConcurrent
                            $ Async.pooledMapConcurrentlyN
                                workerCount
                                ( \(name, streamValue) ->
                                    StreamAuditResult name <$> auditStream target streamValue
                                )
                                resolved
                    let nextCursor = Prelude.snd (Vector.last page)
                        pageCount = Vector.length page
                        nextRemaining = (Prelude.- pageCount) <$> remaining
                    go
                        workerCount
                        nextRemaining
                        nextCursor
                        (Prelude.reverse audited <> resultAcc)
                        (Prelude.reverse rejected <> rejectedAcc)
                        (selectedAcc Prelude.+ pageCount)

auditTargets ::
    (IOE :> es, Store :> es) =>
    AuditMode ->
    AuditBudget ->
    [SomeAuditTarget] ->
    Eff es [AuditReport]
auditTargets auditMode budget =
    traverse $ \(SomeAuditTarget target) -> auditStreams auditMode budget target

renderAuditReport :: AuditReport -> Text
renderAuditReport report =
    Text.intercalate
        " "
        [ "replay-audit"
        , "category=" <> report ^. #targetCategory
        , "mode=" <> report ^. #mode
        , "selected=" <> textShow (report ^. #streamsSelected)
        , "skipped=" <> textShow (report ^. #streamsSkipped)
        , "failures=" <> textShow (report ^. #failures)
        , "divergences=" <> textShow (report ^. #divergences)
        , "checkpoint=" <> maybe "none" textShow (report ^. #checkpoint)
        ]
  where
    textShow :: (Show a) => a -> Text
    textShow = Text.pack . show

auditExitCode :: [AuditReport] -> Int
auditExitCode reports
    | Prelude.any
        (\report -> report ^. #failures > 0 Prelude.|| report ^. #divergences > 0)
        reports =
        1
    | otherwise = 0

-- Selection -----------------------------------------------------------------

categoryStreamCount :: (Store :> es) => Text -> Eff es Int
categoryStreamCount category =
    Prelude.fromIntegral <$> runTransaction (Tx.statement category categoryStreamCountStmt)

selectedStreamCount :: (Store :> es) => AuditMode -> Text -> Eff es Int
selectedStreamCount auditMode category =
    case auditMode of
        AuditFull -> categoryStreamCount category
        AuditTargeted affected ->
            Prelude.fromIntegral
                <$> runTransaction
                    ( Tx.statement
                        (category, eventTypeTexts affected, affected ^. #includeSnapshotStreams)
                        targetedStreamCountStmt
                    )

selectStreamPage ::
    (Store :> es) =>
    AuditMode ->
    Text ->
    GlobalPosition ->
    Int ->
    Eff es (Vector (StreamId, GlobalPosition))
selectStreamPage auditMode category (GlobalPosition cursor) limit =
    case auditMode of
        AuditFull ->
            runTransaction
                $ Tx.statement
                    (category, cursor, Prelude.fromIntegral limit)
                    fullStreamPageStmt
        AuditTargeted affected ->
            runTransaction
                $ Tx.statement
                    ( category
                    , eventTypeTexts affected
                    , affected ^. #includeSnapshotStreams
                    , cursor
                    , Prelude.fromIntegral limit
                    )
                    targetedStreamPageStmt

eventTypeTexts :: AffectedSet -> Vector Text
eventTypeTexts affected =
    Vector.fromList
        [ text
        | EventType text <- Set.toAscList (affected ^. #affectedEventTypes)
        ]

categoryStreamCountStmt :: Statement Text Int64
categoryStreamCountStmt =
    preparable
        """
        SELECT count(*)
        FROM kiroku.streams
        WHERE category = $1
          AND stream_version > 0
        """
        (E.param (E.nonNullable E.text))
        (D.singleRow (D.column (D.nonNullable D.int8)))

fullStreamPageStmt :: Statement (Text, Int64, Int32) (Vector (StreamId, GlobalPosition))
fullStreamPageStmt =
    preparable
        """
        SELECT s.stream_id, max(all_events.stream_version) AS watermark
        FROM kiroku.streams s
        JOIN kiroku.stream_events all_events
          ON all_events.stream_id = 0
         AND all_events.original_stream_id = s.stream_id
        WHERE s.category = $1
        GROUP BY s.stream_id
        HAVING max(all_events.stream_version) > $2
        ORDER BY watermark ASC
        LIMIT $3
        """
        ( contrazip3
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable E.int8))
            (E.param (E.nonNullable E.int4))
        )
        streamPageDecoder

targetedStreamCountStmt :: Statement (Text, Vector Text, Bool) Int64
targetedStreamCountStmt =
    preparable
        (targetedSelectionCte <> "SELECT count(*) FROM selected")
        targetedSelectionEncoder
        (D.singleRow (D.column (D.nonNullable D.int8)))

targetedStreamPageStmt ::
    Statement
        (Text, Vector Text, Bool, Int64, Int32)
        (Vector (StreamId, GlobalPosition))
targetedStreamPageStmt =
    preparable
        ( targetedSelectionCte
            <> """
               SELECT stream_id, watermark
               FROM selected
               WHERE watermark > $4
               ORDER BY watermark ASC
               LIMIT $5
               """
        )
        ( contrazip5
            (E.param (E.nonNullable E.text))
            (E.param (E.nonNullable (E.foldableArray (E.nonNullable E.text))))
            (E.param (E.nonNullable E.bool))
            (E.param (E.nonNullable E.int8))
            (E.param (E.nonNullable E.int4))
        )
        streamPageDecoder

targetedSelectionEncoder :: E.Params (Text, Vector Text, Bool)
targetedSelectionEncoder =
    contrazip3
        (E.param (E.nonNullable E.text))
        (E.param (E.nonNullable (E.foldableArray (E.nonNullable E.text))))
        (E.param (E.nonNullable E.bool))

targetedSelectionCte :: Text
targetedSelectionCte =
    """
    WITH candidate_streams AS (
      SELECT DISTINCT s.stream_id
      FROM kiroku.streams s
      JOIN kiroku.stream_events source_events
        ON source_events.stream_id = 0
       AND source_events.original_stream_id = s.stream_id
      JOIN kiroku.events e ON e.event_id = source_events.event_id
      WHERE s.category = $1
        AND e.event_type = ANY($2::text[])
      UNION
      SELECT s.stream_id
      FROM kiroku.streams s
      JOIN keiro.keiro_snapshots snapshots ON snapshots.stream_id = s.stream_id
      WHERE $3
        AND s.category = $1
    ),
    selected AS (
      SELECT candidates.stream_id, max(all_events.stream_version) AS watermark
      FROM candidate_streams candidates
      JOIN kiroku.stream_events all_events
        ON all_events.stream_id = 0
       AND all_events.original_stream_id = candidates.stream_id
      GROUP BY candidates.stream_id
    )
    """

streamPageDecoder :: D.Result (Vector (StreamId, GlobalPosition))
streamPageDecoder =
    D.rowVector
        $ (,)
        <$> (StreamId <$> D.column (D.nonNullable D.int8))
        <*> (GlobalPosition <$> D.column (D.nonNullable D.int8))