electric: add experimental support for streaming SSE in live mode.

Author: thruflo

packages/sync-service/lib/electric/plug/serve_shape_plug.ex

@@ -31,6 +31,7 @@ defmodule Electric.Plug.ServeShapePlug do
     all_params =
       Map.merge(conn.query_params, conn.path_params)
       |> Map.update("live", "false", &(&1 != "false"))
+      |> Map.update("experimental_live_sse", "false", &(&1 != "false"))
 
     case Api.validate(api, all_params) do
       {:ok, request} ->

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

@@ -56,6 +56,7 @@ defmodule Electric.Shapes.Api do
     stale_age: 300,
     send_cache_headers?: true,
     encoder: Electric.Shapes.Api.Encoder.JSON,
+    sse_encoder: Electric.Shapes.Api.Encoder.SSE,
     configured: false
   ]
 
@@ -484,7 +485,7 @@ defmodule Electric.Shapes.Api do
         if live? && Enum.take(log, 1) == [] do
           request
           |> update_attrs(%{ot_is_immediate_response: false})
-          |> hold_until_change()
+          |> handle_live_request()
         else
           up_to_date_lsn =
             if live? do
@@ -497,9 +498,9 @@ defmodule Electric.Shapes.Api do
               max(global_last_seen_lsn, chunk_end_offset.tx_offset)
             end
 
-          body = Stream.concat([log, maybe_up_to_date(request, up_to_date_lsn)])
+          log_stream = Stream.concat(log, maybe_up_to_date(request, up_to_date_lsn))
 
-          %{response | chunked: true, body: encode_log(request, body)}
+          %{response | chunked: true, body: encode_log(request, log_stream)}
         end
 
       {:error, error} ->
@@ -513,6 +514,13 @@ defmodule Electric.Shapes.Api do
     end
   end
 
+  defp handle_live_request(%Request{params: %{experimental_live_sse: true}} = request) do
+    stream_sse_events(request)
+  end
+  defp handle_live_request(%Request{} = request) do
+    hold_until_change(request)
+  end
+
   defp hold_until_change(%Request{} = request) do
     %{
       new_changes_ref: ref,
@@ -549,10 +557,107 @@ defmodule Electric.Shapes.Api do
     end
   end
 
-  defp clean_up_change_listener(%Request{handle: shape_handle} = request)
-       when not is_nil(shape_handle) do
-    %{api: %{registry: registry}} = request
-    Registry.unregister(registry, shape_handle)
+  defp stream_sse_events(%Request{} = request) do
+    %{
+      new_changes_ref: ref,
+      handle: shape_handle,
+      api: %{sse_timeout: sse_timeout}
+    } = request
+
+    Logger.debug("Client #{inspect(self())} is streaming SSE for changes to #{shape_handle}")
+
+    # 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
+    )
+
+    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
+    %{
+      api: api,
+      handle: shape_handle,
+      new_changes_ref: ref
+    } = request
+
+    receive do
+      {^ref, :new_changes, latest_log_offset} ->
+        updated_request =
+          %{request | last_offset: latest_log_offset}
+          |> determine_global_last_seen_lsn()
+          |> 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)
+
+            message_stream = Stream.concat(log, up_to_date_messages)
+            messages = Enum.to_list(encode_log(updated_request, message_stream))
+
+            {messages, updated_request}
+
+          {:error, _error} ->
+            {[], request}
+        end
+
+      {^ref, :shape_rotation} ->
+        must_refetch = %{headers: %{control: "must-refetch"}}
+        message = encode_message(api, must_refetch)
+
+        {message, :done}
+
+      {:sse_timeout, ^ref} ->
+        {[], :done}
+    end
+  end
+
+  defp clean_up_change_listener(%Request{handle: shape_handle} = request) when not is_nil(shape_handle) do
+    %{
+      api: %{
+        registry: registry,
+        sse_timeout: sse_timeout
+      },
+      params: %{
+        live: live?,
+        experimental_live_sse: live_sse?
+      }
+    } = request
+
+    # When handling SSE requests, the response body is a stream that listens for
+    # :new_changes events. If we unregister the shape_handle event listener immediately,
+    # we don't receive the events. So, in this case, we unregister the shape_handle
+    # listener after the sse_timeout, when we can be sure that the request is over.
+    if live? and live_sse? do
+      spawn(fn ->
+        :timer.sleep(sse_timeout)
+
+        Registry.unregister(registry, shape_handle)
+      end)
+    else
+      Registry.unregister(registry, shape_handle)
+    end
+
     request
   end
 
@@ -600,6 +705,10 @@ defmodule Electric.Shapes.Api do
   def stack_id(%Api{stack_id: stack_id}), do: stack_id
   def stack_id(%{api: %{stack_id: stack_id}}), do: stack_id
 
+  defp encode_log(%Request{api: api, params: %{live: true, experimental_live_sse: true}}, stream) do
+    encode_sse(api, :log, stream)
+  end
+
   defp encode_log(%Request{api: api}, stream) do
     encode(api, :log, stream)
   end
@@ -609,6 +718,10 @@ defmodule Electric.Shapes.Api do
     encode(api, :message, message)
   end
 
+  def encode_message(%Request{api: api, params: %{live: true, experimental_live_sse: true}}, message) do
+    encode_sse(api, :message, message)
+  end
+
   def encode_message(%Request{api: api}, message) do
     encode(api, :message, message)
   end
@@ -617,6 +730,10 @@ defmodule Electric.Shapes.Api do
     apply(encoder, type, [message])
   end
 
+  defp encode_sse(%Api{sse_encoder: sse_encoder}, type, message) when type in [:message, :log] do
+    apply(sse_encoder, type, [message])
+  end
+
   def schema(%Response{
         api: %Api{inspector: inspector},
         shape_definition: %Shapes.Shape{} = shape

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

@@ -56,6 +56,36 @@ defmodule Electric.Shapes.Api.Encoder.JSON do
   end
 end
 
+defmodule Electric.Shapes.Api.Encoder.SSE do
+  @behaviour Electric.Shapes.Api.Encoder
+
+  @impl Electric.Shapes.Api.Encoder
+  def log(item_stream) do
+    # Note that, unlike the JSON log encoder, this doesn't currently use
+    # `Stream.chunk_every/1`.
+    #
+    # This is because it's only handling live events and is usually used
+    # for small updates (the point of enabling SSE mode is to avoid request
+    # overhead when consuming small changes).
+
+    item_stream
+    |> Stream.flat_map(&message/1)
+  end
+
+  @impl Electric.Shapes.Api.Encoder
+  def message(message) do
+    ["data: ", ensure_json(message), "\n\n"]
+  end
+
+  defp ensure_json(json) when is_binary(json) do
+    json
+  end
+
+  defp ensure_json(term) do
+    Jason.encode_to_iodata!(term)
+  end
+end
+
 defmodule Electric.Shapes.Api.Encoder.Term do
   @behaviour Electric.Shapes.Api.Encoder
 

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

@@ -8,6 +8,7 @@ defmodule Electric.Shapes.Api.Params do
   import Ecto.Changeset
 
   @tmp_compaction_flag :experimental_compaction
+  @tmp_sse_flag :experimental_live_sse
 
   @primary_key false
   defmodule ColumnList do
@@ -44,13 +45,13 @@ defmodule Electric.Shapes.Api.Params do
     field(:offset, :string)
     field(:handle, :string)
     field(:live, :boolean, default: false)
-    field(:sse, :boolean, default: false)
     field(:where, :string)
     field(:columns, ColumnList)
     field(:shape_definition, :string)
     field(:replica, Ecto.Enum, values: [:default, :full], default: :default)
     field(:params, {:map, :string}, default: %{})
     field(@tmp_compaction_flag, :boolean, default: false)
+    field(@tmp_sse_flag, :boolean, default: false)
   end
 
   @type t() :: %__MODULE__{}
@@ -62,7 +63,7 @@ defmodule Electric.Shapes.Api.Params do
     |> cast_offset()
     |> validate_handle_with_offset()
     |> validate_live_with_offset()
-    |> validate_sse_with_live()
+    |> validate_live_sse()
     |> cast_root_table(api)
     |> apply_action(:validate)
     |> convert_error(api)
@@ -152,15 +153,15 @@ defmodule Electric.Shapes.Api.Params do
     end
   end
 
-  def validate_sse_with_live(%Ecto.Changeset{valid?: false} = changeset), do: changeset
+  def validate_live_sse(%Ecto.Changeset{valid?: false} = changeset), do: changeset
 
-  def validate_sse_with_live(%Ecto.Changeset{} = changeset) do
+  def validate_live_sse(%Ecto.Changeset{} = changeset) do
     live = get_field(changeset, :live)
 
     if live do
       changeset
     else
-      validate_exclusion(changeset, :sse, [true], message: "can't be true unless live is also true")
+      validate_exclusion(changeset, @tmp_sse_flag, [true], message: "can't be true unless live is also true")
     end
   end
 

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

@@ -112,6 +112,7 @@ defmodule Electric.Shapes.Api.Response do
     |> put_schema_header(response)
     |> put_up_to_date_header(response)
     |> put_offset_header(response)
+    |> put_sse_headers(response)
   end
 
   defp put_location_header(conn, %__MODULE__{status: 409} = response) do
@@ -180,37 +181,32 @@ defmodule Electric.Shapes.Api.Response do
   end
 
   defp put_cache_headers(conn, %__MODULE__{api: api} = response) do
-    case response do
-      # If the offset is -1, set a 1 week max-age, 1 hour s-maxage (shared cache)
-      # and 1 month stale-while-revalidate We want private caches to cache the
-      # initial offset for a long time but for shared caches to frequently
-      # revalidate so they're serving a fairly fresh copy of the initials shape
-      # log.
-      %{params: %{offset: @before_all_offset}} ->
-        conn
-        |> put_cache_header(
-          "cache-control",
-          "public, max-age=604800, s-maxage=3600, stale-while-revalidate=2629746",
-          api
-        )
-
-      # For live requests we want short cache lifetimes and to update the live cursor
-      %{params: %{live: true}, api: api} ->
-        conn
-        |> put_cache_header(
-          "cache-control",
-          "public, max-age=5, stale-while-revalidate=5",
-          api
-        )
+    header_value =
+      case response do
+        # If the offset is -1, set a 1 week max-age, 1 hour s-maxage (shared cache)
+        # and 1 month stale-while-revalidate We want private caches to cache the
+        # initial offset for a long time but for shared caches to frequently
+        # revalidate so they're serving a fairly fresh copy of the initials shape
+        # log.
+        %{params: %{offset: @before_all_offset}} ->
+          "public, max-age=604800, s-maxage=3600, stale-while-revalidate=2629746"
+
+        # For live SSE requests we want to cache for just under the
+        # sse_timeout, in order to enable request collapsing.
+        %{params: %{live: true, experimental_live_sse: true}} ->
+          "public, max-age=#{max(1, div(api.sse_timeout, 1000) - 1)}"
+
+        # For normal live requests we want short cache lifetimes.
+        %{params: %{live: true}} ->
+          "public, max-age=5, stale-while-revalidate=5"
+
+        # Non-live requests have the default cache headers.
+        %{params: %{live: false}} ->
+          "public, max-age=#{api.max_age}, stale-while-revalidate=#{api.stale_age}"
+      end
 
-      %{params: %{live: false}, api: api} ->
-        conn
-        |> put_cache_header(
-          "cache-control",
-          "public, max-age=#{api.max_age}, stale-while-revalidate=#{api.stale_age}",
-          api
-        )
-    end
+    conn
+    |> put_cache_header("cache-control", header_value, api)
   end
 
   defp put_cache_header(conn, header, value, %{send_cache_headers?: true}) do
@@ -263,6 +259,16 @@ defmodule Electric.Shapes.Api.Response do
     Plug.Conn.put_resp_header(conn, "electric-offset", "#{offset}")
   end
 
+  defp put_sse_headers(conn, %__MODULE__{params: %{live: true, experimental_live_sse: true}}) do
+    conn
+    |> Plug.Conn.put_resp_header("content-type", "text/event-stream")
+    |> Plug.Conn.put_resp_header("connection", "keep-alive")
+  end
+
+  defp put_sse_headers(conn, _response) do
+    conn
+  end
+
   defp send_stream(%Plug.Conn{} = conn, %__MODULE__{body: stream, status: status} = response) do
     stack_id = Api.stack_id(response)
     conn = Plug.Conn.send_chunked(conn, status)