fix: subqueries restoration on startup (#3268)

- Fixed JSON deserialization not working at all for shapes with
subqueries
- Fixed startup process not restarting materializers for shapes
- Made startup tiered with regards to shape dependency topology

Closes #3261

Author: icehaunter

.changeset/unlucky-doors-begin.md

@@ -0,0 +1,5 @@
+---
+'@core/sync-service': patch
+---
+
+fix: ensure shapes with subqueries are deserialized correctly when loading from disk, and materializers are properly started

packages/sync-service/lib/electric/replication/eval/expr.ex

@@ -57,14 +57,19 @@ defmodule Electric.Replication.Eval.Expr do
 
   @doc false
   @spec from_json_safe(map()) :: {:ok, t()} | {:error, String.t()}
-  def from_json_safe(%{"version" => 1, "query" => query, "used_refs" => refs}) do
+  def from_json_safe(map, sublink_queries \\ %{})
+
+  def from_json_safe(
+        %{"version" => 1, "query" => query, "used_refs" => refs},
+        sublink_queries
+      ) do
     refs =
       Map.new(refs, fn [k, v] -> {k, type_from_json_safe(v)} end)
 
-    Parser.parse_and_validate_expression(query, refs: refs)
+    Parser.parse_and_validate_expression(query, refs: refs, sublink_queries: sublink_queries)
   end
 
-  def from_json_safe(_),
+  def from_json_safe(_, _),
     do: {:error, "Incorrect serialized format: keys must be `version`, `query`, `used_refs`"}
 
   defp type_from_json_safe(["array", type]), do: {:array, type_from_json_safe(type)}

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

@@ -260,42 +260,43 @@ defmodule Electric.ShapeCache do
     {:reply, !is_nil(ShapeStatus.get_existing_shape(state.stack_id, shape_handle)), state}
   end
 
-  defp shape_handles(state) do
-    ShapeStatus.list_shapes(state.stack_id)
-  end
-
   # Timeout is per-shape, not for the entire function
   defp recover_shapes(state, timeout) do
-    all_handles = shape_handles(state)
+    # we're starting here group after group without parallelization between groups
+    all_handles_and_shapes = ShapeStatus.list_shapes(state.stack_id)
 
     recovered =
-      Task.Supervisor.async_stream_nolink(
-        Electric.ProcessRegistry.name(state.stack_id, Electric.StackTaskSupervisor),
-        all_handles,
-        fn {shape_handle, shape} -> start_and_recover_shape(shape_handle, shape, state) end,
-        ordered: false,
-        timeout: timeout,
-        on_timeout: :kill_task,
-        zip_input_on_exit: true
-      )
-      |> Stream.flat_map(fn
-        {:ok, result} ->
-          result
-
-        # All other exit reasons are caught in the `start_and_recover_shape/3` function
-        {:exit, {{shape_handle, shape}, :timeout}} ->
-          Logger.error(
-            "shape #{inspect(shape)} (#{inspect(shape_handle)}) failed to start within #{timeout}ms"
-          )
+      all_handles_and_shapes
+      |> group_into_layers()
+      |> Enum.flat_map(fn shape_group ->
+        Task.Supervisor.async_stream_nolink(
+          Electric.ProcessRegistry.name(state.stack_id, Electric.StackTaskSupervisor),
+          shape_group,
+          fn {shape_handle, shape} -> start_and_recover_shape(shape_handle, shape, state) end,
+          ordered: true,
+          timeout: timeout,
+          on_timeout: :kill_task,
+          zip_input_on_exit: true
+        )
+        |> Stream.flat_map(fn
+          {:ok, result} ->
+            result
+
+          # All other exit reasons are caught in the `start_and_recover_shape/3` function
+          {:exit, {{shape_handle, shape}, :timeout}} ->
+            Logger.error(
+              "shape #{inspect(shape)} (#{inspect(shape_handle)}) failed to start within #{timeout}ms"
+            )
 
-          ShapeCleaner.remove_shape(shape_handle, stack_id: state.stack_id)
+            ShapeCleaner.remove_shape(shape_handle, stack_id: state.stack_id)
 
-          []
+            []
+        end)
+        |> Enum.to_list()
       end)
-      |> Enum.to_list()
 
     total_recovered = length(recovered)
-    total_failed_to_recover = length(all_handles) - total_recovered
+    total_failed_to_recover = length(all_handles_and_shapes) - total_recovered
 
     {Lsn.max(recovered), total_recovered, total_failed_to_recover}
   end
@@ -331,21 +332,8 @@ defmodule Electric.ShapeCache do
     else
       shape_handles =
         shape.shape_dependencies
-        |> Enum.map(&{&1, maybe_create_shape(&1, otel_ctx, state)})
-        |> Enum.with_index(fn {inner_shape, {shape_handle, _}}, index ->
-          materialized_type =
-            shape.where.used_refs |> Map.fetch!(["$sublink", Integer.to_string(index)])
-
-          ConsumerSupervisor.start_materializer(%{
-            stack_id: state.stack_id,
-            shape_handle: shape_handle,
-            storage: state.storage,
-            columns: inner_shape.explicitly_selected_columns,
-            materialized_type: materialized_type
-          })
-
-          shape_handle
-        end)
+        |> Enum.map(&maybe_create_shape(&1, otel_ctx, state))
+        |> Enum.map(&elem(&1, 0))
 
       shape = %{shape | shape_dependencies_handles: shape_handles}
 
@@ -362,6 +350,20 @@ defmodule Electric.ShapeCache do
   end
 
   defp start_shape(shape_handle, shape, state, otel_ctx, action) do
+    Enum.zip(shape.shape_dependencies_handles, shape.shape_dependencies)
+    |> Enum.with_index(fn {shape_handle, inner_shape}, index ->
+      materialized_type =
+        shape.where.used_refs |> Map.fetch!(["$sublink", Integer.to_string(index)])
+
+      ConsumerSupervisor.start_materializer(%{
+        stack_id: state.stack_id,
+        shape_handle: shape_handle,
+        storage: state.storage,
+        columns: inner_shape.explicitly_selected_columns,
+        materialized_type: materialized_type
+      })
+    end)
+
     case Electric.Shapes.DynamicConsumerSupervisor.start_shape_consumer(
            state.consumer_supervisor,
            stack_id: state.stack_id,
@@ -388,4 +390,19 @@ defmodule Electric.ShapeCache do
         :error
     end
   end
+
+  @spec group_into_layers([{shape_handle(), Shape.t()}]) :: [[{shape_handle(), Shape.t()}], ...]
+  defp group_into_layers(handles_and_shapes, acc \\ [], visited \\ MapSet.new())
+  defp group_into_layers([], acc, _visited), do: Enum.reverse(acc)
+
+  defp group_into_layers(handles_and_shapes, acc, visited) do
+    {appendable, missing_deps} =
+      Enum.split_with(handles_and_shapes, fn {_, shape} ->
+        Enum.all?(shape.shape_dependencies_handles, &MapSet.member?(visited, &1))
+      end)
+
+    visited = MapSet.new(appendable, &elem(&1, 0)) |> MapSet.union(visited)
+
+    group_into_layers(missing_deps, [appendable | acc], visited)
+  end
 end

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

@@ -141,13 +141,30 @@ defmodule Electric.ShapeCache.ShapeStatus do
 
   @impl true
   def list_shapes(stack_ref) do
-    :ets.select(shape_meta_table(stack_ref), [
+    shape_meta_table(stack_ref)
+    |> :ets.select([
       {
         {{@shape_meta_data, :"$1"}, :"$2", :_, :_},
         [],
         [{{:"$1", :"$2"}}]
       }
     ])
+    |> topological_sort()
+  end
+
+  @spec topological_sort([{shape_handle(), Shape.t()}]) :: [{shape_handle(), Shape.t()}]
+  defp topological_sort(handles_and_shapes, acc \\ [], visited \\ MapSet.new())
+  defp topological_sort([], acc, _visited), do: Enum.reverse(acc) |> List.flatten()
+
+  defp topological_sort(handles_and_shapes, acc, visited) do
+    {appendable, missing_deps} =
+      Enum.split_with(handles_and_shapes, fn {_, shape} ->
+        Enum.all?(shape.shape_dependencies_handles, &MapSet.member?(visited, &1))
+      end)
+
+    visited = MapSet.new(appendable, &elem(&1, 0)) |> MapSet.union(visited)
+
+    topological_sort(missing_deps, [appendable | acc], visited)
   end
 
   @impl true
@@ -412,10 +429,33 @@ defmodule Electric.ShapeCache.ShapeStatus do
       :ets.insert(meta_table, meta_tuples)
       :ets.insert(last_used_table, last_used_tuples)
 
+      restore_dependency_handles(shapes, meta_table, storage)
+
       :ok
     end
   end
 
+  defp restore_dependency_handles(shapes, meta_table, storage) do
+    shapes
+    |> Enum.filter(fn {_, shape} ->
+      Shape.has_dependencies?(shape) and not Shape.dependency_handles_known?(shape)
+    end)
+    |> Enum.each(fn {handle, %Shape{shape_dependencies: deps} = shape} ->
+      handles = Enum.map(deps, &get_existing_shape(meta_table, &1))
+
+      if not Enum.any?(handles, &is_nil/1) do
+        handles = Enum.map(handles, &elem(&1, 0))
+        shape = %Shape{shape | shape_dependencies_handles: handles}
+
+        :ets.update_element(meta_table, {@shape_meta_data, handle}, {2, shape})
+      else
+        Logger.warning("Shape #{inspect(handle)} has dependencies but some are unknown")
+        remove_shape(meta_table, handle)
+        Storage.cleanup!(storage, handle)
+      end
+    end)
+  end
+
   defp store_table_backup(meta_table, backup_dir) do
     case backup_dir do
       nil ->

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

@@ -85,6 +85,7 @@ defmodule Electric.Shapes.Consumer do
         buffer: [],
         monitors: [],
         cleaned?: false,
+        terminating?: false,
         txn_offset_mapping: [],
         materializer_subscribed?: false,
         # The existing body of consumer tests made it impossible to replace this dynamic
@@ -128,27 +129,44 @@ defmodule Electric.Shapes.Consumer do
         normalized_latest_offset
       )
 
-    for shape_handle <- state.shape.shape_dependencies_handles do
-      # TODO: handle a case when materializer is down
-      Process.monitor(Materializer.whereis(state.stack_id, shape_handle),
-        tag: {:dependency_materializer_down, shape_handle}
-      )
+    all_materializers_alive? =
+      Enum.all?(state.shape.shape_dependencies_handles, fn shape_handle ->
+        name = Materializer.name(state.stack_id, shape_handle)
 
-      Materializer.subscribe(state.stack_id, shape_handle)
-    end
+        with pid when is_pid(pid) <- GenServer.whereis(name),
+             true <- Process.alive?(pid) do
+          Process.monitor(pid,
+            tag: {:dependency_materializer_down, shape_handle}
+          )
+
+          Materializer.subscribe(state.stack_id, shape_handle)
+
+          true
+        else
+          _ -> false
+        end
+      end)
 
-    ShapeLogCollector.subscribe(state.stack_id, state.shape_handle, state.shape, action)
+    if all_materializers_alive? do
+      ShapeLogCollector.subscribe(state.stack_id, state.shape_handle, state.shape, action)
 
-    Logger.debug("Writer for #{state.shape_handle} initialized")
+      Logger.debug("Writer for #{state.shape_handle} initialized")
 
-    Snapshotter.start_snapshot(state.stack_id, state.shape_handle)
+      Snapshotter.start_snapshot(state.stack_id, state.shape_handle)
 
-    {:noreply,
-     Map.merge(state, %{
-       latest_offset: normalized_latest_offset,
-       writer: writer,
-       pg_snapshot: pg_snapshot
-     }), state.hibernate_after}
+      {:noreply,
+       Map.merge(state, %{
+         latest_offset: normalized_latest_offset,
+         writer: writer,
+         pg_snapshot: pg_snapshot
+       }), state.hibernate_after}
+    else
+      Logger.warning(
+        "Materializers for dependencies of #{state.shape_handle} are not alive, invalidating shape"
+      )
+
+      {:noreply, terminate_safely(state)}
+    end
   end
 
   @impl GenServer
@@ -268,12 +286,22 @@ defmodule Electric.Shapes.Consumer do
     {:noreply, terminate_safely(state)}
   end
 
+  def handle_info({:materializer_shape_invalidated, shape_handle}, state) do
+    Logger.warning("Materializer shape invalidated for #{shape_handle}")
+    {:noreply, terminate_safely(state)}
+  end
+
   def handle_info({{:dependency_materializer_down, handle}, _ref, :process, pid, reason}, state) do
     Logger.warning(
       "Materializer down for a dependency: #{handle} (#{inspect(pid)}) (#{inspect(reason)})"
     )
 
-    {:noreply, terminate_safely(state)}
+    case {reason, state.terminating?} do
+      {_, true} -> {:noreply, state}
+      {{:shutdown, _}, false} -> {:stop, reason, state}
+      {:shutdown, false} -> {:stop, reason, state}
+      _ -> {:noreply, terminate_safely(state)}
+    end
   end
 
   # We're trapping exists so that `terminate` is called to clean up the writer,
@@ -568,7 +596,13 @@ defmodule Electric.Shapes.Consumer do
   # 1. register that we want the shape data to be cleaned up.
   # 2. request a notification when all active shape data reads are complete
   # 3. exit the process when we receive that notification
-  defp terminate_safely(state, reason \\ {:shutdown, :cleanup}) do
+  defp terminate_safely(state, reason \\ {:shutdown, :cleanup})
+
+  defp terminate_safely(%{terminating?: true} = state, _reason) do
+    state
+  end
+
+  defp terminate_safely(state, reason) do
     %{
       stack_id: stack_id,
       shape_handle: shape_handle,
@@ -579,7 +613,7 @@ defmodule Electric.Shapes.Consumer do
 
     :ok = Electric.Shapes.Monitor.notify_reader_termination(stack_id, shape_handle, reason)
 
-    notify_shape_rotation(state)
+    notify_shape_rotation(%{state | terminating?: true})
   end
 
   defp reply_to_snapshot_waiters(%{awaiting_snapshot_start: []} = state, _reply) do

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

@@ -70,6 +70,7 @@ defmodule Electric.Shapes.Consumer.Materializer do
     %{stack_id: stack_id, shape_handle: shape_handle} = opts
 
     Process.set_label({:materializer, shape_handle})
+    Process.flag(:trap_exit, true)
     metadata = [stack_id: stack_id, shape_handle: shape_handle]
     Logger.metadata(metadata)
     Electric.Telemetry.Sentry.set_tags_context(metadata)
@@ -93,6 +94,10 @@ defmodule Electric.Shapes.Consumer.Materializer do
     _ = Consumer.await_snapshot_start(state)
     Consumer.subscribe_materializer(state)
 
+    Process.monitor(GenServer.whereis(Consumer.name(state)),
+      tag: {:consumer_down, state.shape_handle}
+    )
+
     {:noreply, state, {:continue, {:read_stream, storage}}}
   end
 
@@ -171,6 +176,22 @@ defmodule Electric.Shapes.Consumer.Materializer do
     {:reply, :ok, %{state | subscribers: MapSet.put(state.subscribers, pid)}}
   end
 
+  def handle_info({:EXIT, _, reason}, state) do
+    {:stop, reason, state}
+  end
+
+  def handle_info({{:consumer_down, _}, _ref, :process, _pid, {:shutdown, :cleanup}}, state) do
+    for pid <- state.subscribers do
+      send(pid, {:materializer_shape_invalidated, state.shape_handle})
+    end
+
+    {:noreply, state}
+  end
+
+  def handle_info({{:consumer_down, _}, _ref, :process, _pid, _reason}, state) do
+    {:noreply, state}
+  end
+
   def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
     {:noreply, %{state | subscribers: MapSet.delete(state.subscribers, pid)}}
   end