⚡️ SMS.put_messages validation and early GenServer.reply to SP

Author: RTLS

lib/sequin/databases_runtime/message_handler.ex

@@ -426,7 +426,7 @@ defmodule Sequin.DatabasesRuntime.SlotProcessor.MessageHandler do
         put_messages(consumer_id, messages_to_ingest, attempt + 1)
 
       {:error, error} ->
-        Health.put_event(:sink_consumer, consumer_id, %Event{slug: :messages_ingested, status: :failed, error: error})
+        Health.put_event(:sink_consumer, consumer_id, %Event{slug: :messages_ingested, status: :fail, error: error})
         {:error, error}
     end
   end

lib/sequin/databases_runtime/slot_message_store.ex

@@ -328,33 +328,48 @@ defmodule Sequin.DatabasesRuntime.SlotMessageStore do
   end
 
   @impl GenServer
-  def handle_call({:put_messages, messages}, _from, %State{} = state) do
+  def handle_call({:put_messages, messages}, from, %State{} = state) do
     execute_timed(:put_messages, fn ->
       now = Sequin.utc_now()
 
       messages =
         messages
+        |> Stream.reject(&State.message_exists?(state, &1))
         |> Stream.map(&%{&1 | ack_id: Sequin.uuid4(), ingested_at: now})
-        # We may be receiving messages that we've already ingested and persisted, filter out
-        # Do we receive messages that we have ingested but not persisted?
-        |> Enum.filter(&(not State.is_message_persisted?(state, &1)))
-
-      {to_persist, to_put} =
-        if state.consumer.status == :disabled do
-          {Enum.to_list(messages), []}
-        else
-          Enum.split_with(messages, &State.is_message_group_persisted?(state, &1.group_id))
-        end
-
-      with {:ok, state} <- State.put_messages(state, to_put),
-           :ok <- upsert_messages(state, to_persist),
-           {:ok, state} <- State.put_persisted_messages(state, to_persist) do
-        Health.put_event(state.consumer, %Event{slug: :messages_ingested, status: :success})
-        :syn.publish(:consumers, {:messages_ingested, state.consumer.id}, :messages_ingested)
+        |> Enum.to_list()
+
+      # Validate first
+      case State.validate_put_messages(state, messages) do
+        {:ok, _incoming_payload_size_bytes} ->
+          # Reply early since validation passed. This frees up the SlotProcessor to continue
+          # calling other SMSs, accumulating messages, etc.
+          GenServer.reply(from, :ok)
+
+          execute_timed(:put_messages_after_reply, fn ->
+            {to_persist, to_put} =
+              if state.consumer.status == :disabled do
+                {messages, []}
+              else
+                Enum.split_with(messages, &State.is_message_group_persisted?(state, &1.group_id))
+              end
+
+            {:ok, state} = State.put_messages(state, to_put)
+
+            :ok = upsert_messages(state, to_persist)
+            {:ok, state} = State.put_persisted_messages(state, to_persist)
+
+            Health.put_event(state.consumer, %Event{slug: :messages_ingested, status: :success})
+            :syn.publish(:consumers, {:messages_ingested, state.consumer.id}, :messages_ingested)
+
+            if state.test_pid do
+              send(state.test_pid, {:put_messages_done, state.consumer.id})
+            end
+
+            {:noreply, state}
+          end)
 
-        {:reply, :ok, state}
-      else
         {:error, error} ->
+          # Reply with error if validation fails
           {:reply, {:error, error}, state}
       end
     end)

lib/sequin/databases_runtime/slot_message_store_state.ex

@@ -49,24 +49,17 @@ defmodule Sequin.DatabasesRuntime.SlotMessageStore.State do
     :"slot_message_store_state_ordered_cursors_consumer_#{consumer.seq}"
   end
 
-  @spec put_messages(State.t(), list(message()) | Enumerable.t()) ::
-          {:ok, State.t()} | {:error, Error.t()}
-  def put_messages(%State{} = state, messages, opts \\ []) do
+  @spec validate_put_messages(State.t(), list(message()) | Enumerable.t(), keyword()) ::
+          {:ok, non_neg_integer()} | {:error, Error.t()}
+  def validate_put_messages(%State{} = state, messages, opts \\ []) do
     skip_limit_check? = Keyword.get(opts, :skip_limit_check?, false)
 
-    # Create new messages map while filtering out existing messages
-    new_messages =
-      messages
-      |> Stream.reject(&Map.has_key?(state.messages, {&1.commit_lsn, &1.commit_idx}))
-      |> Map.new(&{{&1.commit_lsn, &1.commit_idx}, &1})
-
-    # Only count payload sizes for new messages
-    incoming_payload_size_bytes = Enum.sum_by(Map.values(new_messages), & &1.payload_size_bytes)
+    incoming_payload_size_bytes = Enum.sum_by(messages, & &1.payload_size_bytes)
 
     bytes_exceeded? =
       state.payload_size_bytes + incoming_payload_size_bytes > state.max_memory_bytes
 
-    messages_exceeded? = map_size(new_messages) + map_size(state.messages) > state.setting_max_messages
+    messages_exceeded? = length(messages) + map_size(state.messages) > state.setting_max_messages
 
     cond do
       not skip_limit_check? and bytes_exceeded? ->
@@ -76,28 +69,41 @@ defmodule Sequin.DatabasesRuntime.SlotMessageStore.State do
         {:error, Error.invariant(message: "Message count limit exceeded", code: :payload_size_limit_exceeded)}
 
       true ->
-        # Insert into ETS
-        ets_keys = Enum.map(new_messages, fn {commit_tuple, _msg} -> {commit_tuple} end)
-
-        state.consumer
-        |> ordered_cursors_table()
-        |> :ets.insert(ets_keys)
-
-        ack_ids_to_cursor_tuples =
-          Map.new(new_messages, fn {_commit_tuple, msg} -> {msg.ack_id, {msg.commit_lsn, msg.commit_idx}} end)
-
-        state =
-          %{
-            state
-            | messages: Map.merge(state.messages, new_messages),
-              ack_ids_to_cursor_tuples: Map.merge(state.ack_ids_to_cursor_tuples, ack_ids_to_cursor_tuples),
-              payload_size_bytes: state.payload_size_bytes + incoming_payload_size_bytes
-          }
+        {:ok, incoming_payload_size_bytes}
+    end
+  end
 
-        {:ok, state}
+  @spec put_messages(State.t(), list(message()) | Enumerable.t(), keyword()) ::
+          {:ok, State.t()} | {:error, Error.t()}
+  def put_messages(%State{} = state, messages, opts \\ []) do
+    messages = Enum.reject(messages, &message_exists?(state, &1))
+
+    with {:ok, incoming_payload_size_bytes} <- validate_put_messages(state, messages, opts) do
+      # Insert into ETS
+      ets_keys = Enum.map(messages, fn msg -> {{msg.commit_lsn, msg.commit_idx}} end)
+
+      state.consumer
+      |> ordered_cursors_table()
+      |> :ets.insert(ets_keys)
+
+      cursor_tuples_to_messages = Map.new(messages, fn msg -> {{msg.commit_lsn, msg.commit_idx}, msg} end)
+      ack_ids_to_cursor_tuples = Map.new(messages, fn msg -> {msg.ack_id, {msg.commit_lsn, msg.commit_idx}} end)
+
+      {:ok,
+       %{
+         state
+         | messages: Map.merge(state.messages, cursor_tuples_to_messages),
+           ack_ids_to_cursor_tuples: Map.merge(state.ack_ids_to_cursor_tuples, ack_ids_to_cursor_tuples),
+           payload_size_bytes: state.payload_size_bytes + incoming_payload_size_bytes
+       }}
     end
   end
 
+  @spec message_exists?(State.t(), message()) :: boolean()
+  def message_exists?(%State{} = state, message) do
+    Map.has_key?(state.messages, {message.commit_lsn, message.commit_idx})
+  end
+
   @spec put_persisted_messages(State.t(), list(message()) | Enumerable.t()) :: {:ok, State.t()} | {:error, Error.t()}
   def put_persisted_messages(%State{} = state, messages) do
     persisted_message_groups =

test/sequin/slot_message_store_state_test.exs

@@ -67,6 +67,26 @@ defmodule Sequin.DatabasesRuntime.SlotMessageStoreStateTest do
     end
   end
 
+  describe "validate_put_messages/2" do
+    test "returns payload size if under memory limit", %{state: state} do
+      msg1 = ConsumersFactory.consumer_message(payload_size_bytes: 100)
+      msg2 = ConsumersFactory.consumer_message(payload_size_bytes: 200)
+
+      {:ok, payload_size} = State.validate_put_messages(state, [msg1, msg2])
+
+      assert payload_size == 300
+    end
+
+    test "returns error when exceeding setting_max_accumulated_payload_bytes", %{state: state} do
+      state = %{state | max_memory_bytes: Size.mb(100)}
+
+      msg1 = ConsumersFactory.consumer_message(payload_size_bytes: Size.mb(100))
+      msg2 = ConsumersFactory.consumer_message(payload_size_bytes: Size.mb(100))
+
+      assert {:error, %Error.InvariantError{}} = State.validate_put_messages(state, [msg1, msg2])
+    end
+  end
+
   describe "put_persisted_messages/2" do
     test "updates persisted_message_groups", %{state: state} do
       msg1 = ConsumersFactory.consumer_message(group_id: "group1")

test/sequin/slot_message_store_test.exs

@@ -57,6 +57,9 @@ defmodule Sequin.SlotMessageStoreTest do
 
       :ok = SlotMessageStore.put_messages(consumer.id, [new_message])
 
+      consumer_id = consumer.id
+      assert_receive {:put_messages_done, ^consumer_id}, 1000
+
       persisted_messages = Consumers.list_consumer_messages_for_consumer(consumer)
       assert length(persisted_messages) == 3
 
@@ -380,6 +383,8 @@ defmodule Sequin.SlotMessageStoreTest do
     end
 
     test "persists all messages when consumer is disabled", %{consumer: consumer} do
+      consumer_id = consumer.id
+
       # Put initial message in store while active
       initial_message =
         ConsumersFactory.consumer_message(
@@ -388,6 +393,7 @@ defmodule Sequin.SlotMessageStoreTest do
         )
 
       :ok = SlotMessageStore.put_messages(consumer.id, [initial_message])
+      assert_receive {:put_messages_done, ^consumer_id}, 1000
 
       # Verify message is not persisted yet
       assert [] == Consumers.list_consumer_messages_for_consumer(consumer)
@@ -410,6 +416,7 @@ defmodule Sequin.SlotMessageStoreTest do
       ]
 
       :ok = SlotMessageStore.put_messages(consumer.id, new_messages)
+      assert_receive {:put_messages_done, ^consumer_id}, 1000
 
       # Verify all messages are persisted
       persisted_messages = Consumers.list_consumer_messages_for_consumer(disabled_consumer)