TD-6376: Improve lease management: fix lease races and rework rebalancing to per-worker

[tkusumo]

TD-6376

Summary

Fixes several correctness bugs in the lease subsystem and reworks lease rebalancing from a per-shard tick to a single per-worker Rebalancer process.

Correctness fixes

• Failed renewal now stops the pipeline and releases lease_holder. Previously LeaseV2 logged the error and kept consuming, so a shard whose lease was stolen or expired could be processed by two workers at once (a regression from V1's behavior).
• Steals use the freshly read lease_count instead of the copy in state, which was only synced on the 30s renew tick — a stale count failed the optimistic lock on every steal attempt for up to 30s.
• The Ecto optimistic-lock WHERE clause now pins lease_owner alongside lease_count, closing the read-then-update race and matching the Dynamo adapter's conditional expressions.
• Workers holding zero leases are counted in the balance check. They're absent from the grouped lease counts, so a fresh node looked "balanced" and never claimed its share of shards — scale-out quietly didn't rebalance.

Dynamo adapter

Implements the load balancing queries (get_leases_by_worker, all_incomplete_leases, total_incomplete_lease_counts_by_worker) with filtered scans + ExAws.stream! pagination. Previously they raised BadFunctionError or silently returned [], which made V2 lease balancing a no-op on DynamoDB.

Rebalancing redesign

• New KinesisClient.Stream.Rebalancer — one per Stream supervisor (per worker), same shape as the Java KCL's LeaseTaker. Each jittered tick (rebalance_interval ± 25%) it makes one lease-count query and, when under-loaded, requests at most max_leases_to_steal (default 1) steals. Before: every shard's lease process ran 3-5 full-table queries every 6s and stole unboundedly in the same tick window (100 shards ≈ 5-6k queries/min, now ~tens/min).
• LoadBalance is now pure decision logic (decide/2) with a flip-flop guard: no steal when the lead is a single lease (e.g. 7 shards on 2 workers can never be more even than 4/3).
• LeaseV2 executes steals on :steal_lease messages and stays the single writer for its shard's lease state; its rebalance timer and balancing logic are gone. Crash recovery (take-on-expiry) is unchanged.
• Removes the now-unused lease_owner_with_most_leases adapter callback.

Docs

README now documents the actual lease options (lease_renew_interval, lease_expiry, rebalance_interval, max_leases_to_steal) and the load balancing behavior — lease_renewal_limit was a dead V1-only option.

Test plan

• New pure LoadBalance.decide/2 unit tests, including a step-by-step convergence simulation of a {4, 4, 0} cluster.
• New Rebalancer tests: balanced → no steal; steal routed to the correct local lease process; per-tick cap respected; completed shards skipped.
• LeaseV2 tests reworked around direct :steal_lease messages, plus regression tests for the renewal-failure and stale-count bugs.
• Full suite: 79 tests, 0 failures across repeated runs. The Dynamo adapter tests require localstack (not available locally) — please confirm they pass in CI.
• Note: coordinator_test "don't start child shard until parent shard is closed" previously asserted parent shard supervisors die — an artifact of the old init-time balancing crash-looping on an unstubbed mock. It now asserts parents stay alive; the real intent (children not started early) was always covered by its refute_receive lines.

🤖 Generated with Claude Code

[tkusumo]

Addressed the review in b9d44b5. Point-by-point:

Fix before merge — both addressed:

• Renew-failure wedge: confirmed, including that both adapters reject taking a lease you already own, so the state really was unrecoverable (it even survives a process restart, since the Shard supervisor passes the same lease_owner). Added a reclaim clause to take_or_renew_lease: when the row names this worker as owner but lease_holder is false, renew under the optimistic lock and restart the pipeline. If another worker took the lease in the window, the CAS renewal fails and we keep tracking. Heals within one renew tick; covered by a regression test.
• Rebalancer crash tearing down the stream: the tick body is now wrapped in a rescue — a failed balancing query logs, emits {:rebalance_failed, error}, and retries next tick. Chose rescue-and-continue over re-shaping the supervision tree to keep the Coordinator/ShardSupervisor crash semantics unchanged. Test: a permanently-raising counts query, asserting the process survives and keeps ticking.

Should fix — all three done:

• Steal clamp: LoadBalance.decide/2 now returns {:steal_from, victim, deficit} and steal_from takes min(deficit, max_leases_to_steal). Your {A:4, B:0} max-3 oscillation is a test case now (clamps to 2, lands on {2,2}).
• Mimic get_leases_by_worker delegates to the migration adapters (same call-both-return-from pattern as the other reads).
• Victim in the steal message: it's {:steal_lease, victim} now; maybe_steal pattern-matches the current owner against the chosen victim and skips (with a debug log) if ownership changed since the decision.

Also picked up from the smaller notes: LeaseV2.whereis/3 replaces the re-derived registration name in the Rebalancer, jitter/1 no longer raises for intervals < 2ms, the dead :spread_lease option is gone, and the misleading distributed-:shard_supervisor comment is deleted.

Deferred (agree they're not blockers): steal-at-renewal-boundary to close the ~30s double-consumption window, Dynamo scan efficiency (projection + reusing one scan per tick), and deduplicating notify/2.

84 tests, 0 failures across repeated runs (16 invalid = the localstack-dependent Dynamo tests, unchanged).

[tkusumo]

The three deferred items are now in this PR as well (c289ee9):

Double-consumption window: on inspection, the fetch path was already owner-guarded — Producer.get_records_with_retry/2 checks lease ownership before every Kinesis call, so a stolen lease stops new fetches within one poll cycle, not one renew interval. The real gap was that an owner-guarded checkpoint failure only logged: the victim's pipeline kept its retry loops and "started" status until the lease process noticed at its next renewal. The Producer now stops itself (stop_if_lease_lost/2) the moment a checkpoint fails against a lease owned by another worker, bounding post-steal overlap to the messages already in flight. Given that, I didn't adopt steal-at-renewal-boundary — it would add a consumption gap on every rebalance to close a window that's now one in-flight batch.

Dynamo scan efficiency: went with your "one scan per tick" suggestion, which supersedes the projection idea (a projection_expression reduces payload but not consumed RCU, and the single remaining scan needs full items for the steal candidates anyway). The Rebalancer fetches the incomplete leases once and derives both the per-worker counts and the victim's shards from that list — unbalanced ticks now cost exactly the same as balanced ones, on both adapters. This also removes the client-side completed filtering you flagged (the incomplete-leases query excludes them server-side) and made total_incomplete_lease_counts_by_worker dead code, so the callback is gone from the behaviour, facade, and all three adapters.

notify/2 dedup: moved to KinesisClient.Util (the Coordinator keeps its own — it notifies off a different state key). Also renamed Producer.is_lease_owner? to lease_owner? while in the file, per Credo.

Net for the PR: -41 lines, one new Producer regression test (checkpoint fails → thief named → producer stops). 83 tests, 0 failures across repeated runs; the 15 invalid are the localstack-dependent Dynamo tests as before.

[tkusumo]

Gap clamp added in 4e1587a — confirmed your counter-example by hand first ({2, 4, 4}: deficit 2, but taking 2 flips the A↔B gap from −2 to +2 and the pair trades the same leases every round).

decide/2 now returns min(target - my_count, div(victim_count - my_count, 2)) as the steal count, and the Rebalancer applies the max_leases_to_steal cap on top, exactly as you suggested. The victim_count - my_count > 1 steal guard guarantees the half-gap term is ≥ 1, so the pos_integer contract holds.

Independently re-verified convergence with a simulation against the real decide/2: 2,500 randomized states (500 seeds × max_leases_to_steal ∈ {1, 2, 3, 4, 8}, 2–7 workers, counts 0–11, worker order shuffled every round) plus your explicit {2,4,4} @ 2, {5,5,0} @ 4, and {4,0} @ 3 cases — zero non-converging states within 60 rounds.

Tests: the {2,4,4} live-lock is a named regression case, a {5,5,0} round-by-round convergence case is added, and the existing decide/2 expectations were updated for the tighter steal counts (e.g. {1,4,4} now steals 1, not 2). 85 tests, 0 failures.

[tkusumo]

Reviewed the full change plus the three follow-up commits — the changes look good.

Original review findings, all resolved and regression-tested:

• Renew-failure wedge — reclaim_shard_lease recovers the "DB says I own it but I'm not holding" state under the optimistic lock (also heals the pre-existing restart wedge via handle_continue).
• Rebalancer blast radius — the tick is wrapped in a rescue with the reschedule happening before the work, so a throttled scan or DB blip logs and retries instead of restarting the whole :one_for_all tree.
• Steal convergence — decide/2 now returns min(target - my_count, div(victim_count - my_count, 2)). Verified independently with a simulation against the real module: the {2,4,4} and {5,5,0} live-locks from review now converge, and 500+ randomized states/orderings all settle balanced.
• Victim race — {:steal_lease, victim} with an owner re-check skips steals when the lease changed hands after the decision.
• Mimic migration adapter — get_leases_by_worker delegates, so rebalancing works during store migrations.
• Bounded steal overlap — the producer stops itself on an owner-guarded checkpoint failure, limiting double consumption to in-flight messages (nicer fix than the one suggested in review).
• Single-scan tick + cleanup — one all_incomplete_leases query serves both counts and steal candidates; the redundant adapter callback, dead spread_lease option, and duplicated notify/2 helpers are gone.

Suite: 85 tests, 0 failures locally. The 15 invalid are the localstack-dependent Dynamo tests — worth confirming those pass in CI before merge, as the PR description notes.