From 7c497f26410e4b2c8a4f7378a7cc1171770fb2be Mon Sep 17 00:00:00 2001
From: Jon Kusumo <tahir.kusumo@paradox.ai>
Date: Wed, 1 Jul 2026 19:00:08 -0400
Subject: [PATCH 1/4] Improve lease management: fix races, rework rebalancing
 to per-worker
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Lease correctness fixes:
- Stop the pipeline and release lease_holder when a lease renewal fails
  (was silently continuing to consume, risking duplicate processing)
- Steal with the freshly read lease_count instead of the stale state
  copy, which failed the optimistic lock on every steal attempt
- Pin lease_owner (not just lease_count) in the Ecto update WHERE
  clause, matching the Dynamo adapter's conditional expressions
- Count workers holding zero leases in the balance check so new nodes
  actually claim their share of shards

Dynamo adapter:
- Implement the load balancing queries (get_leases_by_worker,
  all_incomplete_leases, total_incomplete_lease_counts_by_worker) with
  filtered scans; previously they raised or silently returned []

Rebalancing redesign:
- New per-worker KinesisClient.Stream.Rebalancer replaces the per-shard
  rebalance tick: one lease-count query per worker per tick instead of
  3-5 queries per shard per tick
- Steals are capped (max_leases_to_steal, default 1) and the tick is
  jittered +/- 25%, so workers converge without stampeding or flapping
- LoadBalance is now pure decision logic with a flip-flop guard (no
  steal when the lead is a single lease)
- LeaseV2 executes steals on :steal_lease messages, staying the single
  writer for its shard; drops its balancing logic and rebalance timer
- Remove the now-unused lease_owner_with_most_leases adapter callback
- Use the injected pipeline module consistently in LeaseV2

Also documents the current lease options and load balancing behavior in
the README (lease_renewal_limit was a dead V1-only option).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 README.md                                     |  27 +-
 lib/kinesis_client/stream.ex                  |  20 +-
 lib/kinesis_client/stream/app_state.ex        |   6 -
 .../stream/app_state/adapter.ex               |   7 -
 lib/kinesis_client/stream/app_state/dynamo.ex |  44 +--
 lib/kinesis_client/stream/app_state/ecto.ex   |  15 -
 .../stream/app_state/ecto/shard_leases.ex     |  26 +-
 lib/kinesis_client/stream/app_state/mimic.ex  |   8 -
 lib/kinesis_client/stream/rebalancer.ex       | 138 ++++++++
 lib/kinesis_client/stream/shard.ex            |   1 -
 lib/kinesis_client/stream/shard/lease_v2.ex   | 158 +++------
 .../stream/shard/load_balance.ex              | 132 +++----
 .../stream/app_state/dynamo_test.exs          |  48 +++
 .../stream/coordinator_test.exs               |   4 +-
 .../kinesis_client/stream/rebalancer_test.exs | 157 +++++++++
 .../stream/shard/lease_v2_test.exs            | 328 +++++++++---------
 .../stream/shard/load_balance_test.exs        |  61 ++++
 test/support/test_pipeline.ex                 |   4 +
 18 files changed, 743 insertions(+), 441 deletions(-)
 create mode 100644 lib/kinesis_client/stream/rebalancer.ex
 create mode 100644 test/kinesis_client/stream/rebalancer_test.exs
 create mode 100644 test/kinesis_client/stream/shard/load_balance_test.exs

diff --git a/README.md b/README.md
index 47e0595..c2779ad 100644
--- a/README.md
+++ b/README.md
@@ -35,8 +35,18 @@ opts = [
     app_name: app_name,
     stream_name: stream_name
   ],
-  # optional to limit the amount of times a lease can be renewed
-  lease_renewal_limit: 10,
+  # optional, how often (ms) a held lease is renewed. Default: 30_000
+  lease_renew_interval: 30_000,
+  # optional, how long (ms) a lease can go unrenewed before another worker
+  # may take it. Must be greater than lease_renew_interval. Default: 45_001
+  lease_expiry: 45_001,
+  # optional, how often (ms) this worker checks whether leases are spread
+  # evenly across workers and steals one from an overloaded worker if not.
+  # Default: 6_000
+  rebalance_interval: 6_000,
+  # optional, the maximum number of leases to steal per rebalance check.
+  # Default: 1
+  max_leases_to_steal: 1,
   # optional poll_interval for getting records from kinesis
   poll_interval: 500,
   processors: [
@@ -90,8 +100,17 @@ doing it currently:
 SERVICES=kinesis,dynamodb localstack start --host
 ```
 
+## Load balancing
+
+Each worker runs a single rebalancer process that periodically (every
+`rebalance_interval`, jittered +/- 25%) compares how many incomplete leases
+each worker holds. When this worker is below the target load
+(`ceil(total_shards / total_workers)`) and another worker leads it by more
+than one lease, it steals a lease from the worker holding the most — at most
+`max_leases_to_steal` (default 1) per check, so workers converge on an even
+spread without overshooting. Workers that crash stop renewing their leases,
+and after `lease_expiry` the remaining workers take those leases over.
+
 ## TODO
 - [ ] Test shard merges and splits more thoroughly
-- [ ] Implement a work stealing algorithim to help distribute the load among
-  different Elixir nodes processing the same app.
 
diff --git a/lib/kinesis_client/stream.ex b/lib/kinesis_client/stream.ex
index a6b2cdc..0061b67 100644
--- a/lib/kinesis_client/stream.ex
+++ b/lib/kinesis_client/stream.ex
@@ -7,6 +7,7 @@ defmodule KinesisClient.Stream do
   import KinesisClient.Util
 
   alias KinesisClient.Stream.Coordinator
+  alias KinesisClient.Stream.Rebalancer
 
   require Logger
 
@@ -26,6 +27,11 @@ defmodule KinesisClient.Stream do
     * `:lease_expiry`(optional) - The length of time in milliseconds that least lasts for. If a
       lease is not renewed within this time frame, then that lease is considered expired and can be
       taken by another process.
+    * `:rebalance_interval`(optional) - How often (in milliseconds) this worker checks whether
+      leases are spread evenly across workers and steals one from an overloaded worker if not.
+      The interval is jittered +/- 25%. Defaults to 6,000.
+    * `:max_leases_to_steal`(optional) - The maximum number of leases to steal per rebalance
+      check. Defaults to 1.
   """
   def start_link(opts) do
     Supervisor.start_link(__MODULE__, opts, name: Keyword.get(opts, :name, __MODULE__))
@@ -61,10 +67,19 @@ defmodule KinesisClient.Stream do
       |> optional_kw(:lease_expiry, Keyword.get(opts, :lease_expiry))
       |> optional_kw(:spread_lease, Keyword.get(opts, :spread_lease))
       |> optional_kw(:poll_interval, Keyword.get(opts, :poll_interval))
-      |> optional_kw(:rebalance_interval, Keyword.get(opts, :rebalance_interval))
       |> optional_kw(:shard_iterator_type, Keyword.get(opts, :shard_iterator_type))
       |> optional_kw(:timestamp, Keyword.get(opts, :timestamp))
 
+    rebalancer_args =
+      [
+        app_name: app_name,
+        stream_name: stream_name,
+        lease_owner: worker_ref,
+        app_state_opts: Keyword.get(opts, :app_state_opts, [])
+      ]
+      |> optional_kw(:rebalance_interval, Keyword.get(opts, :rebalance_interval))
+      |> optional_kw(:max_leases_to_steal, Keyword.get(opts, :max_leases_to_steal))
+
     coordinator_args = [
       name: coordinator_name,
       stream_name: stream_name,
@@ -77,7 +92,8 @@ defmodule KinesisClient.Stream do
 
     children = [
       shard_supervisor_spec,
-      {Coordinator, coordinator_args}
+      {Coordinator, coordinator_args},
+      {Rebalancer, rebalancer_args}
     ]
 
     Logger.info(
diff --git a/lib/kinesis_client/stream/app_state.ex b/lib/kinesis_client/stream/app_state.ex
index bfbe445..936a9f4 100644
--- a/lib/kinesis_client/stream/app_state.ex
+++ b/lib/kinesis_client/stream/app_state.ex
@@ -61,12 +61,6 @@ defmodule KinesisClient.Stream.AppState do
   def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts \\ []),
     do: adapter(opts).total_incomplete_lease_counts_by_worker(app_name, stream_name, opts)
 
-  @doc """
-  Get lease owner that has the most leases.
-  """
-  def lease_owner_with_most_leases(app_name, stream_name, opts \\ []),
-    do: adapter(opts).lease_owner_with_most_leases(app_name, stream_name, opts)
-
   defp adapter(opts) do
     case Keyword.get(opts, :adapter) do
       :ecto -> KinesisClient.Stream.AppState.Ecto
diff --git a/lib/kinesis_client/stream/app_state/adapter.ex b/lib/kinesis_client/stream/app_state/adapter.ex
index 1900c48..d6a6df5 100644
--- a/lib/kinesis_client/stream/app_state/adapter.ex
+++ b/lib/kinesis_client/stream/app_state/adapter.ex
@@ -84,11 +84,4 @@ defmodule KinesisClient.Stream.AppState.Adapter do
               opts :: keyword
             ) ::
               list({worker :: String.t(), count :: integer})
-
-  @callback lease_owner_with_most_leases(
-              app_name :: String.t(),
-              stream_name :: String.t(),
-              opts :: keyword
-            ) ::
-              list(ShardLease.t())
 end
diff --git a/lib/kinesis_client/stream/app_state/dynamo.ex b/lib/kinesis_client/stream/app_state/dynamo.ex
index 15fab96..136bd17 100644
--- a/lib/kinesis_client/stream/app_state/dynamo.ex
+++ b/lib/kinesis_client/stream/app_state/dynamo.ex
@@ -47,10 +47,11 @@ defmodule KinesisClient.Stream.AppState.Dynamo do
   end
 
   @impl true
-  def get_leases_by_worker(_app_name, _stream_name, _lease_owner, _opts) do
-    raise BadFunctionError,
-      message:
-        "get_leases_by_worker/4 is not currently implemented for DynamoDB. Please implement the callback if you want to use DynamoDB."
+  def get_leases_by_worker(app_name, _stream_name, lease_owner, _opts) do
+    scan_shard_leases(app_name,
+      filter_expression: "lease_owner = :lease_owner",
+      expression_attribute_values: [lease_owner: lease_owner]
+    )
   end
 
   @impl true
@@ -192,30 +193,29 @@ defmodule KinesisClient.Stream.AppState.Dynamo do
   end
 
   @impl true
-  def all_incomplete_leases(_app_name, _stream_name, _opts) do
-    Logger.error(
-      "all_incomplete_leases/3 is not currently implemented for DynamoDB. Please implement the callback if you want to use DynamoDB."
+  def all_incomplete_leases(app_name, _stream_name, _opts) do
+    scan_shard_leases(app_name,
+      filter_expression: "completed = :completed",
+      expression_attribute_values: [completed: false]
     )
-
-    []
   end
 
   @impl true
-  def lease_owner_with_most_leases(_app_name, _stream_name, _opts) do
-    Logger.error(
-      "lease_owner_with_most_leases/3 is not currently implemented for DynamoDB. Please implement the callback if you want to use DynamoDB."
-    )
-
-    []
+  def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts) do
+    app_name
+    |> all_incomplete_leases(stream_name, opts)
+    |> Enum.frequencies_by(& &1.lease_owner)
+    |> Map.to_list()
   end
 
-  @impl true
-  def total_incomplete_lease_counts_by_worker(_app_name, _stream_name, _opts) do
-    Logger.error(
-      "total_incomplete_lease_counts_by_worker/3 is not currently implemented for DynamoDB. Please implement the callback if you want to use DynamoDB."
-    )
-
-    []
+  # Lease tables hold one row per shard, so a filtered Scan is cheap enough for
+  # the load balancing queries (this mirrors how the Java KCL reads its lease
+  # table). ExAws.stream!/1 follows LastEvaluatedKey pagination for us.
+  defp scan_shard_leases(app_name, scan_opts) do
+    app_name
+    |> Dynamo.scan(scan_opts)
+    |> ExAws.stream!()
+    |> Enum.map(&decode_item/1)
   end
 
   defp decode_item(item) do
diff --git a/lib/kinesis_client/stream/app_state/ecto.ex b/lib/kinesis_client/stream/app_state/ecto.ex
index c496e0c..9d830fd 100644
--- a/lib/kinesis_client/stream/app_state/ecto.ex
+++ b/lib/kinesis_client/stream/app_state/ecto.ex
@@ -190,21 +190,6 @@ defmodule KinesisClient.Stream.AppState.Ecto do
     |> ShardLeases.get_shard_leases(repo)
   end
 
-  @impl true
-  def lease_owner_with_most_leases(app_name, stream_name, opts) do
-    repo = Keyword.get(opts, :repo)
-
-    app_name
-    |> ShardLeases.get_owner_with_most_leases(stream_name, repo)
-    |> case do
-      nil ->
-        []
-
-      worker ->
-        get_leases_by_worker(app_name, stream_name, worker, opts)
-    end
-  end
-
   @impl true
   def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts) do
     repo = Keyword.get(opts, :repo)
diff --git a/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex b/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
index bab5319..bc8be0d 100644
--- a/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
+++ b/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
@@ -1,4 +1,5 @@
 defmodule KinesisClient.Stream.AppState.Ecto.ShardLeases do
+  @moduledoc false
   alias KinesisClient.Stream.AppState.Ecto.ShardLease, as: ShardLeaseEcto
   alias KinesisClient.Stream.AppState.ShardLease
 
@@ -62,24 +63,10 @@ defmodule KinesisClient.Stream.AppState.Ecto.ShardLeases do
     |> repo.all()
   end
 
-  @spec get_owner_with_most_leases(String.t(), String.t(), Ecto.Repo.t()) ::
-          owner :: String.t() | nil
-  def get_owner_with_most_leases(app_name, stream_name, repo) do
-    owner_counts = incomplete_group_by_owner(app_name, stream_name, repo)
-
-    case owner_counts do
-      [] ->
-        nil
-
-      counts ->
-        max_count = counts |> Enum.map(fn {_owner, count} -> count end) |> Enum.max()
-
-        {owner, _count} = Enum.find(counts, fn {_owner, count} -> count == max_count end)
-
-        owner
-    end
-  end
-
+  # Pins the row to the freshly read lease_count AND lease_owner so a
+  # concurrent renew/take between our read and this update makes the
+  # update_all match zero rows instead of clobbering the other worker's
+  # lease. This matches the Dynamo adapter's conditional expressions.
   defp build_where_clause(query, shard_lease) do
     query
     |> where(
@@ -87,7 +74,8 @@ defmodule KinesisClient.Stream.AppState.Ecto.ShardLeases do
       sl.shard_id == ^shard_lease.shard_id and
         sl.app_name == ^shard_lease.app_name and
         sl.stream_name == ^shard_lease.stream_name and
-        sl.lease_count == ^shard_lease.lease_count
+        sl.lease_count == ^shard_lease.lease_count and
+        sl.lease_owner == ^shard_lease.lease_owner
     )
   end
 
diff --git a/lib/kinesis_client/stream/app_state/mimic.ex b/lib/kinesis_client/stream/app_state/mimic.ex
index 59ede12..e7ad60b 100644
--- a/lib/kinesis_client/stream/app_state/mimic.ex
+++ b/lib/kinesis_client/stream/app_state/mimic.ex
@@ -102,14 +102,6 @@ defmodule KinesisClient.Stream.AppState.Mimic do
     from.total_incomplete_lease_counts_by_worker(app_name, stream_name, opts)
   end
 
-  @impl true
-  def lease_owner_with_most_leases(app_name, stream_name, opts) do
-    {from, to} = modules(opts)
-
-    to.lease_owner_with_most_leases(app_name, stream_name, opts)
-    from.lease_owner_with_most_leases(app_name, stream_name, opts)
-  end
-
   defp modules(opts) do
     migration = Keyword.get(opts, :migration)
 
diff --git a/lib/kinesis_client/stream/rebalancer.ex b/lib/kinesis_client/stream/rebalancer.ex
new file mode 100644
index 0000000..569f7cb
--- /dev/null
+++ b/lib/kinesis_client/stream/rebalancer.ex
@@ -0,0 +1,138 @@
+defmodule KinesisClient.Stream.Rebalancer do
+  @moduledoc """
+  Periodically rebalances shard leases across workers.
+
+  One Rebalancer runs per `KinesisClient.Stream` (i.e. per worker). Each tick
+  it reads the per-worker lease counts once and, when this worker holds less
+  than its share, picks a lease held by the most loaded worker and tells the
+  local `KinesisClient.Stream.Shard.LeaseV2` process for that shard to steal
+  it. The lease process stays the single writer for its shard's lease state.
+
+  Steals are capped at `:max_leases_to_steal` per tick and the tick interval
+  is jittered, so workers converge on an even distribution without several of
+  them dog-piling the same lease in the same instant.
+  """
+  use GenServer
+
+  import KinesisClient.Util
+
+  alias KinesisClient.Stream.AppState
+  alias KinesisClient.Stream.Shard.LeaseV2
+  alias KinesisClient.Stream.Shard.LoadBalance
+
+  require Logger
+
+  @default_rebalance_interval 6_000
+  @default_max_leases_to_steal 1
+
+  def start_link(opts) do
+    GenServer.start_link(__MODULE__, opts,
+      name: register_name(__MODULE__, opts[:app_name], opts[:stream_name])
+    )
+  end
+
+  defstruct [
+    :app_name,
+    :stream_name,
+    :lease_owner,
+    :app_state_opts,
+    :rebalance_interval,
+    :max_leases_to_steal,
+    :notify
+  ]
+
+  @type t :: %__MODULE__{}
+
+  @impl GenServer
+  def init(opts) do
+    state = %__MODULE__{
+      app_name: opts[:app_name],
+      stream_name: opts[:stream_name],
+      lease_owner: opts[:lease_owner],
+      app_state_opts: Keyword.get(opts, :app_state_opts, []),
+      rebalance_interval: Keyword.get(opts, :rebalance_interval, @default_rebalance_interval),
+      max_leases_to_steal: Keyword.get(opts, :max_leases_to_steal, @default_max_leases_to_steal),
+      notify: Keyword.get(opts, :notify)
+    }
+
+    schedule_rebalance(state)
+
+    Logger.metadata(
+      kcl_app_name: state.app_name,
+      kcl_stream_name: state.stream_name,
+      kcl_lease_owner: state.lease_owner
+    )
+
+    Logger.info("Initializing KinesisClient.Stream.Rebalancer: #{inspect(state)}")
+
+    {:ok, state}
+  end
+
+  @impl GenServer
+  def handle_info(:rebalance, state) do
+    schedule_rebalance(state)
+
+    state.app_name
+    |> AppState.total_incomplete_lease_counts_by_worker(state.stream_name, state.app_state_opts)
+    |> LoadBalance.decide(state.lease_owner)
+    |> case do
+      :balanced ->
+        notify({:all_balanced, state}, state)
+
+      {:steal_from, victim} ->
+        steal_from(victim, state)
+    end
+
+    {:noreply, state}
+  end
+
+  defp steal_from(victim, state) do
+    state.app_name
+    |> AppState.get_leases_by_worker(state.stream_name, victim, state.app_state_opts)
+    |> Enum.reject(& &1.completed)
+    |> Enum.shuffle()
+    |> Enum.map(&local_lease_process(&1, state))
+    |> Enum.reject(&is_nil/1)
+    |> Enum.take(state.max_leases_to_steal)
+    |> Enum.each(fn {shard_id, pid} ->
+      Logger.debug(
+        "Rebalancer: Requesting steal of shard #{shard_id} from #{victim}: " <>
+          "[lease_owner: #{state.lease_owner}]"
+      )
+
+      send(pid, :steal_lease)
+      notify({:steal_requested, shard_id}, state)
+    end)
+  end
+
+  # The steal is executed by the shard's lease process so there is exactly one
+  # writer per shard on this worker. Shards without a running local lease
+  # process (not started yet, or already shut down) are skipped this tick.
+  defp local_lease_process(%{shard_id: shard_id}, state) do
+    LeaseV2
+    |> register_name(state.app_name, state.stream_name, [shard_id])
+    |> Process.whereis()
+    |> case do
+      nil -> nil
+      pid -> {shard_id, pid}
+    end
+  end
+
+  defp schedule_rebalance(%{rebalance_interval: interval}) do
+    Process.send_after(self(), :rebalance, jitter(interval))
+  end
+
+  # +/- 25% so workers don't tick in lockstep and stampede the same lease.
+  defp jitter(interval) do
+    interval + :rand.uniform(div(interval, 2)) - div(interval, 4)
+  end
+
+  defp notify(_msg, %{notify: nil}) do
+    :ok
+  end
+
+  defp notify(msg, %{notify: notify}) do
+    send(notify, msg)
+    :ok
+  end
+end
diff --git a/lib/kinesis_client/stream/shard.ex b/lib/kinesis_client/stream/shard.ex
index 2a25316..b8489c6 100644
--- a/lib/kinesis_client/stream/shard.ex
+++ b/lib/kinesis_client/stream/shard.ex
@@ -23,7 +23,6 @@ defmodule KinesisClient.Stream.Shard do
       |> optional_kw(:app_state_opts, Keyword.get(opts, :app_state_opts))
       |> optional_kw(:renew_interval, Keyword.get(opts, :lease_renew_interval))
       |> optional_kw(:lease_expiry, Keyword.get(opts, :lease_expiry))
-      |> optional_kw(:rebalance_interval, Keyword.get(opts, :rebalance_interval))
       |> optional_kw(:pipeline, Keyword.get(opts, :pipeline))
       |> optional_kw(:spread_lease, Keyword.get(opts, :spread_lease))
 
diff --git a/lib/kinesis_client/stream/shard/lease_v2.ex b/lib/kinesis_client/stream/shard/lease_v2.ex
index 10d8293..5e033bc 100644
--- a/lib/kinesis_client/stream/shard/lease_v2.ex
+++ b/lib/kinesis_client/stream/shard/lease_v2.ex
@@ -1,26 +1,26 @@
 defmodule KinesisClient.Stream.Shard.LeaseV2 do
   @moduledoc """
-  Load-balanced lease management for Kinesis shards.
-
-  This module implements a new load balancing mechanism where:
-  1. Each shard has a corresponding "lease" entry in the shard_lease table
-  2. Workers can steal leases from overloaded workers
-  3. Load balancing algorithm distributes shards evenly across workers
-  4. Periodic rebalancing and crash detection ensure optimal distribution
+  Lease management for a single Kinesis shard.
+
+  Each shard has a corresponding "lease" entry in the shard_lease table. This
+  process creates the lease if missing, renews it while held, and takes over
+  expired leases when their owner stops renewing (crash detection). It also
+  executes lease steals on behalf of `KinesisClient.Stream.Rebalancer`, which
+  makes the load balancing decisions once per worker and sends a
+  `:steal_lease` message to the shard it wants: this process stays the single
+  writer for its shard's lease state.
   """
   use GenServer
 
   import KinesisClient.Util
 
   alias KinesisClient.Stream.AppState
-  alias KinesisClient.Stream.Shard.LoadBalance
   alias KinesisClient.Stream.Shard.Pipeline
 
   require Logger
 
   @default_renew_interval 30_000
   @default_lease_expiry 45_001
-  @default_rebalance_interval 6_000
 
   def start_link(opts) do
     GenServer.start_link(__MODULE__, opts,
@@ -40,8 +40,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     :notify,
     :lease_expiry,
     :lease_holder,
-    :pipeline,
-    :rebalance_interval
+    :pipeline
   ]
 
   @type t :: %__MODULE__{}
@@ -59,8 +58,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
       lease_holder: Keyword.get(opts, :lease_holder, false),
       lease_count_increment_time: current_time(),
       notify: Keyword.get(opts, :notify),
-      pipeline: Keyword.get(opts, :pipeline, Pipeline),
-      rebalance_interval: Keyword.get(opts, :rebalance_interval, @default_rebalance_interval)
+      pipeline: Keyword.get(opts, :pipeline, Pipeline)
     }
 
     Process.send_after(self(), :take_or_renew_lease, state.renew_interval)
@@ -92,21 +90,17 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
 
         shard_lease ->
           Logger.debug(
-            "ShardLease: Found existing lease record in AppState, attempting load balancing: " <>
+            "ShardLease: Found existing lease record in AppState: " <>
               "[shard_id: #{state.shard_id}, lease_owner: #{shard_lease.lease_owner}]"
           )
 
-          shard_lease.lease_count
-          |> set_lease_count(false, state)
-          |> then(&load_balancing(shard_lease, &1))
+          set_lease_count(shard_lease.lease_count, false, state)
       end
 
     if new_state.lease_holder do
       :ok = state.pipeline.start(state)
     end
 
-    Process.send_after(self(), :rebalance, new_state.rebalance_interval)
-
     notify({:initialized, new_state}, state)
 
     {:noreply, new_state}
@@ -136,10 +130,14 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
+  # Sent by KinesisClient.Stream.Rebalancer when it decides this worker should
+  # steal this shard's lease from an overloaded worker.
   @impl GenServer
-  def handle_info(:rebalance, state) do
-    Process.send_after(self(), :rebalance, state.rebalance_interval)
+  def handle_info(:steal_lease, %{lease_holder: true} = state) do
+    {:noreply, state}
+  end
 
+  def handle_info(:steal_lease, state) do
     state
     |> get_shard_lease()
     |> case do
@@ -152,11 +150,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
         {:noreply, state}
 
       shard_lease ->
-        Logger.debug(
-          "ShardLease: Running rebalance process for shard #{state.shard_id}, lease_owner: #{state.lease_owner}, current_owner: #{shard_lease.lease_owner}"
-        )
-
-        {:noreply, load_balancing(shard_lease, state)}
+        {:noreply, maybe_steal(shard_lease, state)}
     end
   end
 
@@ -193,9 +187,20 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
         |> set_lease_count(true, state)
         |> tap(&notify({:lease_renewed, &1}, &1))
 
+      {:error, :lease_renew_failed} ->
+        Logger.error(
+          "ShardLease: Failed to renew lease, stopping pipeline: [app_name: #{app_name}, " <>
+            "shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}, current_owner: #{shard_lease.lease_owner}]"
+        )
+
+        :ok = state.pipeline.stop(state)
+
+        %{state | lease_holder: false, lease_count_increment_time: current_time()}
+        |> tap(&notify({:lease_renew_failed, &1}, &1))
+
       {:error, error} ->
         Logger.error(
-          "ShardLease: Failed to renew lease, error: #{inspect(error)}, [app_name: #{app_name}, " <>
+          "ShardLease: Error trying to renew lease, error: #{inspect(error)}, [app_name: #{app_name}, " <>
             "shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}], current_owner: #{shard_lease.lease_owner}"
         )
 
@@ -222,7 +227,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
         expected
         |> set_lease_count(true, state)
         |> tap(fn state -> notify({:lease_taken, state}, state) end)
-        |> tap(fn state -> Pipeline.start(state) end)
+        |> tap(fn state -> state.pipeline.start(state) end)
 
       {:error, error} ->
         Logger.error(
@@ -234,13 +239,32 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
+  defp maybe_steal(%{completed: true}, state), do: state
+
+  defp maybe_steal(%{lease_owner: current_owner} = shard_lease, %{lease_owner: me} = state)
+       when current_owner != me do
+    steal_shard_lease(shard_lease, state)
+  end
+
+  defp maybe_steal(_shard_lease, state) do
+    Logger.debug(
+      "ShardLease: Steal requested but the lease is already owned by this worker: " <>
+        "[shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
+    )
+
+    state
+  end
+
   defp steal_shard_lease(shard_lease, state) do
+    # Use the lease_count from the freshly read shard_lease rather than the copy
+    # in state: state.lease_count is only synced on the (slower) renew tick, and
+    # a stale count would fail the optimistic-lock check on every steal attempt.
     state.app_name
     |> AppState.take_lease(
       state.stream_name,
       state.shard_id,
       state.lease_owner,
-      state.lease_count,
+      shard_lease.lease_count,
       state.app_state_opts
     )
     |> case do
@@ -250,7 +274,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
         new_lease_count
         |> set_lease_count(true, state)
         |> tap(fn state -> notify({:lease_stolen, state}, state) end)
-        |> tap(fn state -> Pipeline.start(state) end)
+        |> tap(fn state -> state.pipeline.start(state) end)
 
       {:error, error} ->
         Logger.error(
@@ -262,27 +286,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
-  defp load_balancing(shard_lease, state) do
-    state
-    |> LoadBalance.check_if_balanced?()
-    |> case do
-      true ->
-        Logger.debug(
-          "ShardLease: Current load is balanced - no action needed: [shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
-        )
-
-        state
-        |> tap(&notify({:all_balanced, &1}, &1))
-
-      false ->
-        Logger.debug(
-          "ShardLease: Load is unbalanced, evaluating lease stealing options: [shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
-        )
-
-        attempt_to_steal(shard_lease, state)
-    end
-  end
-
   defp take_or_renew_lease(
          shard_lease,
          %{lease_expiry: lease_expiry, lease_count_increment_time: lcit} = state
@@ -336,59 +339,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
-  defp attempt_to_steal(shard_lease, state) do
-    case shard_lease.lease_owner != state.lease_owner do
-      true ->
-        Logger.debug(
-          "ShardLease: Current worker does not own the lease, attempting to steal from overloaded worker: [shard_id: #{state.shard_id}, " <>
-            "lease_owner: #{state.lease_owner}, current_owner: #{shard_lease.lease_owner}]"
-        )
-
-        steal_lease_from_overloaded_worker(state)
-
-      false ->
-        Logger.debug(
-          "ShardLease: Current worker already owns the lease, no action needed: [shard_id: #{state.shard_id}, " <>
-            "lease_owner: #{state.lease_owner}, current_owner: #{shard_lease.lease_owner}]"
-        )
-
-        state
-    end
-  end
-
-  defp steal_lease_from_overloaded_worker(state) do
-    state.app_name
-    |> AppState.lease_owner_with_most_leases(state.stream_name, state.app_state_opts)
-    |> case do
-      [] ->
-        Logger.debug(
-          "ShardLease: No overloaded workers found to steal from for shard #{state.shard_id}"
-        )
-
-        state
-
-      shard_leases ->
-        shard_leases
-        |> Enum.find(fn shard_lease -> shard_lease.shard_id == state.shard_id end)
-        |> case do
-          nil ->
-            Logger.debug(
-              "ShardLease: Shard #{state.shard_id} does not belong to an overloaded worker"
-            )
-
-            state
-
-          shard_lease ->
-            Logger.debug(
-              "ShardLease: Attempting to steal lease for shard #{state.shard_id} [lease_owner: #{state.lease_owner}, " <>
-                "current_owner: #{shard_lease.lease_owner}, lease_count: #{state.lease_count}, current_lease_count: #{shard_lease.lease_count}]"
-            )
-
-            steal_shard_lease(shard_lease, state)
-        end
-    end
-  end
-
   defp set_lease_count(lease_count, is_lease_holder, %__MODULE__{} = state) do
     %{
       state
diff --git a/lib/kinesis_client/stream/shard/load_balance.ex b/lib/kinesis_client/stream/shard/load_balance.ex
index c4a960c..d10388c 100644
--- a/lib/kinesis_client/stream/shard/load_balance.ex
+++ b/lib/kinesis_client/stream/shard/load_balance.ex
@@ -1,108 +1,64 @@
 defmodule KinesisClient.Stream.Shard.LoadBalance do
   @moduledoc """
-  This module will implement the load balancing logic for Kinesis shard leases.
-  It will provide functions to distribute shard leases evenly across workers,
-  detect overloaded workers, and facilitate lease stealing.
+  Pure decision logic for balancing shard leases across workers.
+
+  Given the incomplete-lease counts per worker, decides whether the current
+  worker should steal a lease from an overloaded worker. Fetching the counts
+  and executing the steal are the caller's concern (see
+  `KinesisClient.Stream.Rebalancer`).
   """
 
-  alias KinesisClient.Stream.AppState
-  alias KinesisClient.Stream.AppState.Ecto.ShardLease
-  alias KinesisClient.Stream.AppState.Ecto.ShardLeases
+  @doc """
+  Decides whether `lease_owner` should steal a lease.
 
-  @spec find_worker_with_most_leases(map()) :: list(ShardLease.t())
-  def find_worker_with_most_leases(state) do
-    repo = Keyword.get(state.app_state_opts, :repo)
+  Returns `{:steal_from, victim}` when `lease_owner` holds less than its share
+  of the leases (`ceil(total_leases / total_workers)`) and the most loaded
+  worker leads it by more than one lease, so a steal moves the distribution
+  closer to even instead of flipping the imbalance around. Returns `:balanced`
+  otherwise.
 
-    state.app_name
-    |> ShardLeases.get_owner_with_most_leases(state.stream_name, repo)
-    |> case do
-      nil ->
-        []
+  `lease_owner` is counted as a worker even when it holds no leases and is
+  therefore absent from the grouped counts — otherwise a fresh worker would
+  look "balanced" and never claim its share of the shards.
+  """
+  @spec decide(list({String.t(), non_neg_integer()}), String.t()) ::
+          :balanced | {:steal_from, String.t()}
+  def decide(worker_counts, lease_owner) do
+    worker_counts = include_current_worker(worker_counts, lease_owner)
+    target = target_load(worker_counts)
+    {^lease_owner, my_count} = List.keyfind(worker_counts, lease_owner, 0)
 
-      worker ->
-        AppState.get_leases_by_worker(
-          state.app_name,
-          state.stream_name,
-          worker,
-          state.app_state_opts
-        )
-    end
+    worker_counts
+    |> List.keydelete(lease_owner, 0)
+    |> steal_candidate(my_count, target)
   end
 
-  @spec calculate_load_metrics(map()) :: map()
-  def calculate_load_metrics(state) do
-    %{
-      state
-      | target_load: calculate_target_load(state),
-        current_load: get_current_worker_load(state)
-    }
-  end
+  defp steal_candidate([], _my_count, _target), do: :balanced
 
-  @spec total_leases_count_by_worker(map()) :: list({String.t(), integer})
-  def total_leases_count_by_worker(state) do
-    state.app_name
-    |> AppState.total_incomplete_lease_counts_by_worker(
-      state.stream_name,
-      state.app_state_opts
-    )
-    |> then(fn worker_counts ->
-      worker_counts
-      |> Enum.find(fn {owner, _count} -> owner == state.lease_owner end)
-      |> case do
-        nil -> [worker_counts ++ {state.lease_owner, 0}]
-        _ -> worker_counts
-      end
-    end)
-  end
+  defp steal_candidate(other_counts, my_count, target) do
+    other_counts
+    |> Enum.max_by(fn {_owner, count} -> count end)
+    |> case do
+      {victim, victim_count} when my_count < target and victim_count - my_count > 1 ->
+        {:steal_from, victim}
 
-  @spec check_if_balanced?(map()) :: boolean()
-  def check_if_balanced?(state) do
-    state.app_name
-    |> AppState.total_incomplete_lease_counts_by_worker(
-      state.stream_name,
-      state.app_state_opts
-    )
-    |> Enum.all?(fn {_owner, count} ->
-      abs(count - calculate_target_load(state)) <= 1
-    end)
+      _ ->
+        :balanced
+    end
   end
 
-  defp total_workers_count(workers, current_worker) do
-    workers
-    |> Enum.find(fn {owner, _count} -> owner == current_worker end)
+  defp include_current_worker(worker_counts, lease_owner) do
+    worker_counts
+    |> List.keymember?(lease_owner, 0)
     |> case do
-      nil -> length(workers) + 1
-      _ -> length(workers)
+      true -> worker_counts
+      false -> [{lease_owner, 0} | worker_counts]
     end
   end
 
-  defp calculate_target_load(state) do
-    incomplete_leases_count =
-      AppState.all_incomplete_leases(state.app_name, state.stream_name, state.app_state_opts)
-
-    workers_count =
-      AppState.total_incomplete_lease_counts_by_worker(
-        state.app_name,
-        state.stream_name,
-        state.app_state_opts
-      )
+  defp target_load(worker_counts) do
+    total_leases = worker_counts |> Enum.map(fn {_owner, count} -> count end) |> Enum.sum()
 
-    total_shards = length(incomplete_leases_count)
-    total_workers = total_workers_count(workers_count, state.lease_owner)
-
-    if total_workers > 0, do: ceil(total_shards / total_workers), else: 0
-  end
-
-  defp get_current_worker_load(state) do
-    state.app_name
-    |> AppState.get_leases_by_worker(
-      state.stream_name,
-      state.lease_owner,
-      state.app_state_opts
-    )
-    |> case do
-      [] -> 0
-      leases -> length(leases)
-    end
+    ceil(total_leases / length(worker_counts))
   end
 end
diff --git a/test/kinesis_client/stream/app_state/dynamo_test.exs b/test/kinesis_client/stream/app_state/dynamo_test.exs
index faea028..8f0766a 100644
--- a/test/kinesis_client/stream/app_state/dynamo_test.exs
+++ b/test/kinesis_client/stream/app_state/dynamo_test.exs
@@ -144,6 +144,54 @@ defmodule KinesisClient.Stream.AppState.DynamoTest do
     end
   end
 
+  describe "load balancing queries" do
+    setup do
+      app_name = "foo_app_#{random_string()}"
+      :ok = AppState.initialize(app_name, [])
+      %{lb_app: app_name}
+    end
+
+    test "get_leases_by_worker/4 returns only the given worker's leases", %{lb_app: app_name} do
+      lease_owner = worker_ref()
+      other_owner = worker_ref()
+      assert :ok == AppState.create_lease(app_name, "", "shard-000001", lease_owner, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000002", lease_owner, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000003", other_owner, [])
+
+      leases = AppState.get_leases_by_worker(app_name, "", lease_owner, [])
+
+      assert length(leases) == 2
+      assert Enum.all?(leases, &(&1.lease_owner == lease_owner))
+    end
+
+    test "all_incomplete_leases/3 excludes completed shards", %{lb_app: app_name} do
+      lease_owner = worker_ref()
+      assert :ok == AppState.create_lease(app_name, "", "shard-000001", lease_owner, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000002", lease_owner, [])
+      assert :ok == AppState.close_shard(app_name, "", "shard-000002", lease_owner, [])
+
+      assert [%ShardLease{shard_id: "shard-000001", completed: false}] =
+               AppState.all_incomplete_leases(app_name, "", [])
+    end
+
+    test "total_incomplete_lease_counts_by_worker/3 groups counts by owner", %{lb_app: app_name} do
+      worker_1 = worker_ref()
+      worker_2 = worker_ref()
+      assert :ok == AppState.create_lease(app_name, "", "shard-000001", worker_1, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000002", worker_1, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000003", worker_2, [])
+      assert :ok == AppState.create_lease(app_name, "", "shard-000004", worker_2, [])
+      assert :ok == AppState.close_shard(app_name, "", "shard-000004", worker_2, [])
+
+      counts =
+        app_name
+        |> AppState.total_incomplete_lease_counts_by_worker("", [])
+        |> Enum.sort()
+
+      assert counts == Enum.sort([{worker_1, 2}, {worker_2, 1}])
+    end
+  end
+
   defp confirm_table_created(app_name, attempts \\ 1) do
     case app_name |> Dynamo.describe_table() |> ExAws.request() do
       {:ok, %{"Table" => %{"TableStatus" => "CREATING"}}} ->
diff --git a/test/kinesis_client/stream/coordinator_test.exs b/test/kinesis_client/stream/coordinator_test.exs
index 020a45c..f61b9b0 100644
--- a/test/kinesis_client/stream/coordinator_test.exs
+++ b/test/kinesis_client/stream/coordinator_test.exs
@@ -149,9 +149,9 @@ defmodule KinesisClient.Stream.CoordinatorTest do
     assert Process.alive?(pid_1) == true
 
     refute_receive {:shard_started, %{pid: _, shard_id: "shardId-000000000002"}}, 200
-    assert Process.alive?(pid_0) == false
+    assert Process.alive?(pid_0) == true
     refute_receive {:shard_started, %{pid: _, shard_id: "shardId-000000000003"}}, 200
-    assert Process.alive?(pid_1) == false
+    assert Process.alive?(pid_1) == true
 
     assert Enum.empty?(shards) == false
   end
diff --git a/test/kinesis_client/stream/rebalancer_test.exs b/test/kinesis_client/stream/rebalancer_test.exs
new file mode 100644
index 0000000..092d181
--- /dev/null
+++ b/test/kinesis_client/stream/rebalancer_test.exs
@@ -0,0 +1,157 @@
+defmodule KinesisClient.Stream.RebalancerTest do
+  use KinesisClient.Case
+
+  import KinesisClient.Util
+
+  alias KinesisClient.Stream.AppState.ShardLease
+  alias KinesisClient.Stream.Rebalancer
+  alias KinesisClient.Stream.Shard.LeaseV2
+
+  test "notifies :all_balanced and requests no steals when the load is balanced" do
+    opts = build_rebalancer_opts()
+    lease_owner = opts[:lease_owner]
+
+    stub(AppStateMock, :total_incomplete_lease_counts_by_worker, fn _app_name,
+                                                                    _stream_name,
+                                                                    _opts ->
+      [{lease_owner, 2}, {worker_ref(), 2}]
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:all_balanced, _state}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  test "requests a steal from the local lease process of an overloaded worker's shard" do
+    opts = build_rebalancer_opts()
+    victim = worker_ref()
+
+    victim_leases = [
+      build_shard_lease(shard_id: "shard-000001", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000002", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000003", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000004", lease_owner: victim)
+    ]
+
+    # Only shard-000002 has a local lease process, so despite the random
+    # candidate order the steal request can only go there.
+    register_lease_process(opts, "shard-000002")
+
+    AppStateMock
+    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
+      [{victim, 4}]
+    end)
+    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, in_lease_owner, _opts ->
+      assert in_lease_owner == victim
+      victim_leases
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:steal_requested, "shard-000002"}, 1_000
+    assert_receive {:lease_message, "shard-000002", :steal_lease}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  test "requests at most max_leases_to_steal steals per tick" do
+    # Interval sized so the first tick (jittered to 450-750ms) lands inside
+    # the assert window, while the second tick can't land inside the refute
+    # window below — the next tick would legitimately steal again.
+    opts = build_rebalancer_opts(rebalance_interval: 600)
+    victim = worker_ref()
+
+    victim_leases = [
+      build_shard_lease(shard_id: "shard-000001", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000002", lease_owner: victim)
+    ]
+
+    register_lease_process(opts, "shard-000001")
+    register_lease_process(opts, "shard-000002")
+
+    AppStateMock
+    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
+      [{victim, 4}]
+    end)
+    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
+      victim_leases
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:lease_message, _shard_id, :steal_lease}, 1_000
+    refute_receive {:lease_message, _shard_id, :steal_lease}, 200
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  test "skips completed shards when picking a steal candidate" do
+    opts = build_rebalancer_opts()
+    victim = worker_ref()
+
+    victim_leases = [
+      build_shard_lease(shard_id: "shard-000001", lease_owner: victim, completed: true),
+      build_shard_lease(shard_id: "shard-000002", lease_owner: victim)
+    ]
+
+    register_lease_process(opts, "shard-000001")
+    register_lease_process(opts, "shard-000002")
+
+    AppStateMock
+    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
+      [{victim, 4}]
+    end)
+    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
+      victim_leases
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:lease_message, "shard-000002", :steal_lease}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  # Registers a stand-in for the shard's LeaseV2 process that forwards any
+  # message it receives back to the test process, tagged with the shard_id.
+  defp register_lease_process(opts, shard_id) do
+    test_pid = self()
+
+    pid =
+      spawn_link(fn ->
+        receive do
+          msg -> send(test_pid, {:lease_message, shard_id, msg})
+        end
+      end)
+
+    name = register_name(LeaseV2, opts[:app_name], opts[:stream_name], [shard_id])
+    Process.register(pid, name)
+  end
+
+  defp build_shard_lease(overrides) do
+    default = [
+      shard_id: "shard-000001",
+      lease_owner: worker_ref(),
+      lease_count: 1,
+      completed: false
+    ]
+
+    struct(ShardLease, Keyword.merge(default, overrides))
+  end
+
+  defp build_rebalancer_opts(overrides \\ []) do
+    Keyword.merge(
+      [
+        app_name: "my_streaming_app",
+        stream_name: "stream-#{:rand.uniform(100_000)}",
+        lease_owner: worker_ref(),
+        app_state_opts: [adapter: :test],
+        rebalance_interval: 200,
+        notify: self()
+      ],
+      overrides
+    )
+  end
+end
diff --git a/test/kinesis_client/stream/shard/lease_v2_test.exs b/test/kinesis_client/stream/shard/lease_v2_test.exs
index a567e34..d74e3a5 100644
--- a/test/kinesis_client/stream/shard/lease_v2_test.exs
+++ b/test/kinesis_client/stream/shard/lease_v2_test.exs
@@ -4,6 +4,19 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
   alias KinesisClient.Stream.AppState.ShardLease
   alias KinesisClient.Stream.Shard.LeaseV2
 
+  defmodule NotifyingPipeline do
+    @moduledoc false
+    def start(state) do
+      send(state.notify, {:pipeline_started, state.shard_id})
+      :ok
+    end
+
+    def stop(state) do
+      send(state.notify, {:pipeline_stopped, state.shard_id})
+      :ok
+    end
+  end
+
   test "creates and takes AppState.ShardLease if none already exists" do
     lease_opts = build_lease_opts(pipeline: KinesisClient.TestPipeline)
 
@@ -53,169 +66,204 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
     stop_supervised(LeaseV2)
   end
 
-  describe "when shard_lease already exists" do
-    test "and lease owner does not currently have any leases and it's all balanced so move on" do
-      shard_lease_count = 12
-      lease_opts = build_lease_opts()
-      other_worker = lease_opts[:lease_owner]
-      shard_lease = build_shard_lease(lease_count: shard_lease_count, lease_owner: other_worker)
+  test "tracks an existing lease owned by another worker without taking it" do
+    shard_lease_count = 12
+    lease_opts = build_lease_opts()
+    shard_lease = build_shard_lease(lease_count: shard_lease_count, lease_owner: worker_ref())
+
+    stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      shard_lease
+    end)
+
+    {:ok, pid} = start_supervised({LeaseV2, lease_opts})
+
+    assert_receive {:initialized, lease_state}, 1_000
+    assert lease_state.lease_holder == false
+    assert lease_state.lease_count == shard_lease_count
+    assert Process.alive?(pid)
+    stop_supervised(LeaseV2)
+  end
+
+  test "stops the pipeline and releases lease_holder when renewal fails" do
+    current_worker = worker_ref()
+
+    lease_opts =
+      build_lease_opts(
+        lease_owner: current_worker,
+        renew_interval: 200,
+        pipeline: NotifyingPipeline
+      )
+
+    owned_shard_lease =
+      build_shard_lease(
+        lease_count: 1,
+        lease_owner: current_worker,
+        shard_id: lease_opts[:shard_id]
+      )
+
+    AppStateMock
+    |> expect(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      :not_found
+    end)
+    |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      owned_shard_lease
+    end)
+    |> stub(:create_lease, fn _app_name, _stream_name, _shard_id, _lease_owner, _opts ->
+      :ok
+    end)
+    |> stub(:renew_lease, fn _app_name, _stream_name, _shard_lease, _opts ->
+      {:error, :lease_renew_failed}
+    end)
+
+    {:ok, pid} = start_supervised({LeaseV2, lease_opts})
+
+    assert_receive {:initialized, lease_state}, 1_000
+    assert lease_state.lease_holder == true
+    assert_receive {:pipeline_started, _shard_id}, 1_000
+
+    assert_receive {:lease_renew_failed, lease_state}, 1_000
+    assert lease_state.lease_holder == false
+    assert_receive {:pipeline_stopped, _shard_id}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(LeaseV2)
+  end
+
+  describe ":steal_lease message" do
+    test "steals the lease from its current owner and starts the pipeline" do
       current_worker = worker_ref()
 
+      lease_opts =
+        build_lease_opts(lease_owner: current_worker, pipeline: NotifyingPipeline)
+
+      shard_lease =
+        build_shard_lease(
+          lease_count: 8,
+          lease_owner: worker_ref(),
+          shard_id: lease_opts[:shard_id]
+        )
+
       stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
         shard_lease
       end)
-      |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-        []
-      end)
-      |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-        [{other_worker, 1}]
-      end)
-      |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-        [shard_lease]
-      end)
-      |> stub(:lease_owner_with_most_leases, fn _app_name, _stream_name, _opts ->
-        [shard_lease]
-      end)
       |> stub(:take_lease, fn app_name, stream_name, shard_id, new_owner, lc, _opts ->
         assert app_name == lease_opts[:app_name]
         assert stream_name == lease_opts[:stream_name]
         assert shard_id == lease_opts[:shard_id]
         assert new_owner == current_worker
-        assert lc == 12
+        assert lc == shard_lease.lease_count
 
         {:ok, lc + 1}
       end)
 
-      {:ok, pid} =
-        start_supervised(
-          {LeaseV2, build_lease_opts(shard_id: shard_lease.shard_id, lease_owner: current_worker)}
-        )
+      {:ok, pid} = start_supervised({LeaseV2, lease_opts})
 
-      assert_receive {:all_balanced, lease_state}, 1_000
+      assert_receive {:initialized, lease_state}, 1_000
       assert lease_state.lease_holder == false
-      assert lease_state.lease_count == shard_lease_count
+
+      send(pid, :steal_lease)
+
+      assert_receive {:lease_stolen, lease_state}, 1_000
+      assert lease_state.lease_holder == true
+      assert lease_state.lease_count == shard_lease.lease_count + 1
+      assert lease_state.lease_owner == current_worker
+      assert_receive {:pipeline_started, _shard_id}, 1_000
       assert Process.alive?(pid)
       stop_supervised(LeaseV2)
     end
 
-    test "and all is balanced so move on" do
-      shard_lease_count = 12
-      lease_opts = build_lease_opts()
-      current_worker = lease_opts[:lease_owner]
-      shard_lease_1 = build_shard_lease(lease_count: shard_lease_count, lease_owner: current_worker)
+    test "steals with the freshly read lease_count when the state count is stale" do
+      current_worker = worker_ref()
       other_worker = worker_ref()
-      shard_lease_2 = build_shard_lease(lease_count: 5, lease_owner: other_worker)
+      lease_opts = build_lease_opts(lease_owner: current_worker, pipeline: NotifyingPipeline)
+
+      # The count read at init (synced into state) is 8, but by the time the
+      # steal request arrives the owner has renewed the lease to 9. The steal
+      # must use the fresh count or the optimistic lock will always fail.
+      stale_shard_lease =
+        build_shard_lease(
+          lease_count: 8,
+          lease_owner: other_worker,
+          shard_id: lease_opts[:shard_id]
+        )
 
-      stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
-        shard_lease_1
-      end)
-      |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-        [shard_lease_1]
+      fresh_shard_lease =
+        build_shard_lease(
+          lease_count: 9,
+          lease_owner: other_worker,
+          shard_id: lease_opts[:shard_id]
+        )
+
+      AppStateMock
+      |> expect(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+        stale_shard_lease
       end)
-      |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-        [{current_worker, 1}, {worker_ref(), 1}]
+      |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+        fresh_shard_lease
       end)
-      |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-        [shard_lease_1, shard_lease_2]
+      |> stub(:take_lease, fn _app_name, _stream_name, _shard_id, new_owner, lc, _opts ->
+        assert new_owner == current_worker
+        assert lc == fresh_shard_lease.lease_count
+
+        {:ok, lc + 1}
       end)
 
       {:ok, pid} = start_supervised({LeaseV2, lease_opts})
 
-      assert_receive {:all_balanced, lease_state}, 1_000
-      assert lease_state.shard_id == shard_lease_1.shard_id
-      assert lease_state.lease_count == shard_lease_count
-      assert lease_state.lease_owner == current_worker
+      assert_receive {:initialized, lease_state}, 1_000
+      assert lease_state.lease_count == stale_shard_lease.lease_count
+
+      send(pid, :steal_lease)
+
+      assert_receive {:lease_stolen, lease_state}, 1_000
+      assert lease_state.lease_holder == true
+      assert lease_state.lease_count == fresh_shard_lease.lease_count + 1
       assert Process.alive?(pid)
       stop_supervised(LeaseV2)
     end
 
-    test "and it's not balanced so steal from overloaded worker" do
-      lease_opts = build_lease_opts()
-      worker_1 = worker_ref()
-
-      shard_lease_1 =
-        build_shard_lease(lease_count: 12, lease_owner: worker_1, shard_id: "shard-000001")
-
-      worker_2 = worker_ref()
-
-      shard_lease_2 =
-        build_shard_lease(lease_count: 10, lease_owner: worker_2, shard_id: "shard-000002")
-
-      shard_lease_3 =
-        build_shard_lease(lease_count: 10, lease_owner: worker_2, shard_id: "shard-000003")
+    test "ignores the request when already the lease holder" do
+      lease_opts = build_lease_opts(pipeline: KinesisClient.TestPipeline)
 
-      shard_lease_4 =
-        build_shard_lease(lease_count: 10, lease_owner: worker_2, shard_id: "shard-000004")
+      AppStateMock
+      |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+        :not_found
+      end)
+      |> stub(:create_lease, fn _app_name, _stream_name, _shard_id, _lease_owner, _opts ->
+        :ok
+      end)
 
-      shard_lease_5 =
-        build_shard_lease(lease_count: 10, lease_owner: worker_2, shard_id: "shard-000005")
+      {:ok, pid} = start_supervised({LeaseV2, lease_opts})
 
-      worker_3 = worker_ref()
+      assert_receive {:initialized, %{lease_holder: true}}, 1_000
 
-      shard_lease_6 =
-        build_shard_lease(lease_count: 8, lease_owner: worker_3, shard_id: "shard-000006")
+      send(pid, :steal_lease)
 
-      shard_lease_7 =
-        build_shard_lease(lease_count: 8, lease_owner: worker_3, shard_id: "shard-000007")
+      refute_receive {:lease_stolen, _}, 200
+      assert Process.alive?(pid)
+      stop_supervised(LeaseV2)
+    end
 
-      shard_lease_8 =
-        build_shard_lease(lease_count: 8, lease_owner: worker_3, shard_id: "shard-000008")
+    test "ignores the request when the shard is completed" do
+      lease_opts = build_lease_opts()
 
-      shard_lease_9 =
-        build_shard_lease(lease_count: 8, lease_owner: worker_3, shard_id: "shard-000009")
+      shard_lease =
+        build_shard_lease(
+          lease_owner: worker_ref(),
+          shard_id: lease_opts[:shard_id],
+          completed: true
+        )
 
       stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
-        shard_lease_6
-      end)
-      |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-        [shard_lease_1]
-      end)
-      |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-        [{worker_1, 1}, {worker_2, 4}, {worker_3, 4}]
-      end)
-      |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-        [
-          shard_lease_1,
-          shard_lease_2,
-          shard_lease_3,
-          shard_lease_4,
-          shard_lease_5,
-          shard_lease_6,
-          shard_lease_7,
-          shard_lease_8,
-          shard_lease_9
-        ]
-      end)
-      |> stub(:lease_owner_with_most_leases, fn _app_name, _stream_name, _opts ->
-        [shard_lease_6, shard_lease_7, shard_lease_8, shard_lease_9]
+        shard_lease
       end)
-      |> stub(:take_lease, fn app_name, stream_name, shard_id, new_owner, lc, _opts ->
-        assert app_name == lease_opts[:app_name]
-        assert stream_name == lease_opts[:stream_name]
 
-        assert shard_id in [
-                 shard_lease_6.shard_id,
-                 shard_lease_7.shard_id,
-                 shard_lease_8.shard_id,
-                 shard_lease_9.shard_id
-               ]
-
-        assert new_owner == worker_1
-        assert lc == 8
+      {:ok, pid} = start_supervised({LeaseV2, lease_opts})
 
-        {:ok, lc + 1}
-      end)
+      assert_receive {:initialized, %{lease_holder: false}}, 1_000
 
-      {:ok, pid} =
-        start_supervised(
-          {LeaseV2, build_lease_opts(shard_id: shard_lease_6.shard_id, lease_owner: worker_1)}
-        )
+      send(pid, :steal_lease)
 
-      assert_receive {:lease_stolen, lease_state}, 1_000
-      assert lease_state.lease_holder == true
-      assert lease_state.lease_count == 9
-      assert lease_state.lease_owner == worker_1
-      assert lease_state.shard_id == shard_lease_6.shard_id
+      refute_receive {:lease_stolen, _}, 200
       assert Process.alive?(pid)
       stop_supervised(LeaseV2)
     end
@@ -230,15 +278,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
     |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
       shard_lease
     end)
-    |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-      [shard_lease]
-    end)
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{lease_opts[:application], 1}]
-    end)
-    |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-      [shard_lease]
-    end)
     |> stub(:take_lease, fn app_name, stream_name, shard_id, new_owner, lc, _opts ->
       assert app_name == lease_opts[:app_name]
       assert stream_name == lease_opts[:stream_name]
@@ -269,15 +308,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
     stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
       shard_lease
     end)
-    |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-      [shard_lease]
-    end)
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{lease_opts[:application], 1}]
-    end)
-    |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-      [shard_lease]
-    end)
 
     {:ok, _pid} = start_supervised({LeaseV2, lease_opts})
 
@@ -292,34 +322,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
     stop_supervised(LeaseV2)
   end
 
-  test "run load balancing" do
-    shard_lease_count = 12
-    lease_opts = build_lease_opts(rebalance_interval: 500)
-    shard_lease = build_shard_lease(lease_count: shard_lease_count)
-
-    stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
-      shard_lease
-    end)
-    |> stub(:get_leases_by_worker, fn _in_app_name, _in_stream_name, _lease_owner, _ ->
-      [shard_lease]
-    end)
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{lease_opts[:application], 1}]
-    end)
-    |> stub(:all_incomplete_leases, fn _app_name, _stream_name, _opts ->
-      [shard_lease]
-    end)
-
-    {:ok, pid} = start_supervised({LeaseV2, lease_opts})
-
-    assert_receive {:initialized, %{lease_count_increment_time: _lcit} = lease_state}, 1_000
-    assert lease_state.lease_holder == false
-    assert_receive {:all_balanced, lease_state}, 1_000
-    assert lease_state.lease_count == shard_lease.lease_count
-    Process.alive?(pid)
-    stop_supervised(LeaseV2)
-  end
-
   defp build_lease_opts(overrides \\ []) do
     Keyword.merge(
       [
diff --git a/test/kinesis_client/stream/shard/load_balance_test.exs b/test/kinesis_client/stream/shard/load_balance_test.exs
new file mode 100644
index 0000000..a7b4217
--- /dev/null
+++ b/test/kinesis_client/stream/shard/load_balance_test.exs
@@ -0,0 +1,61 @@
+defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
+  use ExUnit.Case, async: true
+
+  alias KinesisClient.Stream.Shard.LoadBalance
+
+  describe "decide/2" do
+    test "balanced when there are no leases at all" do
+      assert LoadBalance.decide([], "worker-1") == :balanced
+    end
+
+    test "balanced when the current worker is the only worker" do
+      assert LoadBalance.decide([{"worker-1", 5}], "worker-1") == :balanced
+    end
+
+    test "balanced when leases are spread evenly" do
+      counts = [{"worker-1", 3}, {"worker-2", 3}, {"worker-3", 3}]
+
+      assert LoadBalance.decide(counts, "worker-1") == :balanced
+    end
+
+    test "steals from the most loaded worker when under target" do
+      counts = [{"worker-1", 1}, {"worker-2", 4}, {"worker-3", 4}]
+
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2"}
+    end
+
+    test "a worker holding no leases counts itself and steals" do
+      # worker-1 holds nothing so it is absent from the grouped counts. It
+      # must still count itself as a worker, see the imbalance, and steal.
+      counts = [{"worker-2", 4}, {"worker-3", 4}]
+
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2"}
+    end
+
+    test "does not steal when the lead is only one lease" do
+      # 7 leases over 2 workers can never be more even than 4/3: stealing
+      # would just flip the imbalance back and forth forever.
+      counts = [{"worker-1", 3}, {"worker-2", 4}]
+
+      assert LoadBalance.decide(counts, "worker-1") == :balanced
+    end
+
+    test "does not steal when already at target, even if another worker is over" do
+      # worker-1 is at the target of 3; the deficit is worker-3's to fix.
+      counts = [{"worker-1", 3}, {"worker-2", 6}, {"worker-3", 0}]
+
+      assert LoadBalance.decide(counts, "worker-1") == :balanced
+      assert LoadBalance.decide(counts, "worker-3") == {:steal_from, "worker-2"}
+    end
+
+    test "converges to an even spread as steals are applied" do
+      # Simulate the {4, 4, 0} cluster rebalancing one steal at a time.
+      assert LoadBalance.decide([{"w2", 4}, {"w3", 4}], "w1") == {:steal_from, "w2"}
+      assert LoadBalance.decide([{"w1", 1}, {"w2", 3}, {"w3", 4}], "w1") == {:steal_from, "w3"}
+
+      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w1") == :balanced
+      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w2") == :balanced
+      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w3") == :balanced
+    end
+  end
+end
diff --git a/test/support/test_pipeline.ex b/test/support/test_pipeline.ex
index fbdd5c4..7c2b90c 100644
--- a/test/support/test_pipeline.ex
+++ b/test/support/test_pipeline.ex
@@ -3,4 +3,8 @@ defmodule KinesisClient.TestPipeline do
   def start(_opts) do
     :ok
   end
+
+  def stop(_opts) do
+    :ok
+  end
 end

From b9d44b5820a5de045f060644ac2f1e292f2244ef Mon Sep 17 00:00:00 2001
From: Jon Kusumo <tahir.kusumo@paradox.ai>
Date: Wed, 1 Jul 2026 21:53:04 -0400
Subject: [PATCH 2/4] Address review: fix reclaim wedge, isolate rebalancer
 faults, clamp steals
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Reclaim a lease the AppState names us as owner of while lease_holder is
  false (e.g. a renewal that "failed" after actually being applied via a
  lost-response retry). Previously this state was unrecoverable: renewing
  requires lease_holder and both adapters reject taking a lease you
  already own, so the shard sat unconsumed forever.
- Rescue the rebalance tick: adapters raise on transient failures
  (throttled Dynamo scan, DB blip) and the Rebalancer shares a
  :one_for_all supervisor with the Coordinator and every pipeline — a
  failed best-effort balance check must not tear down the data path.
- Clamp steals to the worker's lease deficit: LoadBalance.decide/2 now
  returns the deficit and steal_from takes min(deficit,
  max_leases_to_steal), so max_leases_to_steal > 1 no longer overshoots
  the target and oscillates.
- Carry the chosen victim in the :steal_lease message and skip the steal
  if the lease changed owners since the decision.
- Delegate Mimic.get_leases_by_worker to the migration adapters instead
  of hard-coding [] (rebalancing silently no-oped during migrations).
- Expose LeaseV2.whereis/3 instead of re-deriving the registration name
  in the Rebalancer; guard jitter/1 against sub-2ms intervals; drop the
  dead :spread_lease option and the misleading :global comment.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 README.md                                     |  6 +-
 lib/kinesis_client/stream.ex                  |  3 -
 lib/kinesis_client/stream/app_state/mimic.ex  |  7 +-
 lib/kinesis_client/stream/rebalancer.ex       | 32 +++++---
 lib/kinesis_client/stream/shard.ex            |  1 -
 lib/kinesis_client/stream/shard/lease_v2.ex   | 72 ++++++++++++++---
 .../stream/shard/load_balance.ex              | 16 ++--
 .../kinesis_client/stream/rebalancer_test.exs | 58 ++++++++++++-
 .../stream/shard/lease_v2_test.exs            | 81 +++++++++++++++++--
 .../stream/shard/load_balance_test.exs        | 22 +++--
 10 files changed, 246 insertions(+), 52 deletions(-)

diff --git a/README.md b/README.md
index c2779ad..2b4c6f5 100644
--- a/README.md
+++ b/README.md
@@ -106,9 +106,9 @@ Each worker runs a single rebalancer process that periodically (every
 `rebalance_interval`, jittered +/- 25%) compares how many incomplete leases
 each worker holds. When this worker is below the target load
 (`ceil(total_shards / total_workers)`) and another worker leads it by more
-than one lease, it steals a lease from the worker holding the most — at most
-`max_leases_to_steal` (default 1) per check, so workers converge on an even
-spread without overshooting. Workers that crash stop renewing their leases,
+than one lease, it steals from the worker holding the most — up to its lease
+deficit per check, capped at `max_leases_to_steal` (default 1) — so workers
+converge on an even spread without overshooting. Workers that crash stop renewing their leases,
 and after `lease_expiry` the remaining workers take those leases over.
 
 ## TODO
diff --git a/lib/kinesis_client/stream.ex b/lib/kinesis_client/stream.ex
index 0061b67..3f219a1 100644
--- a/lib/kinesis_client/stream.ex
+++ b/lib/kinesis_client/stream.ex
@@ -65,7 +65,6 @@ defmodule KinesisClient.Stream do
       |> optional_kw(:app_state_opts, fetch_value_for_key!(opts, :app_state_opts))
       |> optional_kw(:lease_renew_interval, Keyword.get(opts, :lease_renew_interval))
       |> optional_kw(:lease_expiry, Keyword.get(opts, :lease_expiry))
-      |> optional_kw(:spread_lease, Keyword.get(opts, :spread_lease))
       |> optional_kw(:poll_interval, Keyword.get(opts, :poll_interval))
       |> optional_kw(:shard_iterator_type, Keyword.get(opts, :shard_iterator_type))
       |> optional_kw(:timestamp, Keyword.get(opts, :timestamp))
@@ -108,8 +107,6 @@ defmodule KinesisClient.Stream do
       nil ->
         register_name(KinesisClient.Stream.Coordinator, opts[:app_name], opts[:stream_name])
 
-      # Shard processes may be running on nodes different from the Coordinator if passed
-      # :shard_supervisor is distributed,so use :global to allow inter-node communication.
       _ ->
         {:global,
          register_name(KinesisClient.Stream.Coordinator, opts[:app_name], opts[:stream_name])}
diff --git a/lib/kinesis_client/stream/app_state/mimic.ex b/lib/kinesis_client/stream/app_state/mimic.ex
index e7ad60b..fa40448 100644
--- a/lib/kinesis_client/stream/app_state/mimic.ex
+++ b/lib/kinesis_client/stream/app_state/mimic.ex
@@ -42,8 +42,11 @@ defmodule KinesisClient.Stream.AppState.Mimic do
   end
 
   @impl true
-  def get_leases_by_worker(_app_name, _stream_name, _lease_owner, _opts) do
-    []
+  def get_leases_by_worker(app_name, stream_name, lease_owner, opts) do
+    {from, to} = modules(opts)
+
+    to.get_leases_by_worker(app_name, stream_name, lease_owner, opts)
+    from.get_leases_by_worker(app_name, stream_name, lease_owner, opts)
   end
 
   @impl true
diff --git a/lib/kinesis_client/stream/rebalancer.ex b/lib/kinesis_client/stream/rebalancer.ex
index 569f7cb..0fae485 100644
--- a/lib/kinesis_client/stream/rebalancer.ex
+++ b/lib/kinesis_client/stream/rebalancer.ex
@@ -71,7 +71,20 @@ defmodule KinesisClient.Stream.Rebalancer do
   @impl GenServer
   def handle_info(:rebalance, state) do
     schedule_rebalance(state)
+    run_rebalance(state)
+    {:noreply, state}
+  rescue
+    # Rebalancing is best-effort and must never take down the stream: the
+    # adapters raise on transient failures (a throttled Dynamo scan, a DB
+    # blip), and this process shares a :one_for_all supervisor with the
+    # Coordinator and every shard pipeline. Log it and try again next tick.
+    error ->
+      Logger.error("Rebalancer: Rebalance tick failed: #{inspect(error)}")
+      notify({:rebalance_failed, error}, state)
+      {:noreply, state}
+  end
 
+  defp run_rebalance(state) do
     state.app_name
     |> AppState.total_incomplete_lease_counts_by_worker(state.stream_name, state.app_state_opts)
     |> LoadBalance.decide(state.lease_owner)
@@ -79,28 +92,26 @@ defmodule KinesisClient.Stream.Rebalancer do
       :balanced ->
         notify({:all_balanced, state}, state)
 
-      {:steal_from, victim} ->
-        steal_from(victim, state)
+      {:steal_from, victim, deficit} ->
+        steal_from(victim, deficit, state)
     end
-
-    {:noreply, state}
   end
 
-  defp steal_from(victim, state) do
+  defp steal_from(victim, deficit, state) do
     state.app_name
     |> AppState.get_leases_by_worker(state.stream_name, victim, state.app_state_opts)
     |> Enum.reject(& &1.completed)
     |> Enum.shuffle()
     |> Enum.map(&local_lease_process(&1, state))
     |> Enum.reject(&is_nil/1)
-    |> Enum.take(state.max_leases_to_steal)
+    |> Enum.take(min(deficit, state.max_leases_to_steal))
     |> Enum.each(fn {shard_id, pid} ->
       Logger.debug(
         "Rebalancer: Requesting steal of shard #{shard_id} from #{victim}: " <>
           "[lease_owner: #{state.lease_owner}]"
       )
 
-      send(pid, :steal_lease)
+      send(pid, {:steal_lease, victim})
       notify({:steal_requested, shard_id}, state)
     end)
   end
@@ -109,9 +120,8 @@ defmodule KinesisClient.Stream.Rebalancer do
   # writer per shard on this worker. Shards without a running local lease
   # process (not started yet, or already shut down) are skipped this tick.
   defp local_lease_process(%{shard_id: shard_id}, state) do
-    LeaseV2
-    |> register_name(state.app_name, state.stream_name, [shard_id])
-    |> Process.whereis()
+    state.app_name
+    |> LeaseV2.whereis(state.stream_name, shard_id)
     |> case do
       nil -> nil
       pid -> {shard_id, pid}
@@ -124,7 +134,7 @@ defmodule KinesisClient.Stream.Rebalancer do
 
   # +/- 25% so workers don't tick in lockstep and stampede the same lease.
   defp jitter(interval) do
-    interval + :rand.uniform(div(interval, 2)) - div(interval, 4)
+    interval + :rand.uniform(max(div(interval, 2), 1)) - div(interval, 4)
   end
 
   defp notify(_msg, %{notify: nil}) do
diff --git a/lib/kinesis_client/stream/shard.ex b/lib/kinesis_client/stream/shard.ex
index b8489c6..a9d5f54 100644
--- a/lib/kinesis_client/stream/shard.ex
+++ b/lib/kinesis_client/stream/shard.ex
@@ -24,7 +24,6 @@ defmodule KinesisClient.Stream.Shard do
       |> optional_kw(:renew_interval, Keyword.get(opts, :lease_renew_interval))
       |> optional_kw(:lease_expiry, Keyword.get(opts, :lease_expiry))
       |> optional_kw(:pipeline, Keyword.get(opts, :pipeline))
-      |> optional_kw(:spread_lease, Keyword.get(opts, :spread_lease))
 
     pipeline_opts =
       [
diff --git a/lib/kinesis_client/stream/shard/lease_v2.ex b/lib/kinesis_client/stream/shard/lease_v2.ex
index 5e033bc..f1374ed 100644
--- a/lib/kinesis_client/stream/shard/lease_v2.ex
+++ b/lib/kinesis_client/stream/shard/lease_v2.ex
@@ -28,6 +28,17 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     )
   end
 
+  @doc """
+  Returns the pid of the locally registered lease process for `shard_id`, or
+  `nil` if none is running.
+  """
+  @spec whereis(String.t(), String.t(), String.t()) :: pid() | nil
+  def whereis(app_name, stream_name, shard_id) do
+    __MODULE__
+    |> register_name(app_name, stream_name, [shard_id])
+    |> Process.whereis()
+  end
+
   defstruct [
     :app_name,
     :stream_name,
@@ -131,13 +142,13 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
   end
 
   # Sent by KinesisClient.Stream.Rebalancer when it decides this worker should
-  # steal this shard's lease from an overloaded worker.
+  # steal this shard's lease from `victim`, the overloaded worker.
   @impl GenServer
-  def handle_info(:steal_lease, %{lease_holder: true} = state) do
+  def handle_info({:steal_lease, _victim}, %{lease_holder: true} = state) do
     {:noreply, state}
   end
 
-  def handle_info(:steal_lease, state) do
+  def handle_info({:steal_lease, victim}, state) do
     state
     |> get_shard_lease()
     |> case do
@@ -150,7 +161,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
         {:noreply, state}
 
       shard_lease ->
-        {:noreply, maybe_steal(shard_lease, state)}
+        {:noreply, maybe_steal(shard_lease, victim, state)}
     end
   end
 
@@ -174,6 +185,40 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
+  # The AppState row names this worker as owner while lease_holder is false —
+  # e.g. a renewal that "failed" after actually being applied (lost response
+  # + AWS retry fails the conditional check), or this process restarted after
+  # taking the lease. Neither renewing (requires lease_holder) nor taking
+  # (both adapters reject taking a lease you already own) can recover from
+  # here, so without this clause the shard would sit unconsumed forever.
+  # Renew under the optimistic lock and resume the pipeline; if another
+  # worker took the lease in the meantime, the renewal fails and we keep
+  # tracking.
+  defp reclaim_shard_lease(shard_lease, %{app_state_opts: opts, app_name: app_name} = state) do
+    expected = shard_lease.lease_count + 1
+
+    case AppState.renew_lease(app_name, state.stream_name, shard_lease, opts) do
+      {:ok, ^expected} ->
+        Logger.info(
+          "ShardLease: Reclaimed own lease that was not being held: " <>
+            "[shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
+        )
+
+        expected
+        |> set_lease_count(true, state)
+        |> tap(&notify({:lease_reclaimed, &1}, &1))
+        |> tap(fn state -> state.pipeline.start(state) end)
+
+      {:error, error} ->
+        Logger.error(
+          "ShardLease: Failed to reclaim own lease, error: #{inspect(error)}, " <>
+            "[shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
+        )
+
+        state
+    end
+  end
+
   defp renew_shard_lease(shard_lease, %{app_state_opts: opts, app_name: app_name} = state) do
     expected = shard_lease.lease_count + 1
 
@@ -239,17 +284,21 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     end
   end
 
-  defp maybe_steal(%{completed: true}, state), do: state
+  defp maybe_steal(%{completed: true}, _victim, state), do: state
 
-  defp maybe_steal(%{lease_owner: current_owner} = shard_lease, %{lease_owner: me} = state)
-       when current_owner != me do
+  defp maybe_steal(%{lease_owner: victim} = shard_lease, victim, %{lease_owner: me} = state)
+       when victim != me do
     steal_shard_lease(shard_lease, state)
   end
 
-  defp maybe_steal(_shard_lease, state) do
+  # The lease changed hands between the rebalancer's decision and this
+  # message (or is already ours) — stealing now would hit the wrong worker,
+  # possibly one that is not overloaded at all. Let the next tick re-decide.
+  defp maybe_steal(shard_lease, victim, state) do
     Logger.debug(
-      "ShardLease: Steal requested but the lease is already owned by this worker: " <>
-        "[shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}]"
+      "ShardLease: Skipping steal, lease is no longer owned by the chosen victim: " <>
+        "[shard_id: #{state.shard_id}, lease_owner: #{state.lease_owner}, " <>
+        "victim: #{victim}, current_owner: #{shard_lease.lease_owner}]"
     )
 
     state
@@ -298,6 +347,9 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
 
         renew_shard_lease(shard_lease, state)
 
+      shard_lease.lease_owner == state.lease_owner ->
+        reclaim_shard_lease(shard_lease, state)
+
       shard_lease.lease_owner != state.lease_owner and state.lease_holder ->
         Logger.debug(
           "ShardLease: Lease lost to another worker, stopping pipeline: [shard_id: #{state.shard_id}, " <>
diff --git a/lib/kinesis_client/stream/shard/load_balance.ex b/lib/kinesis_client/stream/shard/load_balance.ex
index d10388c..30972d5 100644
--- a/lib/kinesis_client/stream/shard/load_balance.ex
+++ b/lib/kinesis_client/stream/shard/load_balance.ex
@@ -11,18 +11,20 @@ defmodule KinesisClient.Stream.Shard.LoadBalance do
   @doc """
   Decides whether `lease_owner` should steal a lease.
 
-  Returns `{:steal_from, victim}` when `lease_owner` holds less than its share
-  of the leases (`ceil(total_leases / total_workers)`) and the most loaded
-  worker leads it by more than one lease, so a steal moves the distribution
-  closer to even instead of flipping the imbalance around. Returns `:balanced`
-  otherwise.
+  Returns `{:steal_from, victim, deficit}` when `lease_owner` holds less than
+  its share of the leases (`ceil(total_leases / total_workers)`) and the most
+  loaded worker leads it by more than one lease, so a steal moves the
+  distribution closer to even instead of flipping the imbalance around.
+  `deficit` is how many leases `lease_owner` is short of the target — the
+  caller must not steal more than that in one round, or the imbalance flips
+  to the other side and oscillates. Returns `:balanced` otherwise.
 
   `lease_owner` is counted as a worker even when it holds no leases and is
   therefore absent from the grouped counts — otherwise a fresh worker would
   look "balanced" and never claim its share of the shards.
   """
   @spec decide(list({String.t(), non_neg_integer()}), String.t()) ::
-          :balanced | {:steal_from, String.t()}
+          :balanced | {:steal_from, String.t(), pos_integer()}
   def decide(worker_counts, lease_owner) do
     worker_counts = include_current_worker(worker_counts, lease_owner)
     target = target_load(worker_counts)
@@ -40,7 +42,7 @@ defmodule KinesisClient.Stream.Shard.LoadBalance do
     |> Enum.max_by(fn {_owner, count} -> count end)
     |> case do
       {victim, victim_count} when my_count < target and victim_count - my_count > 1 ->
-        {:steal_from, victim}
+        {:steal_from, victim, target - my_count}
 
       _ ->
         :balanced
diff --git a/test/kinesis_client/stream/rebalancer_test.exs b/test/kinesis_client/stream/rebalancer_test.exs
index 092d181..ab4d797 100644
--- a/test/kinesis_client/stream/rebalancer_test.exs
+++ b/test/kinesis_client/stream/rebalancer_test.exs
@@ -51,7 +51,7 @@ defmodule KinesisClient.Stream.RebalancerTest do
     {:ok, pid} = start_supervised({Rebalancer, opts})
 
     assert_receive {:steal_requested, "shard-000002"}, 1_000
-    assert_receive {:lease_message, "shard-000002", :steal_lease}, 1_000
+    assert_receive {:lease_message, "shard-000002", {:steal_lease, ^victim}}, 1_000
     assert Process.alive?(pid)
     stop_supervised(Rebalancer)
   end
@@ -81,8 +81,41 @@ defmodule KinesisClient.Stream.RebalancerTest do
 
     {:ok, pid} = start_supervised({Rebalancer, opts})
 
-    assert_receive {:lease_message, _shard_id, :steal_lease}, 1_000
-    refute_receive {:lease_message, _shard_id, :steal_lease}, 200
+    assert_receive {:lease_message, _shard_id, {:steal_lease, _victim}}, 1_000
+    refute_receive {:lease_message, _shard_id, {:steal_lease, _victim}}, 200
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  test "clamps steals to the deficit even when max_leases_to_steal is higher" do
+    # Counts of {victim: 4, me: 0} mean a target of 2 and a deficit of 2:
+    # stealing max_leases_to_steal (3) would overshoot to {1, 3} and the
+    # imbalance would flip back and forth forever.
+    opts = build_rebalancer_opts(rebalance_interval: 600, max_leases_to_steal: 3)
+    victim = worker_ref()
+
+    victim_leases = [
+      build_shard_lease(shard_id: "shard-000001", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000002", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000003", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000004", lease_owner: victim)
+    ]
+
+    Enum.each(victim_leases, &register_lease_process(opts, &1.shard_id))
+
+    AppStateMock
+    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
+      [{victim, 4}]
+    end)
+    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
+      victim_leases
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:lease_message, _shard_id, {:steal_lease, _victim}}, 1_000
+    assert_receive {:lease_message, _shard_id, {:steal_lease, _victim}}, 1_000
+    refute_receive {:lease_message, _shard_id, {:steal_lease, _victim}}, 200
     assert Process.alive?(pid)
     stop_supervised(Rebalancer)
   end
@@ -109,7 +142,24 @@ defmodule KinesisClient.Stream.RebalancerTest do
 
     {:ok, pid} = start_supervised({Rebalancer, opts})
 
-    assert_receive {:lease_message, "shard-000002", :steal_lease}, 1_000
+    assert_receive {:lease_message, "shard-000002", {:steal_lease, ^victim}}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(Rebalancer)
+  end
+
+  test "survives a failing balancing query and keeps ticking" do
+    opts = build_rebalancer_opts()
+
+    stub(AppStateMock, :total_incomplete_lease_counts_by_worker, fn _app_name,
+                                                                    _stream_name,
+                                                                    _opts ->
+      raise "throttled scan"
+    end)
+
+    {:ok, pid} = start_supervised({Rebalancer, opts})
+
+    assert_receive {:rebalance_failed, _error}, 1_000
+    assert_receive {:rebalance_failed, _error}, 1_000
     assert Process.alive?(pid)
     stop_supervised(Rebalancer)
   end
diff --git a/test/kinesis_client/stream/shard/lease_v2_test.exs b/test/kinesis_client/stream/shard/lease_v2_test.exs
index d74e3a5..171b7b1 100644
--- a/test/kinesis_client/stream/shard/lease_v2_test.exs
+++ b/test/kinesis_client/stream/shard/lease_v2_test.exs
@@ -128,9 +128,54 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
     stop_supervised(LeaseV2)
   end
 
+  test "reclaims a lease the AppState says it owns but it is not holding" do
+    current_worker = worker_ref()
+
+    lease_opts =
+      build_lease_opts(
+        lease_owner: current_worker,
+        renew_interval: 200,
+        pipeline: NotifyingPipeline
+      )
+
+    # e.g. a renewal that "failed" after actually being applied (lost
+    # response + AWS retry fails the conditional check), or a restart of this
+    # process after taking the lease: the row names this worker as owner but
+    # lease_holder starts out (or was reset to) false. Without reclaiming,
+    # renewing requires lease_holder and taking rejects the current owner, so
+    # the shard would never be consumed again.
+    owned_shard_lease =
+      build_shard_lease(
+        lease_count: 1,
+        lease_owner: current_worker,
+        shard_id: lease_opts[:shard_id]
+      )
+
+    AppStateMock
+    |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      owned_shard_lease
+    end)
+    |> stub(:renew_lease, fn _app_name, _stream_name, shard_lease, _opts ->
+      {:ok, shard_lease.lease_count + 1}
+    end)
+
+    {:ok, pid} = start_supervised({LeaseV2, lease_opts})
+
+    assert_receive {:initialized, lease_state}, 1_000
+    assert lease_state.lease_holder == false
+
+    assert_receive {:lease_reclaimed, lease_state}, 1_000
+    assert lease_state.lease_holder == true
+    assert lease_state.lease_count == owned_shard_lease.lease_count + 1
+    assert_receive {:pipeline_started, _shard_id}, 1_000
+    assert Process.alive?(pid)
+    stop_supervised(LeaseV2)
+  end
+
   describe ":steal_lease message" do
     test "steals the lease from its current owner and starts the pipeline" do
       current_worker = worker_ref()
+      victim = worker_ref()
 
       lease_opts =
         build_lease_opts(lease_owner: current_worker, pipeline: NotifyingPipeline)
@@ -138,7 +183,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
       shard_lease =
         build_shard_lease(
           lease_count: 8,
-          lease_owner: worker_ref(),
+          lease_owner: victim,
           shard_id: lease_opts[:shard_id]
         )
 
@@ -160,7 +205,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
       assert_receive {:initialized, lease_state}, 1_000
       assert lease_state.lease_holder == false
 
-      send(pid, :steal_lease)
+      send(pid, {:steal_lease, victim})
 
       assert_receive {:lease_stolen, lease_state}, 1_000
       assert lease_state.lease_holder == true
@@ -212,7 +257,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
       assert_receive {:initialized, lease_state}, 1_000
       assert lease_state.lease_count == stale_shard_lease.lease_count
 
-      send(pid, :steal_lease)
+      send(pid, {:steal_lease, other_worker})
 
       assert_receive {:lease_stolen, lease_state}, 1_000
       assert lease_state.lease_holder == true
@@ -221,6 +266,32 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
       stop_supervised(LeaseV2)
     end
 
+    test "skips the steal when the lease changed owners since the decision" do
+      lease_opts = build_lease_opts()
+      chosen_victim = worker_ref()
+      new_owner = worker_ref()
+
+      # The rebalancer picked chosen_victim, but by the time the request
+      # arrives the lease belongs to someone else — stealing now would hit a
+      # worker the balancing decision never targeted.
+      shard_lease =
+        build_shard_lease(lease_owner: new_owner, shard_id: lease_opts[:shard_id])
+
+      stub(AppStateMock, :get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+        shard_lease
+      end)
+
+      {:ok, pid} = start_supervised({LeaseV2, lease_opts})
+
+      assert_receive {:initialized, %{lease_holder: false}}, 1_000
+
+      send(pid, {:steal_lease, chosen_victim})
+
+      refute_receive {:lease_stolen, _}, 200
+      assert Process.alive?(pid)
+      stop_supervised(LeaseV2)
+    end
+
     test "ignores the request when already the lease holder" do
       lease_opts = build_lease_opts(pipeline: KinesisClient.TestPipeline)
 
@@ -236,7 +307,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
 
       assert_receive {:initialized, %{lease_holder: true}}, 1_000
 
-      send(pid, :steal_lease)
+      send(pid, {:steal_lease, worker_ref()})
 
       refute_receive {:lease_stolen, _}, 200
       assert Process.alive?(pid)
@@ -261,7 +332,7 @@ defmodule KinesisClient.Stream.Shard.LeaseV2Test do
 
       assert_receive {:initialized, %{lease_holder: false}}, 1_000
 
-      send(pid, :steal_lease)
+      send(pid, {:steal_lease, shard_lease.lease_owner})
 
       refute_receive {:lease_stolen, _}, 200
       assert Process.alive?(pid)
diff --git a/test/kinesis_client/stream/shard/load_balance_test.exs b/test/kinesis_client/stream/shard/load_balance_test.exs
index a7b4217..975b635 100644
--- a/test/kinesis_client/stream/shard/load_balance_test.exs
+++ b/test/kinesis_client/stream/shard/load_balance_test.exs
@@ -18,10 +18,10 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       assert LoadBalance.decide(counts, "worker-1") == :balanced
     end
 
-    test "steals from the most loaded worker when under target" do
+    test "steals from the most loaded worker when under target, reporting the deficit" do
       counts = [{"worker-1", 1}, {"worker-2", 4}, {"worker-3", 4}]
 
-      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2"}
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 2}
     end
 
     test "a worker holding no leases counts itself and steals" do
@@ -29,7 +29,7 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       # must still count itself as a worker, see the imbalance, and steal.
       counts = [{"worker-2", 4}, {"worker-3", 4}]
 
-      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2"}
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 3}
     end
 
     test "does not steal when the lead is only one lease" do
@@ -45,13 +45,23 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       counts = [{"worker-1", 3}, {"worker-2", 6}, {"worker-3", 0}]
 
       assert LoadBalance.decide(counts, "worker-1") == :balanced
-      assert LoadBalance.decide(counts, "worker-3") == {:steal_from, "worker-2"}
+      assert LoadBalance.decide(counts, "worker-3") == {:steal_from, "worker-2", 3}
+    end
+
+    test "stealing exactly the deficit lands on a balanced spread" do
+      # {A: 4, B: 0}: taking more than the deficit of 2 would overshoot to
+      # {1, 3} and oscillate forever. Taking exactly the deficit settles it.
+      assert LoadBalance.decide([{"worker-a", 4}], "worker-b") == {:steal_from, "worker-a", 2}
+
+      settled = [{"worker-a", 2}, {"worker-b", 2}]
+      assert LoadBalance.decide(settled, "worker-a") == :balanced
+      assert LoadBalance.decide(settled, "worker-b") == :balanced
     end
 
     test "converges to an even spread as steals are applied" do
       # Simulate the {4, 4, 0} cluster rebalancing one steal at a time.
-      assert LoadBalance.decide([{"w2", 4}, {"w3", 4}], "w1") == {:steal_from, "w2"}
-      assert LoadBalance.decide([{"w1", 1}, {"w2", 3}, {"w3", 4}], "w1") == {:steal_from, "w3"}
+      assert LoadBalance.decide([{"w2", 4}, {"w3", 4}], "w1") == {:steal_from, "w2", 3}
+      assert LoadBalance.decide([{"w1", 1}, {"w2", 3}, {"w3", 4}], "w1") == {:steal_from, "w3", 2}
 
       assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w1") == :balanced
       assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w2") == :balanced

From c289ee9966cc8114fefca78323080223ae7a4da8 Mon Sep 17 00:00:00 2001
From: Jon Kusumo <tahir.kusumo@paradox.ai>
Date: Wed, 1 Jul 2026 23:00:14 -0400
Subject: [PATCH 3/4] Address review follow-ups: bound steal overlap,
 single-scan tick, dedup notify
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Stop the Producer as soon as an owner-guarded checkpoint fails against
  a lease owned by another worker. Record fetches were already
  owner-checked per poll, so this closes the remaining post-steal
  overlap: the victim now halts at its first checkpoint after the steal
  instead of idling until the lease process's next renewal.
- One AppState query per rebalance tick: the Rebalancer fetches the
  incomplete leases once and derives both the per-worker counts and the
  steal candidates from that list, instead of a counts query plus a
  second full scan on unbalanced ticks. Removes the now-unused
  total_incomplete_lease_counts_by_worker callback from the behaviour,
  facade, and all three adapters.
- Move the triplicated notify/2 test hook into KinesisClient.Util.
- Rename Producer.is_lease_owner? to lease_owner? per naming standards.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 lib/kinesis_client/stream/app_state.ex        |  6 --
 .../stream/app_state/adapter.ex               |  7 --
 lib/kinesis_client/stream/app_state/dynamo.ex |  8 --
 lib/kinesis_client/stream/app_state/ecto.ex   |  7 --
 .../stream/app_state/ecto/shard_leases.ex     | 11 ---
 lib/kinesis_client/stream/app_state/mimic.ex  |  8 --
 lib/kinesis_client/stream/coordinator.ex      |  3 +-
 lib/kinesis_client/stream/rebalancer.ex       | 29 +++---
 lib/kinesis_client/stream/shard/lease.ex      |  9 --
 lib/kinesis_client/stream/shard/lease_v2.ex   |  9 --
 lib/kinesis_client/stream/shard/producer.ex   | 91 ++++++++++++-------
 lib/kinesis_client/util.ex                    | 15 +++
 .../stream/app_state/dynamo_test.exs          | 17 ----
 .../kinesis_client/stream/rebalancer_test.exs | 71 ++++-----------
 .../stream/shard/producer_test.exs            | 38 ++++++++
 15 files changed, 144 insertions(+), 185 deletions(-)

diff --git a/lib/kinesis_client/stream/app_state.ex b/lib/kinesis_client/stream/app_state.ex
index 936a9f4..cdecd49 100644
--- a/lib/kinesis_client/stream/app_state.ex
+++ b/lib/kinesis_client/stream/app_state.ex
@@ -55,12 +55,6 @@ defmodule KinesisClient.Stream.AppState do
   def all_incomplete_leases(app_name, stream_name, opts \\ []),
     do: adapter(opts).all_incomplete_leases(app_name, stream_name, opts)
 
-  @doc """
-  Get total lease counts per owner for load balancing.
-  """
-  def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts \\ []),
-    do: adapter(opts).total_incomplete_lease_counts_by_worker(app_name, stream_name, opts)
-
   defp adapter(opts) do
     case Keyword.get(opts, :adapter) do
       :ecto -> KinesisClient.Stream.AppState.Ecto
diff --git a/lib/kinesis_client/stream/app_state/adapter.ex b/lib/kinesis_client/stream/app_state/adapter.ex
index d6a6df5..fa1ad0a 100644
--- a/lib/kinesis_client/stream/app_state/adapter.ex
+++ b/lib/kinesis_client/stream/app_state/adapter.ex
@@ -77,11 +77,4 @@ defmodule KinesisClient.Stream.AppState.Adapter do
               opts :: keyword
             ) ::
               list(ShardLease.t())
-
-  @callback total_incomplete_lease_counts_by_worker(
-              app_name :: String.t(),
-              stream_name :: String.t(),
-              opts :: keyword
-            ) ::
-              list({worker :: String.t(), count :: integer})
 end
diff --git a/lib/kinesis_client/stream/app_state/dynamo.ex b/lib/kinesis_client/stream/app_state/dynamo.ex
index 136bd17..845720a 100644
--- a/lib/kinesis_client/stream/app_state/dynamo.ex
+++ b/lib/kinesis_client/stream/app_state/dynamo.ex
@@ -200,14 +200,6 @@ defmodule KinesisClient.Stream.AppState.Dynamo do
     )
   end
 
-  @impl true
-  def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts) do
-    app_name
-    |> all_incomplete_leases(stream_name, opts)
-    |> Enum.frequencies_by(& &1.lease_owner)
-    |> Map.to_list()
-  end
-
   # Lease tables hold one row per shard, so a filtered Scan is cheap enough for
   # the load balancing queries (this mirrors how the Java KCL reads its lease
   # table). ExAws.stream!/1 follows LastEvaluatedKey pagination for us.
diff --git a/lib/kinesis_client/stream/app_state/ecto.ex b/lib/kinesis_client/stream/app_state/ecto.ex
index 9d830fd..5e8c2f6 100644
--- a/lib/kinesis_client/stream/app_state/ecto.ex
+++ b/lib/kinesis_client/stream/app_state/ecto.ex
@@ -190,13 +190,6 @@ defmodule KinesisClient.Stream.AppState.Ecto do
     |> ShardLeases.get_shard_leases(repo)
   end
 
-  @impl true
-  def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts) do
-    repo = Keyword.get(opts, :repo)
-
-    ShardLeases.incomplete_group_by_owner(app_name, stream_name, repo)
-  end
-
   def create_lease(attrs, opts) when is_map(attrs) do
     repo = Keyword.get(opts, :repo)
 
diff --git a/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex b/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
index bc8be0d..d4f0021 100644
--- a/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
+++ b/lib/kinesis_client/stream/app_state/ecto/shard_leases.ex
@@ -52,17 +52,6 @@ defmodule KinesisClient.Stream.AppState.Ecto.ShardLeases do
     end
   end
 
-  @spec incomplete_group_by_owner(String.t(), String.t(), Ecto.Repo.t()) ::
-          [{lease_owner :: String.t(), count :: integer}]
-  def incomplete_group_by_owner(app_name, stream_name, repo) do
-    from(sl in ShardLeaseEcto,
-      where: sl.app_name == ^app_name and sl.stream_name == ^stream_name and not sl.completed,
-      group_by: sl.lease_owner,
-      select: {sl.lease_owner, count(sl.shard_id)}
-    )
-    |> repo.all()
-  end
-
   # Pins the row to the freshly read lease_count AND lease_owner so a
   # concurrent renew/take between our read and this update makes the
   # update_all match zero rows instead of clobbering the other worker's
diff --git a/lib/kinesis_client/stream/app_state/mimic.ex b/lib/kinesis_client/stream/app_state/mimic.ex
index fa40448..acb9cde 100644
--- a/lib/kinesis_client/stream/app_state/mimic.ex
+++ b/lib/kinesis_client/stream/app_state/mimic.ex
@@ -97,14 +97,6 @@ defmodule KinesisClient.Stream.AppState.Mimic do
     from.all_incomplete_leases(app_name, stream_name, opts)
   end
 
-  @impl true
-  def total_incomplete_lease_counts_by_worker(app_name, stream_name, opts) do
-    {from, to} = modules(opts)
-
-    to.total_incomplete_lease_counts_by_worker(app_name, stream_name, opts)
-    from.total_incomplete_lease_counts_by_worker(app_name, stream_name, opts)
-  end
-
   defp modules(opts) do
     migration = Keyword.get(opts, :migration)
 
diff --git a/lib/kinesis_client/stream/coordinator.ex b/lib/kinesis_client/stream/coordinator.ex
index 43f14cf..7cf5b37 100644
--- a/lib/kinesis_client/stream/coordinator.ex
+++ b/lib/kinesis_client/stream/coordinator.ex
@@ -7,7 +7,8 @@ defmodule KinesisClient.Stream.Coordinator do
   use GenServer
   use Retry.Annotation
 
-  import KinesisClient.Util
+  # The local notify/2 works off :notify_pid rather than Util's :notify key.
+  import KinesisClient.Util, except: [notify: 2]
 
   alias KinesisClient.Kinesis
   alias KinesisClient.Stream.AppState
diff --git a/lib/kinesis_client/stream/rebalancer.ex b/lib/kinesis_client/stream/rebalancer.ex
index 0fae485..54c3951 100644
--- a/lib/kinesis_client/stream/rebalancer.ex
+++ b/lib/kinesis_client/stream/rebalancer.ex
@@ -84,23 +84,29 @@ defmodule KinesisClient.Stream.Rebalancer do
       {:noreply, state}
   end
 
+  # One query per tick: the incomplete leases give both the per-worker counts
+  # for the balance decision and the steal candidates, so an unbalanced tick
+  # costs no more than a balanced one.
   defp run_rebalance(state) do
-    state.app_name
-    |> AppState.total_incomplete_lease_counts_by_worker(state.stream_name, state.app_state_opts)
+    leases =
+      AppState.all_incomplete_leases(state.app_name, state.stream_name, state.app_state_opts)
+
+    leases
+    |> Enum.frequencies_by(& &1.lease_owner)
+    |> Map.to_list()
     |> LoadBalance.decide(state.lease_owner)
     |> case do
       :balanced ->
         notify({:all_balanced, state}, state)
 
       {:steal_from, victim, deficit} ->
-        steal_from(victim, deficit, state)
+        steal_from(victim, deficit, leases, state)
     end
   end
 
-  defp steal_from(victim, deficit, state) do
-    state.app_name
-    |> AppState.get_leases_by_worker(state.stream_name, victim, state.app_state_opts)
-    |> Enum.reject(& &1.completed)
+  defp steal_from(victim, deficit, leases, state) do
+    leases
+    |> Enum.filter(&(&1.lease_owner == victim))
     |> Enum.shuffle()
     |> Enum.map(&local_lease_process(&1, state))
     |> Enum.reject(&is_nil/1)
@@ -136,13 +142,4 @@ defmodule KinesisClient.Stream.Rebalancer do
   defp jitter(interval) do
     interval + :rand.uniform(max(div(interval, 2), 1)) - div(interval, 4)
   end
-
-  defp notify(_msg, %{notify: nil}) do
-    :ok
-  end
-
-  defp notify(msg, %{notify: notify}) do
-    send(notify, msg)
-    :ok
-  end
 end
diff --git a/lib/kinesis_client/stream/shard/lease.ex b/lib/kinesis_client/stream/shard/lease.ex
index fa7ae98..4edef16 100644
--- a/lib/kinesis_client/stream/shard/lease.ex
+++ b/lib/kinesis_client/stream/shard/lease.ex
@@ -260,15 +260,6 @@ defmodule KinesisClient.Stream.Shard.Lease do
     end
   end
 
-  defp notify(_msg, %{notify: nil}) do
-    :ok
-  end
-
-  defp notify(msg, %{notify: notify}) do
-    send(notify, msg)
-    :ok
-  end
-
   defp current_time do
     System.monotonic_time(:millisecond)
   end
diff --git a/lib/kinesis_client/stream/shard/lease_v2.ex b/lib/kinesis_client/stream/shard/lease_v2.ex
index f1374ed..8519ff0 100644
--- a/lib/kinesis_client/stream/shard/lease_v2.ex
+++ b/lib/kinesis_client/stream/shard/lease_v2.ex
@@ -400,15 +400,6 @@ defmodule KinesisClient.Stream.Shard.LeaseV2 do
     }
   end
 
-  defp notify(_msg, %{notify: nil}) do
-    :ok
-  end
-
-  defp notify(msg, %{notify: notify}) do
-    send(notify, msg)
-    :ok
-  end
-
   defp current_time() do
     System.monotonic_time(:millisecond)
   end
diff --git a/lib/kinesis_client/stream/shard/producer.ex b/lib/kinesis_client/stream/shard/producer.ex
index 836ebb0..fb25fa9 100644
--- a/lib/kinesis_client/stream/shard/producer.ex
+++ b/lib/kinesis_client/stream/shard/producer.ex
@@ -128,7 +128,7 @@ defmodule KinesisClient.Stream.Shard.Producer do
   def handle_info(:get_records, state) do
     notify(:poll_timer_executed, state)
 
-    if is_lease_owner?(state) do
+    if lease_owner?(state) do
       Logger.debug(
         "Try to fulfill pending demand #{state.demand}: [stream_name: #{state.stream_name}, shard_id: #{state.shard_id}]"
       )
@@ -162,38 +162,43 @@ defmodule KinesisClient.Stream.Shard.Producer do
   def handle_info({:ack, _ref, successful_msgs, []}, state) do
     %{metadata: %{"SequenceNumber" => checkpoint}} = successful_msgs |> Enum.reverse() |> hd()
 
-    state.app_name
-    |> AppState.update_checkpoint(
-      state.stream_name,
-      state.shard_id,
-      state.lease_owner,
-      checkpoint,
-      state.app_state_opts
-    )
-    |> case do
-      :ok ->
-        notify({:acked, %{checkpoint: checkpoint, success: successful_msgs, failed: []}}, state)
+    state =
+      state.app_name
+      |> AppState.update_checkpoint(
+        state.stream_name,
+        state.shard_id,
+        state.lease_owner,
+        checkpoint,
+        state.app_state_opts
+      )
+      |> case do
+        :ok ->
+          notify({:acked, %{checkpoint: checkpoint, success: successful_msgs, failed: []}}, state)
 
-        Logger.debug(
-          "Acknowledged #{length(successful_msgs)} messages: [app_name: #{state.app_name} " <>
-            "shard_id: #{state.shard_id} data: #{inspect(successful_msgs)}"
-        )
+          Logger.debug(
+            "Acknowledged #{length(successful_msgs)} messages: [app_name: #{state.app_name} " <>
+              "shard_id: #{state.shard_id} data: #{inspect(successful_msgs)}"
+          )
 
-      {:error, error} ->
-        shard_lease =
-          AppState.get_lease(
-            state.app_name,
-            state.stream_name,
-            state.shard_id,
-            state.app_state_opts
+          state
+
+        {:error, error} ->
+          shard_lease =
+            AppState.get_lease(
+              state.app_name,
+              state.stream_name,
+              state.shard_id,
+              state.app_state_opts
+            )
+
+          Logger.error(
+            "Failed to update checkpoint after acknowledging #{length(successful_msgs)} messages: [app_name: #{state.app_name} " <>
+              "shard_id: #{state.shard_id} lease_owner: #{state.lease_owner} current_shard_owner: #{inspect(shard_lease)} " <>
+              "checkpoint: #{checkpoint} error: #{inspect(error)} data: #{inspect(successful_msgs)}"
           )
 
-        Logger.error(
-          "Failed to update checkpoint after acknowledging #{length(successful_msgs)} messages: [app_name: #{state.app_name} " <>
-            "shard_id: #{state.shard_id} lease_owner: #{state.lease_owner} current_shard_owner: #{shard_lease.lease_owner} " <>
-            "checkpoint: #{checkpoint} error: #{inspect(error)} data: #{inspect(successful_msgs)}"
-        )
-    end
+          stop_if_lease_lost(shard_lease, state)
+      end
 
     state =
       state.status
@@ -207,7 +212,7 @@ defmodule KinesisClient.Stream.Shard.Producer do
 
   @impl GenStage
   def handle_info({:ack, _ref, [], failed_msgs}, %{status: :stopped} = state) do
-    if is_lease_owner?(state) do
+    if lease_owner?(state) do
       Logger.debug(
         "Shard #{state.shard_id} - Retrying #{length(failed_msgs)} failed messages - #{inspect(failed_msgs)}"
       )
@@ -241,7 +246,7 @@ defmodule KinesisClient.Stream.Shard.Producer do
 
   @impl GenStage
   def handle_info({:ack, _ref, successful_msgs, failed_msgs}, %{status: :stopped} = state) do
-    if is_lease_owner?(state) do
+    if lease_owner?(state) do
       Logger.debug(
         "Shard #{state.shard_id} - Acknowledged #{length(successful_msgs)} messages, " <>
           "Retrying #{length(failed_msgs)} failed messages - #{inspect(failed_msgs)}"
@@ -391,7 +396,7 @@ defmodule KinesisClient.Stream.Shard.Producer do
 
   @retry with: 500 |> exponential_backoff() |> Stream.take(5)
   defp get_records_with_retry(state, kinesis_opts) do
-    if is_lease_owner?(state) do
+    if lease_owner?(state) do
       Kinesis.get_records(state.shard_iterator, kinesis_opts)
       |> tap(
         &Logger.debug("Shard #{state.shard_id} Kinesis get_records_with_retry: #{inspect(&1)}")
@@ -546,7 +551,27 @@ defmodule KinesisClient.Stream.Shard.Producer do
     Process.send_after(self(), :get_records, interval)
   end
 
-  defp is_lease_owner?(state) do
+  # Checkpoint updates are owner-guarded, so a failure with the lease naming
+  # another worker means the lease was stolen or expired since the last
+  # fetch. Stop producing right away instead of waiting for the lease process
+  # to notice at its next renewal — this bounds double consumption after a
+  # steal to the messages already in flight.
+  defp stop_if_lease_lost(%{lease_owner: owner}, %{lease_owner: owner} = state), do: state
+
+  defp stop_if_lease_lost(%{lease_owner: _other_owner}, state) do
+    Logger.warning(
+      "Shard #{state.shard_id} lease is no longer owned by this worker, stopping producer: " <>
+        "[app_name: #{state.app_name}, lease_owner: #{state.lease_owner}]"
+    )
+
+    notify({:lease_lost, state.shard_id}, state)
+
+    %{state | status: :stopped}
+  end
+
+  defp stop_if_lease_lost(_not_found_or_error, state), do: state
+
+  defp lease_owner?(state) do
     state.app_name
     |> AppState.get_lease(state.stream_name, state.shard_id, state.app_state_opts)
     |> then(fn lease -> lease.lease_owner == state.lease_owner end)
diff --git a/lib/kinesis_client/util.ex b/lib/kinesis_client/util.ex
index 2ad147b..1dcf0b9 100644
--- a/lib/kinesis_client/util.ex
+++ b/lib/kinesis_client/util.ex
@@ -11,4 +11,19 @@ defmodule KinesisClient.Util do
   def register_name(module, app_name, stream_name, addtnl \\ []) do
     Module.concat([module, app_name, stream_name] ++ addtnl)
   end
+
+  @doc """
+  Sends `message` to the pid in the state's `:notify` field, if one is set.
+
+  Used by the lease and rebalancer processes to expose lifecycle events to
+  tests.
+  """
+  def notify(_message, %{notify: nil}) do
+    :ok
+  end
+
+  def notify(message, %{notify: pid}) do
+    send(pid, message)
+    :ok
+  end
 end
diff --git a/test/kinesis_client/stream/app_state/dynamo_test.exs b/test/kinesis_client/stream/app_state/dynamo_test.exs
index 8f0766a..2037f7d 100644
--- a/test/kinesis_client/stream/app_state/dynamo_test.exs
+++ b/test/kinesis_client/stream/app_state/dynamo_test.exs
@@ -173,23 +173,6 @@ defmodule KinesisClient.Stream.AppState.DynamoTest do
       assert [%ShardLease{shard_id: "shard-000001", completed: false}] =
                AppState.all_incomplete_leases(app_name, "", [])
     end
-
-    test "total_incomplete_lease_counts_by_worker/3 groups counts by owner", %{lb_app: app_name} do
-      worker_1 = worker_ref()
-      worker_2 = worker_ref()
-      assert :ok == AppState.create_lease(app_name, "", "shard-000001", worker_1, [])
-      assert :ok == AppState.create_lease(app_name, "", "shard-000002", worker_1, [])
-      assert :ok == AppState.create_lease(app_name, "", "shard-000003", worker_2, [])
-      assert :ok == AppState.create_lease(app_name, "", "shard-000004", worker_2, [])
-      assert :ok == AppState.close_shard(app_name, "", "shard-000004", worker_2, [])
-
-      counts =
-        app_name
-        |> AppState.total_incomplete_lease_counts_by_worker("", [])
-        |> Enum.sort()
-
-      assert counts == Enum.sort([{worker_1, 2}, {worker_2, 1}])
-    end
   end
 
   defp confirm_table_created(app_name, attempts \\ 1) do
diff --git a/test/kinesis_client/stream/rebalancer_test.exs b/test/kinesis_client/stream/rebalancer_test.exs
index ab4d797..1c81c80 100644
--- a/test/kinesis_client/stream/rebalancer_test.exs
+++ b/test/kinesis_client/stream/rebalancer_test.exs
@@ -10,11 +10,15 @@ defmodule KinesisClient.Stream.RebalancerTest do
   test "notifies :all_balanced and requests no steals when the load is balanced" do
     opts = build_rebalancer_opts()
     lease_owner = opts[:lease_owner]
+    other_worker = worker_ref()
 
-    stub(AppStateMock, :total_incomplete_lease_counts_by_worker, fn _app_name,
-                                                                    _stream_name,
-                                                                    _opts ->
-      [{lease_owner, 2}, {worker_ref(), 2}]
+    stub(AppStateMock, :all_incomplete_leases, fn _app_name, _stream_name, _opts ->
+      [
+        build_shard_lease(shard_id: "shard-000001", lease_owner: lease_owner),
+        build_shard_lease(shard_id: "shard-000002", lease_owner: lease_owner),
+        build_shard_lease(shard_id: "shard-000003", lease_owner: other_worker),
+        build_shard_lease(shard_id: "shard-000004", lease_owner: other_worker)
+      ]
     end)
 
     {:ok, pid} = start_supervised({Rebalancer, opts})
@@ -39,12 +43,7 @@ defmodule KinesisClient.Stream.RebalancerTest do
     # candidate order the steal request can only go there.
     register_lease_process(opts, "shard-000002")
 
-    AppStateMock
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{victim, 4}]
-    end)
-    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, in_lease_owner, _opts ->
-      assert in_lease_owner == victim
+    stub(AppStateMock, :all_incomplete_leases, fn _app_name, _stream_name, _opts ->
       victim_leases
     end)
 
@@ -63,19 +62,18 @@ defmodule KinesisClient.Stream.RebalancerTest do
     opts = build_rebalancer_opts(rebalance_interval: 600)
     victim = worker_ref()
 
+    # 4 victim leases and none of ours: a deficit of 2, but the default
+    # max_leases_to_steal of 1 caps the tick at a single steal request.
     victim_leases = [
       build_shard_lease(shard_id: "shard-000001", lease_owner: victim),
-      build_shard_lease(shard_id: "shard-000002", lease_owner: victim)
+      build_shard_lease(shard_id: "shard-000002", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000003", lease_owner: victim),
+      build_shard_lease(shard_id: "shard-000004", lease_owner: victim)
     ]
 
-    register_lease_process(opts, "shard-000001")
-    register_lease_process(opts, "shard-000002")
+    Enum.each(victim_leases, &register_lease_process(opts, &1.shard_id))
 
-    AppStateMock
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{victim, 4}]
-    end)
-    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
+    stub(AppStateMock, :all_incomplete_leases, fn _app_name, _stream_name, _opts ->
       victim_leases
     end)
 
@@ -103,11 +101,7 @@ defmodule KinesisClient.Stream.RebalancerTest do
 
     Enum.each(victim_leases, &register_lease_process(opts, &1.shard_id))
 
-    AppStateMock
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{victim, 4}]
-    end)
-    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
+    stub(AppStateMock, :all_incomplete_leases, fn _app_name, _stream_name, _opts ->
       victim_leases
     end)
 
@@ -120,39 +114,10 @@ defmodule KinesisClient.Stream.RebalancerTest do
     stop_supervised(Rebalancer)
   end
 
-  test "skips completed shards when picking a steal candidate" do
-    opts = build_rebalancer_opts()
-    victim = worker_ref()
-
-    victim_leases = [
-      build_shard_lease(shard_id: "shard-000001", lease_owner: victim, completed: true),
-      build_shard_lease(shard_id: "shard-000002", lease_owner: victim)
-    ]
-
-    register_lease_process(opts, "shard-000001")
-    register_lease_process(opts, "shard-000002")
-
-    AppStateMock
-    |> stub(:total_incomplete_lease_counts_by_worker, fn _app_name, _stream_name, _opts ->
-      [{victim, 4}]
-    end)
-    |> stub(:get_leases_by_worker, fn _app_name, _stream_name, _lease_owner, _opts ->
-      victim_leases
-    end)
-
-    {:ok, pid} = start_supervised({Rebalancer, opts})
-
-    assert_receive {:lease_message, "shard-000002", {:steal_lease, ^victim}}, 1_000
-    assert Process.alive?(pid)
-    stop_supervised(Rebalancer)
-  end
-
   test "survives a failing balancing query and keeps ticking" do
     opts = build_rebalancer_opts()
 
-    stub(AppStateMock, :total_incomplete_lease_counts_by_worker, fn _app_name,
-                                                                    _stream_name,
-                                                                    _opts ->
+    stub(AppStateMock, :all_incomplete_leases, fn _app_name, _stream_name, _opts ->
       raise "throttled scan"
     end)
 
diff --git a/test/kinesis_client/stream/shard/producer_test.exs b/test/kinesis_client/stream/shard/producer_test.exs
index 929ed3f..4286f57 100644
--- a/test/kinesis_client/stream/shard/producer_test.exs
+++ b/test/kinesis_client/stream/shard/producer_test.exs
@@ -111,6 +111,44 @@ defmodule KinesisClient.Stream.Shard.ProducerTest do
     assert_receive {:acked, %{success: _successful, checkpoint: "12345", failed: []}}, 10_000
   end
 
+  test "stops producing when the checkpoint fails because the lease was lost" do
+    opts = producer_opts(status: :started)
+    {:ok, producer} = start_supervised({Producer, opts})
+    {:ok, consumer} = start_supervised({KinesisClient.TestConsumer, self()})
+
+    KinesisMock
+    |> expect(:get_shard_iterator, fn _, _, _, _ ->
+      {:ok, %{"ShardIterator" => "somesharditerator"}}
+    end)
+    |> expect(:get_records, fn _, _ ->
+      records = [%{"Data" => "foo", "SequenceNumber" => "12345"}]
+
+      {:ok, %{"NextShardIterator" => "foo", "MillisBehindLatest" => 1_000, "Records" => records}}
+    end)
+
+    # We own the lease for the fetch, then another worker steals it before
+    # the checkpoint: the owner-guarded checkpoint fails and get_lease names
+    # the thief.
+    AppStateMock
+    |> expect(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      %{lease_owner: opts[:lease_owner]}
+    end)
+    |> stub(:get_lease, fn _in_app_name, _in_stream_name, _in_shard_id, _ ->
+      %{lease_owner: worker_ref()}
+    end)
+    |> stub(:update_checkpoint, fn _app_name, _stream_name, _shard_id, _owner, _checkpoint, _ ->
+      {:error, :lease_owner_match}
+    end)
+
+    GenStage.sync_subscribe(consumer, to: producer, max_demand: 10, min_demand: 0)
+    assert_receive {:consumer_events, events}, 1_000
+
+    send(producer, {:ack, make_ref(), events, []})
+
+    assert_receive {:lease_lost, _shard_id}, 1_000
+    assert :sys.get_state(producer).state.status == :stopped
+  end
+
   test "close the shard when getting ResourceNotFoundException error" do
     opts = producer_opts(status: :started)
     {:ok, producer} = start_supervised({Producer, opts})

From 4e1587a5dfd953ab3a778c29daa1bf31bf2c81ee Mon Sep 17 00:00:00 2001
From: Jon Kusumo <tahir.kusumo@paradox.ai>
Date: Thu, 2 Jul 2026 00:16:58 -0400
Subject: [PATCH 4/4] Clamp steals by half the victim gap to prevent pairwise
 live-lock
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The deficit alone bounds what the thief needs, not what the victim can
give up: with max_leases_to_steal >= 2, {2, 4, 4} stole the deficit of 2
and flipped the pair to {4, 2, 4} every round, forever. The steal count
is now min(target - my_count, div(victim_count - my_count, 2)) — half
the gap is the largest move that cannot invert the pairwise imbalance.

Verified with a simulation against the real decide/2: 2500 randomized
states (2-7 workers, max_leases_to_steal 1-8, shuffled per-round worker
orderings) plus the reported {2,4,4} and {5,5,0} counter-examples all
converge to balanced.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 .../stream/shard/load_balance.ex              | 22 +++++---
 .../stream/shard/load_balance_test.exs        | 50 ++++++++++++++-----
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/lib/kinesis_client/stream/shard/load_balance.ex b/lib/kinesis_client/stream/shard/load_balance.ex
index 30972d5..95b5965 100644
--- a/lib/kinesis_client/stream/shard/load_balance.ex
+++ b/lib/kinesis_client/stream/shard/load_balance.ex
@@ -11,13 +11,12 @@ defmodule KinesisClient.Stream.Shard.LoadBalance do
   @doc """
   Decides whether `lease_owner` should steal a lease.
 
-  Returns `{:steal_from, victim, deficit}` when `lease_owner` holds less than
-  its share of the leases (`ceil(total_leases / total_workers)`) and the most
-  loaded worker leads it by more than one lease, so a steal moves the
+  Returns `{:steal_from, victim, steal_count}` when `lease_owner` holds less
+  than its share of the leases (`ceil(total_leases / total_workers)`) and the
+  most loaded worker leads it by more than one lease, so a steal moves the
   distribution closer to even instead of flipping the imbalance around.
-  `deficit` is how many leases `lease_owner` is short of the target — the
-  caller must not steal more than that in one round, or the imbalance flips
-  to the other side and oscillates. Returns `:balanced` otherwise.
+  `steal_count` is the most the caller may take in one round and still
+  converge. Returns `:balanced` otherwise.
 
   `lease_owner` is counted as a worker even when it holds no leases and is
   therefore absent from the grouped counts — otherwise a fresh worker would
@@ -42,13 +41,22 @@ defmodule KinesisClient.Stream.Shard.LoadBalance do
     |> Enum.max_by(fn {_owner, count} -> count end)
     |> case do
       {victim, victim_count} when my_count < target and victim_count - my_count > 1 ->
-        {:steal_from, victim, target - my_count}
+        {:steal_from, victim, steal_count(my_count, victim_count, target)}
 
       _ ->
         :balanced
     end
   end
 
+  # A steal is bounded by what this worker lacks (target - my_count) AND by
+  # half the gap to the victim: taking more than half flips the pairwise
+  # imbalance instead of settling it, and the pair trades the same leases
+  # back and forth forever (e.g. {2, 4, 4}: the deficit of 2 flips 2/4 to
+  # 4/2 every round; half the gap moves it to 3/3 and it converges).
+  defp steal_count(my_count, victim_count, target) do
+    min(target - my_count, div(victim_count - my_count, 2))
+  end
+
   defp include_current_worker(worker_counts, lease_owner) do
     worker_counts
     |> List.keymember?(lease_owner, 0)
diff --git a/test/kinesis_client/stream/shard/load_balance_test.exs b/test/kinesis_client/stream/shard/load_balance_test.exs
index 975b635..bcaa0ce 100644
--- a/test/kinesis_client/stream/shard/load_balance_test.exs
+++ b/test/kinesis_client/stream/shard/load_balance_test.exs
@@ -18,10 +18,12 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       assert LoadBalance.decide(counts, "worker-1") == :balanced
     end
 
-    test "steals from the most loaded worker when under target, reporting the deficit" do
+    test "steals from the most loaded worker when under target" do
+      # The steal count is half the gap to the victim (1), not the full
+      # deficit of 2 — taking 2 would leave the victim below worker-1.
       counts = [{"worker-1", 1}, {"worker-2", 4}, {"worker-3", 4}]
 
-      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 2}
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 1}
     end
 
     test "a worker holding no leases counts itself and steals" do
@@ -29,7 +31,7 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       # must still count itself as a worker, see the imbalance, and steal.
       counts = [{"worker-2", 4}, {"worker-3", 4}]
 
-      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 3}
+      assert LoadBalance.decide(counts, "worker-1") == {:steal_from, "worker-2", 2}
     end
 
     test "does not steal when the lead is only one lease" do
@@ -48,9 +50,9 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       assert LoadBalance.decide(counts, "worker-3") == {:steal_from, "worker-2", 3}
     end
 
-    test "stealing exactly the deficit lands on a balanced spread" do
-      # {A: 4, B: 0}: taking more than the deficit of 2 would overshoot to
-      # {1, 3} and oscillate forever. Taking exactly the deficit settles it.
+    test "stealing the full steal_count lands on a balanced spread" do
+      # {A: 4, B: 0}: taking more than 2 would overshoot to {1, 3} and
+      # oscillate forever. Taking exactly the steal count settles it.
       assert LoadBalance.decide([{"worker-a", 4}], "worker-b") == {:steal_from, "worker-a", 2}
 
       settled = [{"worker-a", 2}, {"worker-b", 2}]
@@ -58,14 +60,38 @@ defmodule KinesisClient.Stream.Shard.LoadBalanceTest do
       assert LoadBalance.decide(settled, "worker-b") == :balanced
     end
 
+    test "steal_count never flips the pairwise gap ({2, 4, 4} live-lock regression)" do
+      # The deficit here is 2 (target 4), but stealing 2 from worker-b flips
+      # {2, 4} to {4, 2} and the pair trades the same leases forever. Half
+      # the gap (1) moves it to {3, 3, 4}, which is balanced.
+      counts = [{"worker-a", 2}, {"worker-b", 4}, {"worker-c", 4}]
+
+      assert LoadBalance.decide(counts, "worker-a") == {:steal_from, "worker-b", 1}
+
+      settled = [{"worker-a", 3}, {"worker-b", 3}, {"worker-c", 4}]
+      assert LoadBalance.decide(settled, "worker-a") == :balanced
+      assert LoadBalance.decide(settled, "worker-b") == :balanced
+      assert LoadBalance.decide(settled, "worker-c") == :balanced
+    end
+
+    test "converges for {5, 5, 0} without flipping" do
+      assert LoadBalance.decide([{"w1", 5}, {"w2", 5}], "w3") == {:steal_from, "w1", 2}
+      assert LoadBalance.decide([{"w1", 3}, {"w2", 5}, {"w3", 2}], "w3") == {:steal_from, "w2", 1}
+
+      settled = [{"w1", 3}, {"w2", 4}, {"w3", 3}]
+      assert LoadBalance.decide(settled, "w1") == :balanced
+      assert LoadBalance.decide(settled, "w2") == :balanced
+      assert LoadBalance.decide(settled, "w3") == :balanced
+    end
+
     test "converges to an even spread as steals are applied" do
-      # Simulate the {4, 4, 0} cluster rebalancing one steal at a time.
-      assert LoadBalance.decide([{"w2", 4}, {"w3", 4}], "w1") == {:steal_from, "w2", 3}
-      assert LoadBalance.decide([{"w1", 1}, {"w2", 3}, {"w3", 4}], "w1") == {:steal_from, "w3", 2}
+      # Simulate the {4, 4, 0} cluster rebalancing round by round.
+      assert LoadBalance.decide([{"w2", 4}, {"w3", 4}], "w1") == {:steal_from, "w2", 2}
+      assert LoadBalance.decide([{"w1", 2}, {"w2", 2}, {"w3", 4}], "w1") == {:steal_from, "w3", 1}
 
-      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w1") == :balanced
-      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w2") == :balanced
-      assert LoadBalance.decide([{"w1", 2}, {"w2", 3}, {"w3", 3}], "w3") == :balanced
+      assert LoadBalance.decide([{"w1", 3}, {"w2", 2}, {"w3", 3}], "w1") == :balanced
+      assert LoadBalance.decide([{"w1", 3}, {"w2", 2}, {"w3", 3}], "w2") == :balanced
+      assert LoadBalance.decide([{"w1", 3}, {"w2", 2}, {"w3", 3}], "w3") == :balanced
     end
   end
 end