Dynamic queue sharding

lib/consumer.ex

@@ -31,7 +31,7 @@ defmodule Mississippi.Consumer do
 
     channels_per_connection = amqp_consumer_options[:channels]
 
-    queue_count = queue_config[:range_end] - queue_config[:range_start] + 1
+    queue_count = queue_config[:total_count]
 
     # Invariant: we use one channel for one queue.
     connection_number = Kernel.ceil(queue_count / channels_per_connection)

lib/consumer/amqp_data_consumer.ex

@@ -123,6 +123,12 @@ defmodule Mississippi.Consumer.AMQPDataConsumer do
     {:stop, :normal, state}
   end
 
+  def handle_info({:EXIT, _from, {:shutdown, :process_redistribution}}, state) do
+    _ = Logger.info("AMQPDataConsumer shutting down due to process redistribution")
+    # Exit with :normal so we're not restarted
+    {:stop, :normal, state}
+  end
+
   def handle_info({:EXIT, _from, reason}, state) do
     {:stop, reason, state}
   end

lib/consumer/amqp_data_consumer/starter.ex

@@ -0,0 +1,72 @@
+# Copyright 2025 SECO Mind Srl
+# SPDX-License-Identifier: Apache-2.0
+
+defmodule Mississippi.Consumer.AMQPDataConsumer.Starter do
+  @moduledoc false
+  use Task, restart: :transient
+
+  alias Horde.DynamicSupervisor
+  alias Mississippi.Consumer.AMQPDataConsumer
+
+  require Logger
+
+  @restart_backoff :timer.seconds(2)
+
+  def start_link(queues_config) do
+    Task.start_link(__MODULE__, :start_consumers, [queues_config])
+  end
+
+  def start_consumers(queues_config) do
+    start_consumers(queues_config, 10)
+  end
+
+  defp start_consumers(_, 0) do
+    _ = Logger.warning("Cannot start AMQPDataConsumers")
+    {:error, :cannot_start_consumers}
+  end
+
+  defp start_consumers(queues_config, retry) do
+    start_amqp_consumers(queues_config)
+
+    queue_total = queues_config[:total_count]
+
+    child_count =
+      AMQPDataConsumer.Supervisor |> DynamicSupervisor.which_children() |> Enum.count()
+
+    case child_count do
+      ^queue_total ->
+        :ok
+
+      _ ->
+        # TODO: do we want something more refined, e.g. exponential backoff?
+        backoff_delta = :rand.uniform(@restart_backoff)
+        Process.sleep(@restart_backoff + backoff_delta)
+        start_consumers(queues_config, retry - 1)
+    end
+  end
+
+  def start_amqp_consumers(queues_config) do
+    children = amqp_data_consumers_childspecs(queues_config)
+
+    Enum.each(children, fn child ->
+      DynamicSupervisor.start_child(AMQPDataConsumer.Supervisor, child)
+    end)
+  end
+
+  defp amqp_data_consumers_childspecs(queues_config) do
+    queue_prefix = queues_config[:prefix]
+    queue_total = queues_config[:total_count]
+    max_index = queue_total - 1
+
+    for queue_index <- 0..max_index do
+      queue_name = "#{queue_prefix}#{queue_index}"
+
+      init_args = [
+        queue_name: queue_name,
+        queue_index: queue_index
+      ]
+
+      {AMQPDataConsumer, init_args}
+    end
+  end
+end

lib/consumer/consumers_supervisor.ex

@@ -29,9 +29,29 @@ defmodule Mississippi.Consumer.ConsumersSupervisor do
       {Registry, [keys: :unique, name: DataUpdater.Registry, members: :auto]},
       {Registry, [keys: :unique, name: MessageTracker.Registry, members: :auto]},
       {Registry, [keys: :unique, name: AMQPDataConsumer.Registry, members: :auto]},
-      {DataUpdater.Supervisor, message_handler: message_handler},
-      {DynamicSupervisor, strategy: :one_for_one, name: MessageTracker.Supervisor, members: :auto},
-      {AMQPDataConsumer.Supervisor, queues_config: queues_config}
+      {DynamicSupervisor,
+       strategy: :one_for_one,
+       name: DataUpdater.Supervisor,
+       members: :auto,
+       process_redistribution: :active,
+       extra_arguments: [message_handler: message_handler],
+       distribution_strategy: Horde.UniformQuorumDistribution},
+      {DynamicSupervisor,
+       strategy: :one_for_one,
+       name: MessageTracker.Supervisor,
+       members: :auto,
+       process_redistribution: :active,
+       distribution_strategy: Horde.UniformQuorumDistribution},
+      {DynamicSupervisor,
+       strategy: :one_for_one,
+       name: AMQPDataConsumer.Supervisor,
+       members: :auto,
+       process_redistribution: :active,
+       distribution_strategy: Horde.UniformQuorumDistribution},
+      # This will make queue listeners start after re-sharding in a multi-node cluster
+      {NodeListener, queues_config},
+      # This will make queue listeners start in a single-node cluster
+      {AMQPDataConsumer.Starter, queues_config}
     ]
 
     opts = [strategy: :rest_for_one]
@@ -66,16 +86,16 @@ defmodule Mississippi.Consumer.ConsumersSupervisor do
               ],
               range_start: [
                 type: :non_neg_integer,
-                default: 0,
                 doc: """
                 The start index of the range of queues that this Mississippi consumer instance will handle.
+                This option is deprecated and will be ignored.
                 """
               ],
               range_end: [
                 type: :non_neg_integer,
-                default: 127,
                 doc: """
                 The end index of the range of queues that this Mississippi consumer instance will handle.
+                This option is deprecated and will be ignored.
                 """
               ],
               prefix: [

lib/consumer/data_updater.ex

@@ -14,6 +14,7 @@ defmodule Mississippi.Consumer.DataUpdater do
   use GenServer, restart: :temporary
   use Efx
 
+  alias Horde.DynamicSupervisor
   alias Horde.Registry
   alias Mississippi.Consumer.DataUpdater
   alias Mississippi.Consumer.DataUpdater.State
@@ -51,7 +52,10 @@ defmodule Mississippi.Consumer.DataUpdater do
           {:ok, pid()} | {:error, :data_updater_start_fail}
   defeffect get_data_updater_process(sharding_key) do
     # TODO bring back :offload_start (?)
-    case DataUpdater.Supervisor.start_child({DataUpdater, sharding_key: sharding_key}) do
+    case DynamicSupervisor.start_child(
+           DataUpdater.Supervisor,
+           {DataUpdater, sharding_key: sharding_key}
+         ) do
       {:ok, pid} ->
         {:ok, pid}
 
@@ -71,6 +75,10 @@ defmodule Mississippi.Consumer.DataUpdater do
     end
   end
 
+  def start_link(start_args) do
+    start_link({:message_handler, Mississippi.Consumer.DataUpdater.Handler.Impl}, start_args)
+  end
+
   def start_link(extra_args, start_args) do
     {:message_handler, message_handler} = extra_args
     sharding_key = Keyword.fetch!(start_args, :sharding_key)

lib/consumer/node_listener.ex

@@ -0,0 +1,30 @@
+# Copyright 2025 SECO Mind Srl
+# SPDX-License-Identifier: Apache-2.0
+
+defmodule NodeListener do
+  @moduledoc false
+  use GenServer
+
+  alias Mississippi.Consumer.AMQPDataConsumer.Starter
+
+  require Logger
+
+  def start_link(args), do: GenServer.start_link(__MODULE__, args)
+
+  def init(queues_config) do
+    :net_kernel.monitor_nodes(true, node_type: :visible)
+    {:ok, queues_config}
+  end
+
+  def handle_info({:nodeup, node, node_type}, queues_config) do
+    _ = Logger.info("Node #{inspect(node)} of type #{inspect(node_type)} is up")
+    _ = Starter.start_consumers(queues_config)
+    {:noreply, queues_config}
+  end
+
+  def handle_info({:nodedown, node, node_type}, queues_config) do
+    _ = Logger.info("Node #{inspect(node)} of type #{inspect(node_type)} is down")
+    _ = Starter.start_consumers(queues_config)
+    {:noreply, queues_config}
+  end
+end

test/consumer/data_updater_test.exs

@@ -18,7 +18,13 @@ defmodule Mississippi.Consumer.DataUpdater.Test do
   setup_all do
     start_supervised!({Registry, [keys: :unique, name: DataUpdater.Registry]})
 
-    start_supervised!({DataUpdater.Supervisor, message_handler: MockMessageHandler})
+    start_supervised(
+      {DynamicSupervisor,
+       strategy: :one_for_one,
+       name: DataUpdater.Supervisor,
+       members: :auto,
+       extra_arguments: [message_handler: MockMessageHandler]}
+    )
 
     :ok
   end