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.

[tkusumo]

Heads-up: PR #50 (TD-6395) is landing a fix for the recurring :update_checkpoint_failed in prod, and it interacts with this PR — flagging so we can sequence/rebase cleanly.

The bug: update_checkpoint (Ecto) reads via get_shard_lease then writes via update_shard_lease/3, whose build_where_clause/2 gates on lease_count. Leases renew every 30s (same owner), bumping lease_count. A renewal landing between the checkpoint's read and write makes the UPDATE match 0 rows → :update_checkpoint_failed, even though the worker still owns the lease.

Interaction with this PR:

• This PR keeps update_checkpoint routing through the lease_count-guarded update_shard_lease/3 (and actually tightens build_where_clause by adding lease_owner on top of lease_count), so the same-owner checkpoint race is still present here.
• The new stop_if_lease_lost/2 in the producer treats a checkpoint failure where the lease is still owned by this worker (%{lease_owner: owner}, %{lease_owner: owner} -> state) as benign and carries on — i.e. it silently tolerates exactly this failure rather than preventing it. Its doc comment ("Checkpoint updates are owner-guarded…") assumes a failure implies an ownership change, which isn't true while the lease_count guard remains.

PR #50's approach: give checkpoint its own ShardLeases.update_checkpoint/3 guarded on ownership only (shard_id + app_name + stream_name + lease_owner), never lease_count — mirroring the Dynamo adapter. Renew/take keep their lease_count optimistic lock untouched.

No conflict expected at the file level (PR #50 doesn't touch the producer or build_where_clause), but once both land, stop_if_lease_lost's same-owner branch becomes effectively dead for the checkpoint case since the write won't spuriously fail anymore. Worth a quick look when you rebase.

[tkusumo]

Ran the full kcl_test harness against this branch (90dbd3b): 7 scenarios — single-worker splits (3→6→12), scale-out (worker joining mid-stream), kill -9 failover, splits with 2 live workers, an odd-shard-count idle watch for steal flip-flop, and a 2000-msg burst. 3,200 messages total, 0 missing, 0 duplicates, and steal behavior converged and went quiet in every case. Two observations from the runs, neither data-affecting:

1. Boot-time race: LeaseV2 can start the pipeline before Broadway registers

When a worker wins a lease during init, handle_continue(:initialize) calls state.pipeline.start(state) (lease_v2.ex#L112) which hits Broadway.Topology.config/1 before the shard's Broadway process has registered (pipeline.ex#L67):

[error] GenServer LeaseV2.kcl_test.kcl_test_stream.shardId-000000000002 terminating
** (stop) exited in: Broadway.Topology.config(...Pipeline...shardId-000000000002)
   ** (EXIT) no process
   (broadway 1.1.0) lib/broadway/topology.ex:38: Broadway.Topology.config/1
   (kinesis_client) lib/kinesis_client/stream/shard/pipeline.ex:69: Pipeline.start/1
   (kinesis_client) lib/kinesis_client/stream/shard/lease_v2.ex:112: LeaseV2.handle_continue/2

Reproduced on 2 of 3 cold boots — but only when the worker immediately acquires a lease at init (a worker joining a cluster where all leases are held boots clean, since its lease processes start idle). The supervisor restarts it and the retry succeeds, so it self-heals, but it'll log a crash on essentially every fresh deploy of a lone/first worker. A bounded retry on Pipeline.start/1's no-proc exit (or deferring the initial start one tick) would silence it.

2. Cosmetic: expiry-take path doesn't skip completed shards

After a parent shard completes post-split, the other worker's idle LeaseV2 process still sees the lease as expired (owner stopped renewing on completion) and takes it — so lease_owner flips on rows that are already completed = true:

[debug] ShardLease: Lease expired, attempting to take lease: [shard_id: shardId-000000000000, ... current_owner: <dead-ish owner>]
[debug] ShardLease: Taking lease: [shard_id: shardId-000000000000, lease_owner: <other worker>]

The Rebalancer correctly skips completed shards, but the take-on-expiry path in LeaseV2 (lease_v2.ex#L365 area) doesn't check completed. Harmless — the shard is closed and nothing gets processed — but it's confusing churn when reading the lease table, and a wasted write per completed shard. A completed guard before attempting the take would clean it up.

---

Update: 3-node cluster results (same branch, 6 shards, Ecto adapter)

• Scale-out 1→3: started workers B and C while A held all 6 leases mid-stream. Converged 6/0/0 → 2/2/2 in under 15s with exactly 4 steals (each new worker took precisely its deficit of 2 — no overshoot, no contention between the two Rebalancers stealing from the same victim concurrently), then stayed pinned. 300/300 msgs, 0 missing, 0 dupes.
• Node kill at N=3: kill -9 on a worker holding 2 leases; after lease expiry the two survivors absorbed one each → clean 3/3 in ~90s. 300/300 msgs, 0 missing, 0 dupes.
• Race protection exercised for real: several times two workers raced for the same lease (concurrent scale-out steals, and both survivors going for the dead node's shards). Every race had exactly one winner — the owner-pinned optimistic lock rejected the loser with :lease_take_failed, and 0 duplicates end-to-end confirms no double ownership ever occurred.

One more minor observation from this run:

3. Lost take races log at [error] level

When two workers legitimately race for the same expired lease, the loser logs:

[error] KinesisClient: Error trying to take lease for shardId-000000000003: :not_found
[error] ShardLease: Error trying to take lease for shard shardId-000000000003, ... error: :lease_take_failed

That's the optimistic lock working as designed, but it produced 13 [error] lines across two workers in a single scale event. In production, any deploy/scale-out will emit a burst of these and look alarming on error-rate dashboards. Suggest downgrading the lost-race case (:lease_take_failed / :not_found on take) to warn or info, keeping error for genuine adapter failures.

Running totals across all 8 scenarios: 3,800 messages, 0 missing, 0 duplicates.

[tkusumo]

Addressed all three observations from the harness runs in 2352c08:

1. Boot-time race (LeaseV2 starts pipeline before Broadway registers) — Pipeline.start/1 now retries on the :noproc exit with a bounded backoff (25 × 100ms). Root cause confirmed: Shard supervises [LeaseV2, Pipeline] as :one_for_all with the lease process first, so a lease won in handle_continue can call start/1 before the sibling Broadway topology has registered. Broadway comes up within a tick, so the retry silences the crash; if it genuinely never comes up, the exit still propagates after the cap (self-heals via restart, as before).

2. Expiry-take doesn't skip completed shards — take_or_renew_lease now guards the expiry branch with not shard_lease.completed, so a completed shard's forever-"expired" lease falls through to tracking instead of flipping lease_owner on a completed = true row. Added a regression test (doesn't take an expired lease when the shard is completed).

3. Lost take/steal races log at [error] — downgraded to [warning] in three places: LeaseV2.take_shard_lease, LeaseV2.steal_shard_lease, and the Ecto adapter's take_lease (the KinesisClient: Error trying to take lease… line). Genuine adapter failures raise rather than returning {:error, _}, so real problems still surface at [error] and crash loudly — only the expected optimistic-lock contention is quieted.

Full non-integration suite: 87 tests, 0 failures. Didn't touch the load-balancing / steal-convergence logic, since the runs showed it converging cleanly (0 missing / 0 dupes across all scenarios).