(DRAFT) Oracle restore tests - and example fixes

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

@@ -20,13 +20,16 @@ defmodule Electric.MonitoredCoreSupervisor do
     Logger.metadata(stack_id: stack_id)
     Electric.Telemetry.Sentry.set_tags_context(stack_id: stack_id)
 
+    persistent_kv = Keyword.get(opts, :persistent_kv)
+
     children = [
       {Electric.StatusMonitor, stack_id: stack_id},
       {Electric.ShapeCache.ShapeCleaner.CleanupTaskSupervisor,
        [stack_id: stack_id] ++ Keyword.get(tweaks, :shape_cleaner_opts, [])},
       {Electric.ShapeCache.ShapeStatus.ShapeDb.Supervisor,
        Keyword.take(opts, [:stack_id, :shape_db_opts])},
-      {Electric.ShapeCache.ShapeStatusOwner, [stack_id: stack_id]},
+      {Electric.ShapeCache.ShapeStatusOwner,
+       [stack_id: stack_id, persistent_kv: persistent_kv]},
       {Electric.CoreSupervisor, opts}
     ]
 

packages/sync-service/lib/electric/replication/shape_log_collector.ex

@@ -163,6 +163,36 @@ defmodule Electric.Replication.ShapeLogCollector do
     GenServer.cast(name(stack_id), {:writer_flushed, shape_handle, offset})
   end
 
+  @doc """
+  Stop accepting new replication events and block until every consumer
+  has flushed up to the most recently seen offset.
+
+  Used by `Electric.StackSupervisor.ShutdownCoordinator` to drive a
+  coordinated safe shutdown: callers must first stop the replication
+  client (so no new events arrive) and then call `drain/2`. When this
+  returns `:ok`, every shape's writer has persisted the same
+  `last_persisted_txn_offset`.
+
+  If the timeout expires before all shapes catch up, returns
+  `{:error, :timeout}`. The caller should treat that as a dirty
+  shutdown — no clean-shutdown marker should be written.
+  """
+  @spec drain(Electric.stack_id(), timeout()) :: :ok | {:error, :timeout | :no_collector}
+  def drain(stack_id, timeout) do
+    case GenServer.whereis(name(stack_id)) do
+      nil ->
+        {:error, :no_collector}
+
+      pid ->
+        try do
+          GenServer.call(pid, :drain, timeout)
+        catch
+          :exit, {:timeout, _} -> {:error, :timeout}
+          :exit, {:noproc, _} -> {:error, :no_collector}
+        end
+    end
+  end
+
   @doc """
   Returns the list of currently active shapes being tracked
   in the shape matching filters.
@@ -382,6 +412,20 @@ defmodule Electric.Replication.ShapeLogCollector do
     {:reply, :ok, Map.put(state, :last_processed_offset, offset)}
   end
 
+  def handle_call(:drain, from, state) do
+    state = Map.put(state, :draining?, true)
+
+    if FlushTracker.empty?(state.flush_tracker) do
+      {:reply, :ok, state}
+    else
+      {:noreply, Map.put(state, :drain_waiter, from)}
+    end
+  end
+
+  def handle_call({:handle_event, _, _}, _from, %{draining?: true} = state) do
+    {:reply, {:error, :draining}, state}
+  end
+
   def handle_call({:handle_event, _, _}, _from, state)
       when not is_ready_to_process(state) do
     {:reply, {:error, :not_ready}, state}
@@ -414,9 +458,19 @@ defmodule Electric.Replication.ShapeLogCollector do
   end
 
   def handle_cast({:writer_flushed, shape_id, offset}, state) do
-    {:noreply,
-     state
-     |> Map.update!(:flush_tracker, &FlushTracker.handle_flush_notification(&1, shape_id, offset))}
+    state =
+      state
+      |> Map.update!(:flush_tracker, &FlushTracker.handle_flush_notification(&1, shape_id, offset))
+
+    state =
+      if state[:drain_waiter] && FlushTracker.empty?(state.flush_tracker) do
+        GenServer.reply(state.drain_waiter, :ok)
+        Map.delete(state, :drain_waiter)
+      else
+        state
+      end
+
+    {:noreply, state}
   end
 
   def handle_cast(

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

@@ -286,8 +286,14 @@ defmodule Electric.ShapeCache do
 
     Electric.Replication.PublicationManager.wait_for_restore(state.stack_id)
 
-    # 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.
+    # 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)
+
     ShapeLogCollector.mark_as_ready(state.stack_id)
 
     duration = System.monotonic_time() - start_time
@@ -305,6 +311,39 @@ 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} ->
+          _ = Electric.Shapes.Consumer.await_snapshot_start(state.stack_id, handle)
+
+        _ ->
+          :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/shape_cache/shape_status.ex

@@ -291,6 +291,14 @@ defmodule Electric.ShapeCache.ShapeStatus do
       [] ->
         :error
     end
+  rescue
+    # The shape_meta_table is created by ShapeStatusOwner. During a stack
+    # restart there's a window when the old owner's table has been freed
+    # and the new owner hasn't recreated it yet. An inflight HTTP request
+    # held by Bandit can wake up in that window. Treat a missing table
+    # the same as a missing entry — the caller can fall back to
+    # fetch_handle_by_shape (which goes through the live SQLite store).
+    ArgumentError -> :error
   end
 
   @spec mark_snapshot_started(stack_id(), shape_handle()) :: :ok | :error

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

@@ -14,10 +14,16 @@ defmodule Electric.ShapeCache.ShapeStatusOwner do
   use GenServer
 
   alias Electric.ShapeCache.ShapeStatus
+  alias Electric.ShapeCache.ShapeStatus.ShapeDb
+  alias Electric.ShapeCache.Storage
+  alias Electric.StackSupervisor.ShutdownCoordinator
 
   require Logger
 
-  @schema NimbleOptions.new!(stack_id: [type: :string, required: true])
+  @schema NimbleOptions.new!(
+            stack_id: [type: :string, required: true],
+            persistent_kv: [type: :any, required: false]
+          )
 
   def name(stack_id) do
     Electric.ProcessRegistry.name(stack_id, __MODULE__)
@@ -50,11 +56,14 @@ defmodule Electric.ShapeCache.ShapeStatusOwner do
 
     :ok = Electric.LsnTracker.initialize(stack_id)
 
-    {:ok, %{stack_id: stack_id}, {:continue, :initialize}}
+    {:ok, %{stack_id: stack_id, persistent_kv: Map.get(config, :persistent_kv)},
+     {:continue, :initialize}}
   end
 
   @impl true
   def handle_continue(:initialize, state) do
+    maybe_reset_after_dirty_shutdown(state)
+
     # Initialize shape metadata from SQLite so we can serve shapes in
     # read-only mode before the advisory lock is acquired
     :ok = ShapeStatus.initialize(state.stack_id)
@@ -67,6 +76,45 @@ defmodule Electric.ShapeCache.ShapeStatusOwner do
     {:noreply, state, :hibernate}
   end
 
+  # If the previous stack shutdown wasn't clean, different shapes' on-disk
+  # `last_persisted_txn_offset`s may be at different LSNs (writers flush
+  # independently and the supervisor doesn't coordinate a final flush
+  # before tear-down). Recovering from such a state would mean that an
+  # outer subquery shape's `state.views` (seeded from the materializer,
+  # which is rebuilt from the inner shape's history) and its on-disk
+  # storage land at different LSNs, producing duplicate inserts and
+  # orphan deletes when live events resume.
+  #
+  # Sound recovery from a partial shutdown is hard; the conservative
+  # answer is to throw away every shape's persisted state and force a
+  # re-snapshot. Clients re-request their shapes from scratch on the
+  # next poll.
+  defp maybe_reset_after_dirty_shutdown(%{persistent_kv: nil}), do: :ok
+
+  defp maybe_reset_after_dirty_shutdown(%{persistent_kv: persistent_kv} = state) do
+    if ShutdownCoordinator.consume_clean_shutdown_marker(persistent_kv) do
+      :ok
+    else
+      Logger.warning(
+        "Stack #{state.stack_id} did not shut down cleanly last time; " <>
+          "dropping all shape data to avoid inconsistent recovery"
+      )
+
+      # Wipe before ShapeStatus.initialize so that initialize sees an
+      # empty SQLite store and creates a fresh, empty ETS cache. We
+      # can't call ShapeStatus.reset/1 here because the ETS tables it
+      # tries to clear haven't been created yet (initialize is what
+      # creates them).
+      :ok = ShapeDb.reset(state.stack_id)
+
+      state.stack_id
+      |> Storage.for_stack()
+      |> Storage.cleanup_all!()
+
+      :ok
+    end
+  end
+
   @impl true
   def handle_call(:refresh, _from, state) do
     :ok = ShapeStatus.refresh(state.stack_id)

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

@@ -970,6 +970,13 @@ defmodule Electric.Shapes.Api do
       _ ->
         :no_change
     end
+  rescue
+    # During a stack restart, the per-stack ETS tables this path reads
+    # (shape_meta_table, write_buffer_shapes, …) may be transiently absent
+    # for a held long-poll that wakes up between the old owner dying and
+    # the new owner recreating the table. Treat as no_change rather than
+    # bubbling an ArgumentError up to the request handler.
+    ArgumentError -> :no_change
   end
 
   defp stream_sse_events(%Request{} = request) do

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

@@ -1063,6 +1063,18 @@ defmodule Electric.Shapes.Consumer do
          opts
        ) do
     Materializer.new_changes(Map.take(state, [:stack_id, :shape_handle]), changes_or_bounds, opts)
+  catch
+    # The consumer monitors the materializer; if the materializer died the
+    # :DOWN message is already in our mailbox and handle_materializer_down/2
+    # will run after the current handle_event/handle_call completes.
+    # Treat a `:noproc` (or transient `:normal`/`:shutdown` exit) here as
+    # the same condition: don't crash the consumer (which would route into
+    # the abnormal-shutdown path of handle_writer_termination and remove
+    # the shape from disk).
+    :exit, {:noproc, _} -> :ok
+    :exit, :noproc -> :ok
+    :exit, {:normal, _} -> :ok
+    :exit, {:shutdown, _} -> :ok
   end
 
   defp notify_materializer_of_new_changes(_state, _changes_or_bounds, _opts), do: :ok

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

@@ -171,33 +171,72 @@ 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.
   """
-  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.
+              %{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/lib/electric/stack_supervisor.ex

@@ -427,7 +427,12 @@ defmodule Electric.StackSupervisor do
         {Electric.MonitoredCoreSupervisor,
          stack_id: stack_id,
          connection_manager_opts: connection_manager_opts,
-         shape_db_opts: config.shape_db_opts}
+         shape_db_opts: config.shape_db_opts,
+         persistent_kv: config.persistent_kv},
+        # ShutdownCoordinator must be the LAST child so it terminates
+        # FIRST on graceful stack shutdown — see the module docstring.
+        {Electric.StackSupervisor.ShutdownCoordinator,
+         stack_id: stack_id, persistent_kv: config.persistent_kv}
       ]
       |> Enum.reject(&is_nil/1)