fix(sync-service): promote adopted dependency consumers to transaction writes

.changeset/promote-inner-consumer-write-unit.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+Promote an already-running standalone shape consumer to whole-transaction writes when a subquery materializer subscribes to it. If a fragmented transaction is already in progress, the subscription is completed at its commit boundary so the materializer never misses the already-written head of that transaction.

packages/sync-service/lib/electric/shapes/consumer.ex

@@ -83,7 +83,7 @@ defmodule Electric.Shapes.Consumer do
   def subscribe_materializer(stack_id, shape_handle, pid) do
     stack_id
     |> consumer_pid(shape_handle)
-    |> GenServer.call({:subscribe_materializer, pid})
+    |> GenServer.call({:subscribe_materializer, pid}, :infinity)
   end
 
   @spec whereis(Electric.stack_id(), Electric.shape_handle()) :: pid() | nil
@@ -217,10 +217,21 @@ defmodule Electric.Shapes.Consumer do
     end
   end
 
+  # A materializer's initial history read is bounded by the last committed offset in
+  # storage, while `state.latest_offset` advances with every fragment written in
+  # `write_unit: :txn_fragment` mode. Replying mid-transaction would hand the materializer an
+  # offset past what it can read, silently losing the already-written head of the
+  # transaction. Park the subscription until the commit boundary instead — see
+  # `complete_deferred_materializer_subscription/1`.
+  def handle_call({:subscribe_materializer, pid}, from, %State{pending_txn: pending_txn} = state)
+      when is_write_unit_txn_fragment(state.write_unit) and not is_nil(pending_txn) do
+    Logger.debug("Deferring materializer subscription for #{state.shape_handle} until commit")
+    state = %{state | deferred_materializer_subscription: {pid, from}}
+    {:noreply, state, state.hibernate_after}
+  end
+
   def handle_call({:subscribe_materializer, pid}, _from, state) do
-    Logger.debug("Subscribing materializer for #{state.shape_handle}")
-    Process.monitor(pid, tag: :materializer_down)
-    state = %{state | materializer_subscribed?: true}
+    state = subscribe_materializer(pid, state)
     {:reply, {:ok, state.latest_offset}, state, state.hibernate_after}
   end
 
@@ -467,8 +478,9 @@ defmodule Electric.Shapes.Consumer do
   # ShapeLogCollector holds no pending flush entry for this shape.
   defp consumer_can_suspend?(state) do
     is_snapshot_started(state) and not Shape.has_dependencies(state.shape) and
-      not state.materializer_subscribed? and is_nil(state.pending_txn) and
-      state.txn_offset_mapping == [] and is_nil(state.pending_flush_offset)
+      not state.materializer_subscribed? and is_nil(state.deferred_materializer_subscription) and
+      is_nil(state.pending_txn) and state.txn_offset_mapping == [] and
+      is_nil(state.pending_flush_offset)
   end
 
   defp schedule_suspend_timer(%{suspend_after: nil} = state), do: state
@@ -757,7 +769,7 @@ defmodule Electric.Shapes.Consumer do
   defp skip_txn_fragment(state, %TransactionFragment{} = txn_fragment) do
     %{state | pending_txn: nil}
     |> consider_flushed(txn_fragment.last_log_offset)
-    |> clear_pending_flush_offset()
+    |> finish_pending_txn()
   end
 
   # This function does similar things to do_handle_txn/2 but with the following simplifications:
@@ -916,7 +928,7 @@ defmodule Electric.Shapes.Consumer do
       %{state | pending_txn: nil}
       |> consider_flushed(txn_fragment.last_log_offset)
     end
-    |> clear_pending_flush_offset()
+    |> finish_pending_txn()
   end
 
   def process_buffered_txn_fragments(%State{buffer: buffer} = state) do
@@ -1197,14 +1209,43 @@ defmodule Electric.Shapes.Consumer do
     state
   end
 
-  # After a pending transaction completes and txn_offset_mapping is populated,
-  # process the deferred flushed offset (if any).
+  # After a pending transaction completes and txn_offset_mapping is populated, process the
+  # deferred flushed offset (if any), then complete a materializer subscription that arrived
+  # while the transaction was in progress.
   #
   # Even if the most recent transaction is skipped or no changes from it end up satisfying the
   # shape's `where` condition, Storage may have signaled a flush offset from the previous transaction
   # while we were still processing fragments of the current one. Therefore this function must
   # be called any time `state.pending_txn` is reset to nil in a multi-fragment transaction
   # processing setting.
+  defp finish_pending_txn(state) do
+    state
+    |> clear_pending_flush_offset()
+    |> complete_deferred_materializer_subscription()
+  end
+
+  defp complete_deferred_materializer_subscription(
+         %{deferred_materializer_subscription: nil} = state
+       ),
+       do: state
+
+  # The transaction that was in flight when the materializer called us has now committed, so
+  # `state.latest_offset` sits on a commit boundary and everything up to it is readable from
+  # storage. Subscribe and release the materializer's blocked call.
+  defp complete_deferred_materializer_subscription(
+         %{deferred_materializer_subscription: {pid, from}} = state
+       ) do
+    state = subscribe_materializer(pid, %{state | deferred_materializer_subscription: nil})
+    GenServer.reply(from, {:ok, state.latest_offset})
+    state
+  end
+
+  defp subscribe_materializer(pid, state) do
+    Logger.debug("Subscribing materializer for #{state.shape_handle}")
+    Process.monitor(pid, tag: :materializer_down)
+    State.mark_materializer_subscribed(state)
+  end
+
   defp clear_pending_flush_offset(%{pending_flush_offset: nil} = state), do: state
 
   defp clear_pending_flush_offset(%{pending_flush_offset: flushed_offset} = state) do

packages/sync-service/lib/electric/shapes/consumer/state.ex

@@ -26,6 +26,11 @@ defmodule Electric.Shapes.Consumer.State do
     buffer: [],
     txn_offset_mapping: [],
     materializer_subscribed?: false,
+    # `{pid, from}` of a materializer whose `subscribe_materializer` call arrived while a
+    # fragmented transaction was in progress. The subscription is completed (and the call
+    # replied to) once that transaction commits, so the materializer's initial history read
+    # starts from a commit boundary.
+    deferred_materializer_subscription: nil,
     terminating?: false,
     buffering?: false,
     # Based on the write unit value, consumer will either buffer txn fragments in memory until
@@ -135,6 +140,27 @@ defmodule Electric.Shapes.Consumer.State do
     }
   end
 
+  @doc """
+  Record that a materializer is subscribed to this consumer and promote the consumer to
+  whole-transaction writes: a materializer needs transaction-atomic updates.
+
+  A fragment-streaming consumer must only be promoted with no transaction in progress
+  (`pending_txn: nil`), so that already-written fragments are never mixed with a buffered
+  tail of the same transaction. The consumer guarantees this by deferring subscriptions that
+  arrive mid-transaction until the commit boundary. A consumer already buffering whole
+  transactions can be subscribed at any point.
+  """
+  @spec mark_materializer_subscribed(t()) :: t()
+  def mark_materializer_subscribed(%__MODULE__{write_unit: @write_unit_txn} = state) do
+    %{state | materializer_subscribed?: true}
+  end
+
+  def mark_materializer_subscribed(
+        %__MODULE__{write_unit: @write_unit_txn_fragment, pending_txn: nil} = state
+      ) do
+    %{state | materializer_subscribed?: true, write_unit: @write_unit_txn}
+  end
+
   @doc """
   After the storage is ready, initialize the state with info from storage and writer state.
   """

packages/sync-service/test/electric/shape_cache_test.exs

@@ -1154,6 +1154,37 @@ defmodule Electric.ShapeCacheTest do
       assert %{write_unit: :txn} = :sys.get_state(dependency_pid)
     end
 
+    test "promotes a running standalone consumer when it becomes a dependency", ctx do
+      Support.TestUtils.patch_snapshotter(fn parent,
+                                             shape_handle,
+                                             _shape,
+                                             %{snapshot_fun: snapshot_fun} ->
+        GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_100})
+        snapshot_fun.([])
+        GenServer.cast(parent, {:snapshot_started, shape_handle})
+      end)
+
+      [dependency_shape] = @shape_with_subquery.shape_dependencies
+
+      {dependency_handle, _} =
+        ShapeCache.get_or_create_shape_handle(dependency_shape, ctx.stack_id)
+
+      assert ShapeCache.await_snapshot_start(dependency_handle, ctx.stack_id) == :started
+
+      dependency_pid = Electric.Shapes.Consumer.whereis(ctx.stack_id, dependency_handle)
+      assert %{write_unit: :txn_fragment} = :sys.get_state(dependency_pid)
+
+      {parent_handle, _} =
+        ShapeCache.get_or_create_shape_handle(@shape_with_subquery, ctx.stack_id)
+
+      assert ShapeCache.await_snapshot_start(parent_handle, ctx.stack_id) == :started
+
+      {:ok, parent_shape} = ShapeCache.fetch_shape_by_handle(parent_handle, ctx.stack_id)
+      assert parent_shape.shape_dependencies_handles == [dependency_handle]
+      assert Electric.Shapes.Consumer.whereis(ctx.stack_id, dependency_handle) == dependency_pid
+      assert %{write_unit: :txn} = :sys.get_state(dependency_pid)
+    end
+
     test "should wait for consumer to come up", ctx do
       Support.TestUtils.patch_snapshotter(fn parent, shape_handle, _, _ ->
         GenServer.cast(parent, {:pg_snapshot_known, shape_handle, @pg_snapshot_xmin_100})

packages/sync-service/test/electric/shapes/consumer/state_test.exs

@@ -109,6 +109,30 @@ defmodule Electric.Shapes.Consumer.StateTest do
     end
   end
 
+  describe "mark_materializer_subscribed/1" do
+    setup [:with_stack_id_from_test]
+
+    setup %{stack_id: stack_id} do
+      shape = %Shape{root_table: {"public", "items"}, root_table_id: 1}
+      state = State.new(stack_id, "test-handle", shape)
+      %{state: state}
+    end
+
+    test "promotes a fragment-streaming consumer to transaction writes", %{state: state} do
+      state = %{state | write_unit: :txn_fragment}
+
+      assert %{materializer_subscribed?: true, write_unit: :txn} =
+               State.mark_materializer_subscribed(state)
+    end
+
+    test "subscribes a transaction-buffering consumer even mid-transaction", %{state: state} do
+      state = %{state | write_unit: :txn, pending_txn: %{}}
+
+      assert %{materializer_subscribed?: true, write_unit: :txn, pending_txn: %{}} =
+               State.mark_materializer_subscribed(state)
+    end
+  end
+
   describe "initialize/3" do
     setup [:with_stack_id_from_test, :with_in_memory_storage]
 

packages/sync-service/test/electric/shapes/consumer_test.exs

@@ -13,6 +13,7 @@ defmodule Electric.Shapes.ConsumerTest do
   alias Electric.Shapes
   alias Electric.Shapes.Shape
   alias Electric.Shapes.Consumer
+  alias Electric.Shapes.Consumer.Materializer
 
   alias Support.StubInspector
 
@@ -28,7 +29,8 @@ defmodule Electric.Shapes.ConsumerTest do
       register_as_replication_client: 1,
       complete_txn_fragment: 3,
       txn_fragments: 3,
-      txn_fragment: 4
+      txn_fragment: 4,
+      wait_until: 1
     ]
 
   @receive_timeout 1_000
@@ -1958,6 +1960,69 @@ defmodule Electric.Shapes.ConsumerTest do
       assert [] == :ets.tab2list(table)
     end
 
+    test "materializer subscribing mid-transaction sees the whole transaction", ctx do
+      {shape_handle, _} = ShapeCache.get_or_create_shape_handle(@shape1, ctx.stack_id)
+      :started = ShapeCache.await_snapshot_start(shape_handle, ctx.stack_id)
+
+      consumer_pid = Consumer.whereis(ctx.stack_id, shape_handle)
+      assert %{pending_txn: nil, write_unit: :txn_fragment} = :sys.get_state(consumer_pid)
+
+      xid = 11
+      lsn = Lsn.from_integer(10)
+
+      new_record = fn id, offset ->
+        %Changes.NewRecord{
+          relation: {"public", "test_table"},
+          record: %{"id" => id},
+          log_offset: offset
+        }
+      end
+
+      begin_fragment =
+        txn_fragment(xid, lsn, [new_record.("1", LogOffset.new(lsn, 0))], has_begin?: true)
+
+      assert :ok = ShapeLogCollector.handle_event(begin_fragment, ctx.stack_id)
+      assert %{pending_txn: %{}, write_unit: :txn_fragment} = :sys.get_state(consumer_pid)
+
+      materializer_opts = %{stack_id: ctx.stack_id, shape_handle: shape_handle}
+
+      materializer_pid =
+        start_link_supervised!(
+          {Materializer,
+           Map.merge(materializer_opts, %{
+             storage: ctx.storage,
+             columns: ["id"],
+             materialized_type: {:array, :int8}
+           })}
+        )
+
+      # The subscription is parked until the in-flight transaction commits.
+      assert wait_until(fn ->
+               match?(
+                 %{deferred_materializer_subscription: {^materializer_pid, _}},
+                 :sys.get_state(consumer_pid)
+               )
+             end)
+
+      assert %{materializer_subscribed?: false, write_unit: :txn_fragment} =
+               :sys.get_state(consumer_pid)
+
+      commit_fragment =
+        txn_fragment(xid, lsn, [new_record.("2", LogOffset.new(lsn, 1))], has_commit?: true)
+
+      assert :ok = ShapeLogCollector.handle_event(commit_fragment, ctx.stack_id)
+
+      assert %{
+               pending_txn: nil,
+               write_unit: :txn,
+               materializer_subscribed?: true,
+               deferred_materializer_subscription: nil
+             } = :sys.get_state(consumer_pid)
+
+      assert :ok = Materializer.wait_until_ready(materializer_opts)
+      assert Materializer.get_link_values(materializer_opts) == MapSet.new([1, 2])
+    end
+
     @tag with_pure_file_storage_opts: [flush_period: 1]
     test "writes txn fragments to storage immediately but keeps txn boundaries when flushing",
          ctx do