------------------------------ MODULE RoostApplyBatch ------------------------------
(***************************************************************************)
(* A batch of record writes over several single-writer shards, where a     *)
(* transaction cannot span shards.                                         *)
(*                                                                         *)
(* A batch is a sequence of ops (create or delete of an id on a shard).    *)
(* Some op may be rejected by validation or a before-trigger; that outcome *)
(* is a property of the op, so it is the same on every attempt.           *)
(*                                                                         *)
(* The design under test runs in three phases:                             *)
(*   Plan    walk the ops in order, deciding what each would write; the    *)
(*           first rejection ends the batch with nothing written           *)
(*   Commit  one transaction per shard, applying all of that shard's ops   *)
(*   Finish  audit each committed op                                       *)
(* A crash may hit between any two steps; a retry re-runs the same batch   *)
(* against whatever landed.  Creates and deletes are idempotent (a create  *)
(* of an existing id is a duplicate, a delete of a missing id is missing), *)
(* which is what makes the retry converge.                                 *)
(***************************************************************************)

EXTENDS Naturals, FiniteSets

CONSTANTS
  Shards,  \* the shards a batch may touch
  Ids,     \* record ids
  N        \* ops per batch

Kinds == {"create", "delete"}
Ops   == [kind : Kinds, shard : Shards, id : Ids]
Idx   == 1..N
Phases == {"plan", "commit", "finish", "done", "rejected", "crashed"}

VARIABLES
  ops,        \* Idx -> Ops, the batch
  rejects,    \* Idx -> BOOLEAN, which ops validation refuses
  db0,        \* Shards -> SUBSET Ids, the data before the batch
  db,         \* the data now
  dbStart,    \* the data when the current attempt started planning
  phase,
  i,          \* next op to plan
  committed,  \* shards committed in this attempt
  ever,       \* shards committed in any attempt
  audit       \* ops audited

vars == <<ops, rejects, db0, db, dbStart, phase, i, committed, ever, audit>>

TypeOK ==
  /\ ops \in [Idx -> Ops]
  /\ rejects \in [Idx -> BOOLEAN]
  /\ db0 \in [Shards -> SUBSET Ids]
  /\ db \in [Shards -> SUBSET Ids]
  /\ dbStart \in [Shards -> SUBSET Ids]
  /\ phase \in Phases
  /\ i \in 1..(N + 1)
  /\ committed \subseteq Shards
  /\ ever \subseteq Shards
  /\ audit \subseteq Idx

Init ==
  /\ ops \in [Idx -> Ops]
  /\ rejects \in [Idx -> BOOLEAN]
  /\ db0 \in [Shards -> SUBSET Ids]
  /\ db = db0
  /\ dbStart = db0
  /\ phase = "plan"
  /\ i = 1
  /\ committed = {}
  /\ ever = {}
  /\ audit = {}

(* ---------------------------------------------------------------- helpers *)

Apply(S, k) == IF ops[k].kind = "create" THEN S \cup {ops[k].id} ELSE S \ {ops[k].id}

(* The content of shard s after ops 1..k are applied in order. *)
RECURSIVE ApplyUpTo(_, _, _)
ApplyUpTo(S, s, k) ==
  IF k = 0 THEN S
  ELSE LET P == ApplyUpTo(S, s, k - 1) IN IF ops[k].shard = s THEN Apply(P, k) ELSE P

Expected(D) == [s \in Shards |-> ApplyUpTo(D[s], s, N)]

Rejected == \E k \in Idx : rejects[k]

(* ------------------------------------------------------------------ steps *)

Plan ==
  /\ phase = "plan" /\ i <= N
  /\ IF rejects[i]
       THEN phase' = "rejected" /\ UNCHANGED i
       ELSE /\ i' = i + 1
            /\ phase' = IF i = N THEN "commit" ELSE "plan"
  /\ UNCHANGED <<ops, rejects, db0, db, dbStart, committed, ever, audit>>

Commit(s) ==
  /\ phase = "commit" /\ s \notin committed
  /\ db' = [db EXCEPT ![s] = ApplyUpTo(db[s], s, N)]
  /\ committed' = committed \cup {s}
  /\ ever' = ever \cup {s}
  /\ phase' = IF committed' = Shards THEN "finish" ELSE "commit"
  /\ UNCHANGED <<ops, rejects, db0, dbStart, i, audit>>

Finish(k) ==
  /\ phase = "finish" /\ k \notin audit /\ ops[k].shard \in committed
  /\ audit' = audit \cup {k}
  /\ phase' = IF audit' = Idx THEN "done" ELSE "finish"
  /\ UNCHANGED <<ops, rejects, db0, db, dbStart, i, committed, ever>>

Crash ==
  /\ phase \in {"plan", "commit", "finish"}
  /\ phase' = "crashed"
  /\ UNCHANGED <<ops, rejects, db0, db, dbStart, i, committed, ever, audit>>

(* The client retries the whole batch against whatever landed. *)
Retry ==
  /\ phase = "crashed"
  /\ phase' = "plan" /\ i' = 1 /\ committed' = {} /\ dbStart' = db
  /\ UNCHANGED <<ops, rejects, db0, db, ever, audit>>

Next ==
  \/ Plan
  \/ \E s \in Shards : Commit(s)
  \/ \E k \in Idx : Finish(k)
  \/ Crash \/ Retry

Spec == Init /\ [][Next]_vars

(***************************************************************************)
(* Invariants                                                              *)
(***************************************************************************)

(* A shard holds either what it had when the attempt started or the whole
   batch applied: never some of the batch's ops for that shard. *)
NoPartialShard ==
  \A s \in Shards : db[s] = dbStart[s] \/ db[s] = ApplyUpTo(dbStart[s], s, N)

(* A rejected batch changed nothing. *)
RejectedBatchWritesNothing ==
  phase = "rejected" => db = dbStart

(* Audit rows exist only for ops whose shard committed. *)
AuditOnlyCommitted ==
  \A k \in audit : ops[k].shard \in ever

(* However many crashes and retries, a batch that completes leaves the data
   exactly as one uninterrupted run would have. *)
RetryConverges ==
  phase = "done" => db = Expected(db0)

=============================================================================