sse: avoid materializing inner log.

Author: thruflo

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

@@ -561,33 +561,35 @@ defmodule Electric.Shapes.Api do
     %{
       new_changes_ref: ref,
       handle: shape_handle,
-      api: %{sse_timeout: sse_timeout}
+      api: %{sse_timeout: sse_timeout},
+      params: %{offset: since_offset}
     } = request
 
-    Logger.debug("Client #{inspect(self())} is streaming SSE for changes to #{shape_handle}")
+    Logger.debug(
+      "Client #{inspect(self())} is streaming SSE for changes to #{shape_handle} since #{inspect(since_offset)}"
+    )
 
     # Set up timer for SSE timeout
     timer_ref = Process.send_after(self(), {:sse_timeout, ref}, sse_timeout)
 
     # Stream changes as SSE events for the duration of the timer.
-    sse_event_stream = Stream.resource(
-      fn ->
-        request
-      end,
-      &next_sse_event/1,
-      fn _ ->
-        Process.cancel_timer(timer_ref)
-      end
-    )
+    sse_event_stream =
+      Stream.resource(
+        fn ->
+          {request, since_offset}
+        end,
+        &next_sse_event/1,
+        fn _ ->
+          Process.cancel_timer(timer_ref)
+        end
+      )
 
     response = %{request.response | chunked: true, body: sse_event_stream}
 
     %{response | trace_attrs: Map.put(response.trace_attrs || %{}, :ot_is_sse_response, true)}
   end
 
-  defp next_sse_event(:done), do: {:halt, :done}
-
-  defp next_sse_event(%Request{} = request) do
+  defp next_sse_event({%Request{} = request, since_offset}) do
     %{
       api: api,
       handle: shape_handle,
@@ -602,23 +604,23 @@ defmodule Electric.Shapes.Api do
           |> determine_log_chunk_offset()
           |> determine_up_to_date()
 
-        case Shapes.get_merged_log_stream(
-          updated_request.api,
-          shape_handle,
-          since: updated_request.params.offset,
-          up_to: updated_request.chunk_end_offset
-        ) do
-          {:ok, log} ->
-            up_to_date_lsn = updated_request.chunk_end_offset.tx_offset
-            up_to_date_messages = maybe_up_to_date(updated_request, up_to_date_lsn)
+        # This is usually but not always the `latest_log_offset`
+        # as per `determine_log_chunk_offset/1`.
+        end_offset = updated_request.chunk_end_offset
 
-            message_stream = Stream.concat(log, up_to_date_messages)
-            messages = Enum.to_list(encode_log(updated_request, message_stream))
+        case Shapes.get_merged_log_stream(updated_request.api, shape_handle,
+               since: since_offset,
+               up_to: end_offset
+             ) do
+          {:ok, log} ->
+            control_messages = maybe_up_to_date(updated_request, end_offset.tx_offset)
+            message_stream = Stream.concat(log, control_messages)
+            encoded_stream = encode_log(updated_request, message_stream)
 
-            {messages, updated_request}
+            {[], {:emit, encoded_stream, updated_request, end_offset}}
 
           {:error, _error} ->
-            {[], request}
+            {[], {request, since_offset}}
         end
 
       {^ref, :shape_rotation} ->
@@ -632,7 +634,29 @@ defmodule Electric.Shapes.Api do
     end
   end
 
-  defp clean_up_change_listener(%Request{handle: shape_handle} = request) when not is_nil(shape_handle) do
+  defp next_sse_event({:emit, stream, %Request{} = request, since_offset}) do
+    # Can change the number taken to adjust the grouping. Currently three
+    # because there's typically 3 elements per SSE -- the actual message
+    # and the "data: " and "\n\n" delimiters around it.
+    #
+    # The JSON encoder groups stream elements by 500. So perhaps this
+    # could be a larger number for more efficiency?
+    case Electric.Utils.take_and_drop(stream, 3) do
+      {[], []} ->
+        {[], {request, since_offset}}
+
+      {head, []} ->
+        {head, {request, since_offset}}
+
+      {head, tail} ->
+        {head, {:emit, tail, request, since_offset}}
+    end
+  end
+
+  defp next_sse_event(:done), do: {:halt, :done}
+
+  defp clean_up_change_listener(%Request{handle: shape_handle} = request)
+       when not is_nil(shape_handle) do
     %{
       api: %{
         registry: registry,

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

@@ -717,4 +717,85 @@ defmodule Electric.Utils do
       after_fun
     )
   end
+
+  @doc """
+  This function is a combination of `Enum.take/2` and `Enum.drop/2` returning
+  first `n` dropped elements and the rest of the enum as a stream.
+
+  The important difference is that the enumerable is only iterated once, and
+  only for the required `n` items. The rest of the enumerable may be iterated
+  lazily later from the returned stream.
+
+  ## Examples
+      iex> {head, tail} = take_and_drop(Stream.cycle(1..3), 4)
+      iex> head
+      [1, 2, 3, 1]
+      iex> Enum.take(tail, 7)
+      [2, 3, 1, 2, 3, 1, 2]
+
+  Cribbed from https://github.com/tallakt/stream_split
+  """
+  @spec take_and_drop(Enumerable.t(), pos_integer) :: {List.t(), Enumerable.t()}
+  def take_and_drop(enum, n) when n > 0 do
+    case apply_reduce(enum, n) do
+      {:done, {_, list}} ->
+        {:lists.reverse(list), []}
+
+      {:suspended, {_, list}, cont} ->
+        stream_split = %{continuation: cont, stream: continuation_to_stream(cont)}
+        {:lists.reverse(list), stream_split}
+
+      {:halted, {_, list}} ->
+        {:lists.reverse(list), []}
+    end
+  end
+
+  def take_and_drop(enum, 0) do
+    {[], enum}
+  end
+
+  defp apply_reduce(%{continuation: cont}, n) do
+    cont.({:cont, {n, []}})
+  end
+
+  defp apply_reduce(enum, n) do
+    Enumerable.reduce(enum, {:cont, {n, []}}, &reducer_helper/2)
+  end
+
+  defp reducer_helper(item, :tail) do
+    {:suspend, item}
+  end
+
+  defp reducer_helper(item, {c, list}) when c > 1 do
+    {:cont, {c - 1, [item | list]}}
+  end
+
+  defp reducer_helper(item, {_, list}) do
+    {:suspend, {0, [item | list]}}
+  end
+
+  defp continuation_to_stream(cont) do
+    wrapped = fn {_, _, acc_cont} ->
+      case acc_cont.({:cont, :tail}) do
+        acc = {:suspended, item, _cont} ->
+          {[item], acc}
+
+        {:halted, acc} ->
+          {:halt, acc}
+
+        {:done, acc} ->
+          {:halt, acc}
+      end
+    end
+
+    cleanup = fn
+      {:suspended, _, acc_cont} ->
+        acc_cont.({:halt, nil})
+
+      _ ->
+        nil
+    end
+
+    Stream.resource(fn -> {:suspended, nil, cont} end, wrapped, cleanup)
+  end
 end