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

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 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
 - [ ] 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.
 

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__))
@@ -59,12 +65,20 @@ 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(: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 +91,8 @@ defmodule KinesisClient.Stream do
 
     children = [
       shard_supervisor_spec,
-      {Coordinator, coordinator_args}
+      {Coordinator, coordinator_args},
+      {Rebalancer, rebalancer_args}
     ]
 
     Logger.info(
@@ -92,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])}

lib/kinesis_client/stream/app_state.ex

@@ -55,18 +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)
-
-  @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

lib/kinesis_client/stream/app_state/adapter.ex

@@ -77,18 +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})
-
-  @callback lease_owner_with_most_leases(
-              app_name :: String.t(),
-              stream_name :: String.t(),
-              opts :: keyword
-            ) ::
-              list(ShardLease.t())
 end

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,21 @@ 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."
-    )
-
-    []
-  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

lib/kinesis_client/stream/app_state/ecto.ex

@@ -190,28 +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)
-
-    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)
 

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
 
@@ -51,43 +52,19 @@ 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
-
-  @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(
       [sl],
       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
 

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
@@ -94,22 +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
-
-  @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)
 

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