Restore subquery materializer state from disk after server restart

.changeset/restore-subquery-materializer.md

@@ -0,0 +1,12 @@
+---
+'@core/sync-service': patch
+---
+
+Fix subquery shapes diverging from Postgres after a server restart. Shapes
+whose where clause contains a subquery (e.g. `id IN (SELECT ... WHERE active)`)
+could return stale results or a `409 must-refetch` once the stack restarted and
+restored from disk. Two causes are addressed: the dependency materializer now
+replays the full persisted history (snapshot + main log) on startup instead of
+just the first chunk, and dependent (outer) subquery consumers are eagerly
+restarted during restore so their materializer subscription is re-established
+before live events flow.

packages/sync-service/lib/electric/shape_cache.ex

@@ -286,6 +286,14 @@ defmodule Electric.ShapeCache do
 
     Electric.Replication.PublicationManager.wait_for_restore(state.stack_id)
 
+    # Subquery shapes' consumers must be fully initialized before
+    # ShapeLogCollector starts dispatching events. If events flow first,
+    # the materializer can advance past the outer shape's on-disk storage;
+    # the outer consumer's later init would then seed `state.views` from
+    # the advanced materializer view and a subsequent move-in event for
+    # a value already in that seeded view would be dropped as redundant.
+    eagerly_start_subquery_shape_consumers(state)
+
     # Let ShapeLogCollector that it can start processing after finishing this function so that
     # we're subscribed to the producer before it starts forwarding its demand.
     ShapeLogCollector.mark_as_ready(state.stack_id)
@@ -305,6 +313,61 @@ defmodule Electric.ShapeCache do
     {:noreply, state}
   end
 
+  # Shapes whose where clause contains a subquery (`shape_dependencies != []`)
+  # rely on their materializer subscription to be notified of dependency-side
+  # changes. The router only delivers events for a shape when its own
+  # `root_table` changes, so a subquery dependent stays dormant after a
+  # restart until something writes to its own table — movements driven by
+  # the dependency (e.g. parent rows becoming active) never reach its
+  # on-disk view. Restoring it here re-establishes the materializer
+  # subscription so dependency updates flow in.
+  #
+  # `await_snapshot_start/2` is queued *after* the consumer's
+  # `:initialize_shape` info message, so by the time it returns
+  # `EventHandlerBuilder.build` has run and `state.views` is seeded.
+  defp eagerly_start_subquery_shape_consumers(state) do
+    opts = %{
+      stack_id: state.stack_id,
+      action: :restore,
+      otel_ctx: nil,
+      feature_flags: state.feature_flags
+    }
+
+    for {handle, %Shape{shape_dependencies: [_ | _]} = shape} <-
+          ShapeStatus.list_shapes(state.stack_id),
+        is_nil(Electric.Shapes.ConsumerRegistry.whereis(state.stack_id, handle)) do
+      case restore_shape_and_dependencies(handle, shape, opts) do
+        {:ok, _pid} ->
+          # await_snapshot_start/2 is a GenServer.call into the just-started
+          # consumer. If that consumer dies before/during the call it exits;
+          # left unguarded that would propagate out of handle_continue and
+          # crash ShapeCache before mark_as_ready — turning a single shape
+          # that reliably fails its snapshot into a stack-wide restart loop.
+          # A call timeout (the consumer is alive but wedged) exits the same
+          # way. In either case we can't confirm the shape's consumer came up
+          # subscribed-and-correct, and the eager start exists precisely to
+          # guarantee that consistency. Leaving the shape alive-but-unconfirmed
+          # would silently reintroduce the divergence this restore path fixes,
+          # so we purge it (mirroring restore_shape_and_dependencies' own
+          # clean_shape-on-failure) and let the client refetch from scratch.
+          try do
+            _ = Electric.Shapes.Consumer.await_snapshot_start(state.stack_id, handle)
+          catch
+            :exit, reason ->
+              Logger.warning(
+                "Eager subquery consumer await failed for #{handle}: #{inspect(reason)}; " <>
+                  "purging shape to force a clean refetch"
+              )
+
+              clean_shape(handle, state.stack_id)
+          end
+
+        _ ->
+          :ok
+      end
+    end
+  end
+
   @impl GenServer
   def handle_call({:create_or_wait_shape_handle, shape, otel_ctx}, _from, state) do
     if not is_nil(otel_ctx), do: OpenTelemetry.set_current_context(otel_ctx)

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

@@ -171,33 +171,84 @@ defmodule Electric.Shapes.Consumer.Materializer do
   end
 
   def handle_continue({:read_stream, storage}, state) do
-    {:ok, offset, stream} =
-      get_stream_up_to_offset(state.offset, state.subscribed_offset, storage)
-
-    {state, _} =
-      stream
-      |> decode_json_stream()
-      |> apply_changes(state)
-
+    state = read_history_up_to_subscribed(state, storage)
     write_link_values(state)
-
-    {:noreply, %{state | offset: offset}}
+    {:noreply, state}
   end
 
   @doc """
-  Get a stream of log entries from storage, bounded by the subscribed offset.
-
-  The subscribed_offset is the Consumer's latest_offset at the time of subscription.
-  We only read up to this offset to avoid duplicates - any changes after this offset
-  will be delivered via new_changes messages from the Consumer.
+  Replay all of the source shape's persisted history (snapshot + log) up to
+  `state.subscribed_offset` so the materializer's value_counts reflect the
+  on-disk state on startup.
+
+  `Storage.get_log_stream/3` returns at most one chunk per call **for
+  snapshot chunks** — but for the main log it returns the entire requested
+  range `[min_offset, subscribed_offset]` in one call. So we iterate
+  through snapshot chunks using `Storage.get_chunk_end_log_offset/2`,
+  and as soon as the iteration would step into the main log we stop:
+  the previous call already streamed everything up to the subscribed
+  offset. Iterating into the main log would re-read entries already
+  applied, producing duplicate inserts that crash the materializer.
+
+  The subscribed_offset is the Consumer's latest_offset at the time of
+  subscription. We only read up to this offset to avoid duplicates — any
+  changes after this offset will be delivered via new_changes messages
+  from the Consumer.
   """
-  def get_stream_up_to_offset(min_offset, subscribed_offset, storage) do
-    # If subscribed_offset is nil or at/before min_offset, nothing to read
-    if is_nil(subscribed_offset) or is_log_offset_lte(subscribed_offset, min_offset) do
-      {:ok, min_offset, []}
-    else
-      stream = Storage.get_log_stream(min_offset, subscribed_offset, storage)
-      {:ok, subscribed_offset, stream}
+  def read_history_up_to_subscribed(state, storage) do
+    cond do
+      is_nil(state.subscribed_offset) ->
+        state
+
+      is_log_offset_lte(state.subscribed_offset, state.offset) ->
+        state
+
+      true ->
+        stream = Storage.get_log_stream(state.offset, state.subscribed_offset, storage)
+        {state, _} = stream |> decode_json_stream() |> apply_changes(state)
+
+        # If the read just covered the main log (because either the
+        # current offset is already past the snapshot or the next chunk
+        # boundary jumps into real-offset territory), `stream_main_log`
+        # returned the whole range up to `subscribed_offset` in a single
+        # call and we're done.
+        if is_real_offset(state.offset) or is_last_virtual_offset(state.offset) do
+          %{state | offset: state.subscribed_offset}
+        else
+          next_offset = Storage.get_chunk_end_log_offset(state.offset, storage)
+
+          cond do
+            is_nil(next_offset) ->
+              # No further chunks past this offset — we've reached the end.
+              %{state | offset: state.subscribed_offset}
+
+            is_log_offset_lte(next_offset, state.offset) ->
+              # Defensive: chunk_end did not advance. Stop to avoid an
+              # infinite loop. This shouldn't happen in normal operation.
+              Logger.warning(
+                "Materializer chunk iteration did not advance past " <>
+                  "#{inspect(state.offset)} (chunk_end=#{inspect(next_offset)}); " <>
+                  "stopping replay at subscribed_offset to avoid an infinite loop",
+                shape_handle: state.shape_handle
+              )
+
+              %{state | offset: state.subscribed_offset}
+
+            is_log_offset_lte(state.subscribed_offset, next_offset) ->
+              %{state | offset: state.subscribed_offset}
+
+            is_real_offset(next_offset) ->
+              # The next chunk is in the main log, which means the call
+              # we just made (with `state.offset` past the last snapshot
+              # chunk) already streamed the entire main log up to
+              # `subscribed_offset`. Stop — iterating further would
+              # re-read entries we've already applied.
+              %{state | offset: state.subscribed_offset}
+
+            true ->
+              read_history_up_to_subscribed(%{state | offset: next_offset}, storage)
+          end
+        end
     end
   end
 

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

@@ -1407,6 +1407,144 @@ defmodule Electric.ShapeCacheTest do
     end
   end
 
+  describe "wait_for_restore eager subquery consumer start" do
+    setup [
+      :with_noop_publication_manager,
+      :with_log_chunking,
+      :with_registry,
+      :with_shape_log_collector
+    ]
+
+    test "shapes with subquery dependencies have their consumer eagerly started", ctx do
+      %{stack_id: stack_id} = ctx
+      test_pid = self()
+
+      # Pre-add a subquery shape to ShapeStatus, simulating a shape that
+      # was created in a prior incarnation of the stack and is now being
+      # restored. The shape's consumer is NOT yet running.
+      {:ok, shape_handle} = ShapeStatus.add_shape(stack_id, @shape_with_subquery)
+
+      # We don't have a fully wired-up consumer chain in this test, so
+      # short-circuit `start_shape_consumer` and `start_materializer` while
+      # recording the calls. This proves wait_for_restore reaches the start
+      # path for the subquery shape; full consumer lifecycle is covered by
+      # the integration tests in `oracle_restore_test.exs`.
+      Repatch.patch(
+        Electric.Shapes.DynamicConsumerSupervisor,
+        :start_materializer,
+        [mode: :shared],
+        fn _stack_id, _config -> {:ok, self()} end
+      )
+
+      Repatch.patch(
+        Electric.Shapes.DynamicConsumerSupervisor,
+        :start_shape_consumer,
+        [mode: :shared],
+        fn _stack_id, %{shape_handle: handle} ->
+          send(test_pid, {:start_shape_consumer_called, handle})
+          # Returning :error short-circuits restore_shape_and_dependencies'
+          # follow-up calls (initialize_shape, update_last_read_time) which
+          # would fail without a real consumer pid.
+          {:error, :test_short_circuit}
+        end
+      )
+
+      # clean_shape is called on start_shape_consumer error; stub it so the
+      # follow-up cleanup doesn't interfere with the test assertion.
+      Repatch.patch(
+        Electric.ShapeCache.ShapeCleaner,
+        :remove_shape,
+        [mode: :shared],
+        fn _stack_id, _handle -> :ok end
+      )
+
+      activate_mocks_for_descendant_procs(Electric.ShapeCache)
+
+      with_shape_cache(ctx)
+
+      # The eager-start path runs in handle_continue(:wait_for_restore); the
+      # patched start_shape_consumer captures the call.
+      assert_receive {:start_shape_consumer_called, ^shape_handle}, 5_000
+    end
+
+    test "non-subquery shapes are NOT eagerly started", ctx do
+      %{stack_id: stack_id} = ctx
+      test_pid = self()
+
+      # Add a shape with no dependencies — eager-start should skip it.
+      {:ok, _shape_handle} = ShapeStatus.add_shape(stack_id, @shape)
+
+      Repatch.patch(
+        Electric.Shapes.DynamicConsumerSupervisor,
+        :start_shape_consumer,
+        [mode: :shared],
+        fn _stack_id, %{shape_handle: handle} ->
+          send(test_pid, {:start_shape_consumer_called, handle})
+          {:error, :test_short_circuit}
+        end
+      )
+
+      activate_mocks_for_descendant_procs(Electric.ShapeCache)
+
+      with_shape_cache(ctx)
+
+      # Give wait_for_restore time to complete; eager-start should NOT
+      # have called start_shape_consumer for the simple shape.
+      refute_receive {:start_shape_consumer_called, _}, 200
+    end
+
+    test "purges the shape when the eager consumer await fails", ctx do
+      %{stack_id: stack_id} = ctx
+      test_pid = self()
+
+      {:ok, shape_handle} = ShapeStatus.add_shape(stack_id, @shape_with_subquery)
+
+      # Make the consumer start succeed so restore_shape_and_dependencies
+      # returns {:ok, pid} and the eager-start reaches await_snapshot_start.
+      Repatch.patch(
+        Electric.Shapes.DynamicConsumerSupervisor,
+        :start_materializer,
+        [mode: :shared],
+        fn _stack_id, _config -> {:ok, self()} end
+      )
+
+      Repatch.patch(
+        Electric.Shapes.DynamicConsumerSupervisor,
+        :start_shape_consumer,
+        [mode: :shared],
+        fn _stack_id, _config -> {:ok, self()} end
+      )
+
+      # ...then make the await fail (consumer died, or timed out while
+      # wedged) so the catch fires. We can't confirm the shape came up
+      # subscribed-and-correct, so it must be purged rather than left alive.
+      Repatch.patch(
+        Electric.Shapes.Consumer,
+        :await_snapshot_start,
+        [mode: :shared],
+        fn _stack_id, _handle -> exit(:test_consumer_died) end
+      )
+
+      # clean_shape goes through ShapeCleaner.remove_shape; capture the call
+      # to prove the failed shape is purged so a client refetches from scratch.
+      Repatch.patch(
+        Electric.ShapeCache.ShapeCleaner,
+        :remove_shape,
+        [mode: :shared],
+        fn _stack_id, handle ->
+          send(test_pid, {:remove_shape_called, handle})
+          :ok
+        end
+      )
+
+      activate_mocks_for_descendant_procs(Electric.ShapeCache)
+
+      with_shape_cache(ctx)
+
+      assert_receive {:remove_shape_called, ^shape_handle}, 5_000
+    end
+  end
+
   describe "start_consumer_for_handle/2" do
     setup [
       :with_noop_publication_manager,
@@ -1445,17 +1583,20 @@ defmodule Electric.ShapeCacheTest do
                GenServer.whereis(Electric.Shapes.Consumer.Materializer.name(stack_id, dep_handle))
              )
 
-      # Register this test as the connection manager to get "consumers ready" notification
       restart_shape_cache(ctx)
 
       assert [{^dep_handle, _}, {^shape_handle, _}] = ShapeCache.list_shapes(stack_id)
 
-      refute Electric.Shapes.ConsumerRegistry.whereis(stack_id, shape_handle)
-      refute Electric.Shapes.ConsumerRegistry.whereis(stack_id, dep_handle)
-
-      refute GenServer.whereis(Electric.Shapes.Consumer.Materializer.name(stack_id, dep_handle))
-
-      assert {:ok, _pid1} = ShapeCache.start_consumer_for_handle(shape_handle, stack_id)
+      # After restart, ShapeCache eagerly starts subquery shape consumers
+      # in `handle_continue(:wait_for_restore)` so the outer consumer and
+      # its dependency materializer come back up automatically — no
+      # explicit `start_consumer_for_handle/2` call is required. The
+      # continue runs asynchronously after `start_supervised!` returns,
+      # so we wait for the registry to populate.
+      assert wait_until(1000, fn ->
+               not is_nil(Electric.Shapes.ConsumerRegistry.whereis(stack_id, shape_handle)) and
+                 not is_nil(Electric.Shapes.ConsumerRegistry.whereis(stack_id, dep_handle))
+             end)
 
       # Materializer should be started
       assert Process.alive?(