better buffering

Author: icehaunter

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

@@ -27,6 +27,7 @@ defmodule Electric.ShapeCache.CrashingFileStorage do
   defdelegate hibernate(opts), to: PureFileStorage
   defdelegate compact(opts, keep_complete_chunks), to: PureFileStorage
   defdelegate append_move_in_snapshot_to_log!(name, writer_state), to: PureFileStorage
+  defdelegate append_move_in_snapshot_to_log_filtered!(name, writer_state, touch_tracker, snapshot), to: PureFileStorage
   defdelegate append_control_message!(control_message, writer_state), to: PureFileStorage
   defdelegate write_move_in_snapshot!(stream, name, opts), to: PureFileStorage
 

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

@@ -362,6 +362,49 @@ defmodule Electric.ShapeCache.InMemoryStorage do
     {{initial_offset, resulting_offset}, opts}
   end
 
+  @impl Electric.ShapeCache.Storage
+  def append_move_in_snapshot_to_log_filtered!(
+        name,
+        %MS{log_table: log_table} = opts,
+        touch_tracker,
+        snapshot
+      ) do
+    initial_offset = current_offset(opts)
+    ref = make_ref()
+
+    Stream.unfold(initial_offset, fn offset ->
+      case :ets.next_lookup(log_table, {:movein, {name, nil}}) do
+        {{:movein, {^name, _}}, [{_, {key, _} = item}]} ->
+          # Check if this row should be skipped
+          if Electric.Shapes.Consumer.MoveIns.should_skip_query_row?(
+               touch_tracker,
+               snapshot,
+               key
+             ) do
+            # Skip this row - don't increment offset
+            {[], offset}
+          else
+            offset = LogOffset.increment(offset)
+            {{{:offset, storage_offset(offset)}, item}, offset}
+          end
+
+        _ ->
+          send(self(), {ref, offset})
+          nil
+      end
+    end)
+    |> Stream.reject(&(&1 == []))
+    |> Stream.chunk_every(500)
+    |> Stream.each(&:ets.insert(log_table, &1))
+    |> Stream.run()
+
+    :ets.match_delete(log_table, {{:movein, {name, :_}}, :_})
+
+    resulting_offset = receive(do: ({^ref, offset} -> offset))
+
+    {{initial_offset, resulting_offset}, opts}
+  end
+
   @impl Electric.ShapeCache.Storage
   def cleanup!(%MS{} = opts) do
     for table <- tables(opts),

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

@@ -909,6 +909,60 @@ defmodule Electric.ShapeCache.PureFileStorage do
     {inserted_range, state}
   end
 
+  def append_move_in_snapshot_to_log_filtered!(
+        name,
+        writer_state(opts: opts, writer_acc: acc) = state,
+        touch_tracker,
+        snapshot
+      ) do
+    starting_offset = WriteLoop.last_seen_offset(acc)
+
+    writer_state(writer_acc: acc) =
+      state =
+      Stream.resource(
+        fn ->
+          {File.open!(shape_snapshot_path(opts, name), [:read, :raw, :read_ahead]),
+           LogOffset.increment(starting_offset)}
+        end,
+        fn {file, offset} ->
+          with {:meta, <<key_size::32, json_size::64, op_type::8>>} <-
+                 {:meta, IO.binread(file, 13)},
+               <<key::binary-size(key_size)>> <- IO.binread(file, key_size),
+               <<json::binary-size(json_size)>> <- IO.binread(file, json_size) do
+            # Check if this row should be skipped
+            if Electric.Shapes.Consumer.MoveIns.should_skip_query_row?(
+                 touch_tracker,
+                 snapshot,
+                 key
+               ) do
+              # Skip this row - don't increment offset
+              {[], {file, offset}}
+            else
+              # Include this row
+              {[{offset, key_size, key, op_type, 0, json_size, json}],
+               {file, LogOffset.increment(offset)}}
+            end
+          else
+            {:meta, :eof} ->
+              {:halt, {file, offset}}
+
+            _ ->
+              raise Storage.Error,
+                message: "Incomplete move-in snapshot file at #{shape_snapshot_path(opts, name)}"
+          end
+        end,
+        fn {file, _} ->
+          File.close(file)
+          FileInfo.delete(shape_snapshot_path(opts, name))
+        end
+      )
+      |> append_to_log!(state)
+
+    inserted_range = {starting_offset, WriteLoop.last_seen_offset(acc)}
+
+    {inserted_range, state}
+  end
+
   def append_control_message!(control_message, writer_state(writer_acc: acc) = state)
       when is_binary(control_message) do
     offset = WriteLoop.last_seen_offset(acc)

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

@@ -108,6 +108,22 @@ defmodule Electric.ShapeCache.Storage do
   @callback append_move_in_snapshot_to_log!(name :: String.t(), writer_state()) ::
               {inserted_range :: {LogOffset.t(), LogOffset.t()}, writer_state()} | no_return()
 
+  @doc """
+  Splice a move in snapshot into the main log with filtering.
+
+  Rows are filtered using the touch_tracker: if a row's key has been touched by a transaction
+  that is NOT visible in the snapshot, skip that row (stream has fresher data).
+
+  Returns the inserted range (excluding skipped rows) and updated writer state.
+  """
+  @callback append_move_in_snapshot_to_log_filtered!(
+              name :: String.t(),
+              writer_state(),
+              touch_tracker :: %{String.t() => pos_integer()},
+              snapshot :: {pos_integer(), pos_integer(), [pos_integer()]}
+            ) ::
+              {inserted_range :: {LogOffset.t(), LogOffset.t()}, writer_state()} | no_return()
+
   @doc """
   Append a control message to the log that doesn't have an offset associated with it.
 
@@ -290,6 +306,14 @@ defmodule Electric.ShapeCache.Storage do
     {inserted_range, {mod, new_writer_state}}
   end
 
+  @impl __MODULE__
+  def append_move_in_snapshot_to_log_filtered!(name, {mod, writer_state}, touch_tracker, snapshot) do
+    {inserted_range, new_writer_state} =
+      mod.append_move_in_snapshot_to_log_filtered!(name, writer_state, touch_tracker, snapshot)
+
+    {inserted_range, {mod, new_writer_state}}
+  end
+
   @impl __MODULE__
   def append_control_message!(control_message, state)
       when is_map(control_message) do

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

@@ -313,15 +313,31 @@ defmodule Electric.Shapes.Consumer do
     end
   end
 
-  def handle_info({:query_move_in_complete, name, key_set}, state) do
+  def handle_info({:pg_snapshot_known, name, snapshot}, state) do
+    Logger.debug(fn -> "Snapshot known for move-in #{name}" end)
+
+    # Update the snapshot in waiting_move_ins
+    move_handling_state = MoveIns.set_snapshot(state.move_handling_state, name, snapshot)
+
+    # Garbage collect touches visible in all known snapshots
+    state = %{state | move_handling_state: move_handling_state}
+    state = State.gc_touch_tracker(state)
+
+    {:noreply, state, state.hibernate_after}
+  end
+
+  def handle_info({:query_move_in_complete, name, key_set, snapshot}, state) do
     Logger.debug(fn ->
       "Consumer query move in complete for #{name} with #{length(key_set)} keys"
     end)
 
-    {state, notification} = MoveHandling.query_complete(state, name, key_set)
+    {state, notification} = MoveHandling.query_complete(state, name, key_set, snapshot)
     state = notify_new_changes(state, notification)
 
-    {:noreply, state, {:continue, :consume_buffer}}
+    # Garbage collect touches after query completes (no buffer consumption needed)
+    state = State.gc_touch_tracker(state)
+
+    {:noreply, state, state.hibernate_after}
   end
 
   def handle_info({:query_move_in_error, _, error, stacktrace}, state) do
@@ -438,6 +454,7 @@ defmodule Electric.Shapes.Consumer do
 
   defp handle_txns(txns, %State{} = state), do: Enum.reduce_while(txns, state, &handle_txn/2)
 
+  # Keep buffering for initial snapshot
   defp handle_txn(txn, %State{buffering?: true} = state),
     do: {:cont, State.add_to_buffer(state, txn)}
 
@@ -447,27 +464,12 @@ defmodule Electric.Shapes.Consumer do
         {:cont, consider_flushed(%{state | initial_snapshot_state: initial_snapshot_state}, txn)}
 
       {:continue, new_initial_snapshot_state} ->
-        handle_txn_or_start_buffering(txn, %{
-          state
-          | initial_snapshot_state: new_initial_snapshot_state
-        })
+        handle_txn_in_span(txn, %{state | initial_snapshot_state: new_initial_snapshot_state})
     end
   end
 
-  defp handle_txn(txn, state), do: handle_txn_or_start_buffering(txn, state)
-
-  defp handle_txn_or_start_buffering(
-         txn,
-         %State{move_handling_state: move_handling_state} = state
-       ) do
-    case MoveIns.check_txn(move_handling_state, txn) do
-      {:start_buffering, new_move_handling_state} ->
-        handle_txn(txn, %{state | move_handling_state: new_move_handling_state, buffering?: true})
-
-      {:continue, new_move_handling_state} ->
-        handle_txn_in_span(txn, %{state | move_handling_state: new_move_handling_state})
-    end
-  end
+  # Remove the move-in buffering check - just process immediately
+  defp handle_txn(txn, state), do: handle_txn_in_span(txn, state)
 
   defp handle_txn_in_span(txn, %State{} = state) do
     ot_attrs =
@@ -530,6 +532,12 @@ defmodule Electric.Shapes.Consumer do
       {changes, num_changes, last_log_offset} ->
         timestamp = System.monotonic_time()
 
+        # Track touches for all filtered changes
+        state =
+          Enum.reduce(changes, state, fn change, acc ->
+            State.track_change(acc, xid, change)
+          end)
+
         {lines, total_size} = prepare_log_entries(changes, xid, shape)
         writer = ShapeCache.Storage.append_to_log!(lines, writer)
 

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

@@ -17,33 +17,36 @@ defmodule Electric.Shapes.Consumer.MoveHandling do
 
     storage = state.storage
     name = Electric.Utils.uuid4()
+    consumer_pid = self()
 
-    # We're querying and writing to storage in a separate task, but we're blocking until we know
-    # the insertion conditions to know when to start buffering the changes.
-    pg_snapshot =
-      Electric.ProcessRegistry.name(state.stack_id, Electric.StackTaskSupervisor)
-      |> PartialModes.query_move_in(
-        state.shape_handle,
-        state.shape,
-        formed_where_clause,
-        stack_id: state.stack_id,
-        results_fn: fn stream ->
-          stream
-          |> Stream.transform(
-            fn -> [] end,
-            fn [key, _] = item, acc -> {[item], [key | acc]} end,
-            fn acc -> send(self(), {:acc, acc}) end
-          )
-          |> Storage.write_move_in_snapshot!(name, storage)
+    # Start async query - don't block on snapshot
+    Electric.ProcessRegistry.name(state.stack_id, Electric.StackTaskSupervisor)
+    |> PartialModes.query_move_in_async(
+      state.shape_handle,
+      state.shape,
+      formed_where_clause,
+      stack_id: state.stack_id,
+      consumer_pid: consumer_pid,
+      results_fn: fn stream, pg_snapshot ->
+        task_pid = self()
 
-          receive(do: ({:acc, acc} -> acc))
-        end,
-        move_in_name: name
-      )
+        # Process query results
+        stream
+        |> Stream.transform(
+          fn -> [] end,
+          fn [key, _] = item, acc -> {[item], [key | acc]} end,
+          fn acc -> send(task_pid, {:acc, acc, pg_snapshot}) end
+        )
+        |> Storage.write_move_in_snapshot!(name, storage)
 
-    move_handling_state =
-      MoveIns.add_waiting(state.move_handling_state, name, pg_snapshot)
+        # Return accumulated keys and snapshot
+        receive(do: ({:acc, acc, snapshot} -> {acc, snapshot}))
+      end,
+      move_in_name: name
+    )
 
+    # Add to waiting WITHOUT blocking (snapshot will be set later via message)
+    move_handling_state = MoveIns.add_waiting(state.move_handling_state, name, nil)
     %{state | move_handling_state: move_handling_state}
   end
 
@@ -67,16 +70,26 @@ defmodule Electric.Shapes.Consumer.MoveHandling do
     {%{state | writer: writer}, {[message], upper_bound}}
   end
 
-  @spec query_complete(State.t(), MoveIns.move_in_name(), list(String.t())) ::
+  @spec query_complete(
+          State.t(),
+          MoveIns.move_in_name(),
+          list(String.t()),
+          MoveIns.pg_snapshot()
+        ) ::
           {State.t(), notification :: term()}
-  def query_complete(%State{} = state, name, key_set) do
-    # 1. Splice the stored data into the main log
+  def query_complete(%State{} = state, name, key_set, snapshot) do
+    # 1. Splice stored snapshot into main log with filtering
     {{_, upper_bound} = bounds, writer} =
-      Storage.append_move_in_snapshot_to_log!(name, state.writer)
+      Storage.append_move_in_snapshot_to_log_filtered!(
+        name,
+        state.writer,
+        state.touch_tracker,
+        snapshot
+      )
 
-    # 2. Remove this entry from "waiting" and add to "filtering"
+    # 2. Move from "waiting" to "filtering"
     move_handling_state =
-      MoveIns.change_to_filtering(state.move_handling_state, name, key_set)
+      MoveIns.change_to_filtering(state.move_handling_state, name, MapSet.new(key_set))
 
     state = %{state | move_handling_state: move_handling_state, writer: writer}