diff --git a/.changeset/calm-shapes-multiplex.md b/.changeset/calm-shapes-multiplex.md new file mode 100644 index 0000000000..56c97ea968 --- /dev/null +++ b/.changeset/calm-shapes-multiplex.md @@ -0,0 +1,8 @@ +--- +"@core/sync-service": minor +--- + +Add an authenticated, active-instance-only WebSocket endpoint for multiplexing +silent live shape waits. The endpoint coalesces change subscriptions by shape, +wakes proxies without routing shape data through the socket, and reproduces +normal empty live responses when Electric's configured deadline expires. diff --git a/packages/sync-service/README.md b/packages/sync-service/README.md index 70c50cd7c2..910fa2386b 100644 --- a/packages/sync-service/README.md +++ b/packages/sync-service/README.md @@ -82,3 +82,40 @@ The Electric app will startup along with the rest of your Elixir app. Beyond the required database connection configuration there are a lot of other optional configuration parameters. See the [`Electric` docs for more information](https://hexdocs.pm/electric/Electric.html). + +## Internal live-request multiplexing + +Electric exposes an authenticated WebSocket upgrade at +`GET /v1/shape/multiplex` for trusted proxies that need to park many silent +live waits without retaining one HTTP request process per wait. This endpoint +uses the `electric.shape-multiplex.v1` WebSocket subprotocol. It is only +available on the active Electric instance; read-only instances reject the +upgrade. + +The client sends JSON text frames to add or remove logical waits: + +```json +{"type":"watch","id":"request-1","handle":"...","offset":"12_0","cursor":"123"} +{"type":"unwatch","id":"request-1"} +``` + +`cursor` is the raw value of the previous `electric-cursor` header, or `null` +when there is no previous value. Electric may acknowledge an armed wait with +`{"type":"ready","id":"request-1"}`. It then sends one terminal frame: + +```json +{"type":"wake","id":"request-1","reason":"changes"} +{"type":"wake","id":"request-1","reason":"rotation"} +{"type":"no_change","id":"request-1","response":{"status":200,"headers":{},"body":[]}} +{"type":"error","id":"request-1","code":"...","message":"...","retryable":true} +``` + +Wake frames deliberately carry no shape rows. The proxy must issue the normal +shape HTTP request after a wake. A `no_change` frame contains the lowercase +HTTP response headers and JSON body that the proxy should return to the +unchanged shape client. The server uses its configured long-poll timeout and +removes every wait after a terminal frame. + +Embedded deployments can call `Electric.Plug.ShapeMultiplexPlug` directly. An +optional `:availability_guard` zero-arity function can enforce tenant ownership; +it must return `:ok` while the socket is valid and is rechecked periodically. diff --git a/packages/sync-service/lib/electric/plug/router.ex b/packages/sync-service/lib/electric/plug/router.ex index 923d6af2d1..a5ea8a9f30 100644 --- a/packages/sync-service/lib/electric/plug/router.ex +++ b/packages/sync-service/lib/electric/plug/router.ex @@ -35,6 +35,10 @@ defmodule Electric.Plug.Router do to: PassAssignToOptsPlug, init_opts: [plug: Electric.Plug.ServeShapePlug, assign_key: :config] + get "/v1/shape/multiplex", + to: PassAssignToOptsPlug, + init_opts: [plug: Electric.Plug.ShapeMultiplexPlug, assign_key: :config] + post "/v1/shape", to: PassAssignToOptsPlug, init_opts: [plug: Electric.Plug.ServeShapePlug, assign_key: :config] @@ -55,7 +59,8 @@ defmodule Electric.Plug.Router do # OPTIONS requests should not be authenticated def authenticate(%Plug.Conn{method: "OPTIONS"} = conn, _opts), do: conn - def authenticate(%Plug.Conn{request_path: "/v1/shape"} = conn, _opts) do + def authenticate(%Plug.Conn{request_path: request_path} = conn, _opts) + when request_path in ["/v1/shape", "/v1/shape/multiplex"] do api_secret = conn.assigns.config[:secret] if is_nil(api_secret) do diff --git a/packages/sync-service/lib/electric/plug/shape_multiplex_plug.ex b/packages/sync-service/lib/electric/plug/shape_multiplex_plug.ex new file mode 100644 index 0000000000..84d90281fd --- /dev/null +++ b/packages/sync-service/lib/electric/plug/shape_multiplex_plug.ex @@ -0,0 +1,121 @@ +defmodule Electric.Plug.ShapeMultiplexPlug do + @moduledoc """ + Upgrades an authenticated request to the internal shape live-wait multiplexer. + + The standalone router exposes this Plug at `GET /v1/shape/multiplex`. Embedded + deployments may invoke it directly with the normal Electric API options plus: + + * `:availability_guard` — optional zero-arity function returning `:ok` while + this process owns the tenant and `{:error, reason}` otherwise. It is + checked before upgrade and periodically for established sockets. + * `:subprotocol` — selected WebSocket subprotocol; defaults to + `electric.shape-multiplex.v1`. + + The Plug also requires the underlying Electric stack to be the active + instance. Authentication remains the responsibility of the enclosing router. + """ + + @behaviour Plug + + alias Electric.Shapes.Api.Multiplex + alias Electric.Shapes.Api.Multiplex.Source + alias Electric.Shapes.Api.Multiplex.WebSocket + + @max_frame_size 65_536 + + @impl Plug + def init(opts), do: opts + + @impl Plug + def call(conn, opts) do + api = fetch_opt!(opts, :api) + source = Access.get(opts, :multiplex_source, Source) + source_opts = Access.get(opts, :multiplex_source_opts) + availability_guard = Access.get(opts, :availability_guard) + subprotocol = Access.get(opts, :subprotocol, Multiplex.protocol()) + + cond do + not Multiplex.available?(api, source, source_opts, availability_guard) -> + error_response( + conn, + 503, + "inactive_instance", + "Multiplexing is only available on the active Electric instance", + true, + retry_after: 1 + ) + + not offered_subprotocol?(conn, subprotocol) -> + error_response( + conn, + 400, + "unsupported_subprotocol", + "WebSocket subprotocol #{subprotocol} is required", + false + ) + + true -> + case WebSockAdapter.UpgradeValidation.validate_upgrade(conn) do + :ok -> + socket_opts = %{ + api: api, + availability_guard: availability_guard, + multiplex_source: source, + multiplex_source_opts: source_opts, + multiplex_status_check_interval: + Access.get(opts, :multiplex_status_check_interval, 5_000) + } + + conn + |> Plug.Conn.put_resp_header("sec-websocket-protocol", subprotocol) + |> WebSockAdapter.upgrade(WebSocket, socket_opts, + early_validate_upgrade: false, + max_frame_size: Access.get(opts, :multiplex_max_frame_size, @max_frame_size), + timeout: max(60_000, api.long_poll_timeout * 3) + ) + + {:error, reason} -> + error_response(conn, 400, "invalid_upgrade", reason, false) + end + end + end + + defp offered_subprotocol?(conn, subprotocol) when is_binary(subprotocol) do + conn + |> Plug.Conn.get_req_header("sec-websocket-protocol") + |> Enum.flat_map(&Plug.Conn.Utils.list/1) + |> Enum.any?(&(&1 == subprotocol)) + end + + defp offered_subprotocol?(_conn, _subprotocol), do: false + + defp error_response(conn, status, code, message, retryable, opts \\ []) do + conn = + conn + |> Plug.Conn.put_resp_content_type("application/json") + |> Plug.Conn.put_resp_header("cache-control", "no-store") + |> Plug.Conn.put_resp_header("surrogate-control", "no-store") + + conn = + case Keyword.fetch(opts, :retry_after) do + {:ok, seconds} -> + Plug.Conn.put_resp_header(conn, "retry-after", Integer.to_string(seconds)) + + :error -> + conn + end + + Plug.Conn.send_resp( + conn, + status, + Jason.encode!(%{code: code, message: message, retryable: retryable}) + ) + end + + defp fetch_opt!(opts, key) do + case Access.fetch(opts, key) do + {:ok, value} -> value + :error -> raise KeyError, key: key, term: opts + end + end +end diff --git a/packages/sync-service/lib/electric/shape_cache.ex b/packages/sync-service/lib/electric/shape_cache.ex index f406413a6a..7f84d5aaa4 100644 --- a/packages/sync-service/lib/electric/shape_cache.ex +++ b/packages/sync-service/lib/electric/shape_cache.ex @@ -585,7 +585,8 @@ defmodule Electric.ShapeCache do @spec fetch_latest_offset(stack_id(), shape_handle(), keyword()) :: {:ok, LogOffset.t()} | :error - defp fetch_latest_offset(stack_id, shape_handle, opts \\ []) do + @doc false + def fetch_latest_offset(stack_id, shape_handle, opts \\ []) do storage = Storage.for_shape(shape_handle, Storage.for_stack(stack_id, read_only?: opts[:read_only?])) diff --git a/packages/sync-service/lib/electric/shapes.ex b/packages/sync-service/lib/electric/shapes.ex index e001033a18..dddd0d88b6 100644 --- a/packages/sync-service/lib/electric/shapes.ex +++ b/packages/sync-service/lib/electric/shapes.ex @@ -123,6 +123,16 @@ defmodule Electric.Shapes do ShapeCache.has_shape?(shape_handle, stack_id) end + @doc false + @spec fetch_latest_offset(stack_id(), shape_handle()) :: {:ok, LogOffset.t()} | :error + def fetch_latest_offset(stack_id, shape_handle) do + if ShapeCache.has_shape?(shape_handle, stack_id) do + ShapeCache.fetch_latest_offset(stack_id, shape_handle) + else + :error + end + end + @doc """ Remove and clean up all data (meta data and shape log + snapshot) associated with the given shape handle diff --git a/packages/sync-service/lib/electric/shapes/api/multiplex.ex b/packages/sync-service/lib/electric/shapes/api/multiplex.ex new file mode 100644 index 0000000000..fc4b8bbeca --- /dev/null +++ b/packages/sync-service/lib/electric/shapes/api/multiplex.ex @@ -0,0 +1,124 @@ +defmodule Electric.Shapes.Api.Multiplex do + @moduledoc """ + Protocol helpers for Electric's internal live-request multiplexer. + + The WebSocket subprotocol is `electric.shape-multiplex.v1`. A client adds and + removes logical live requests with JSON text frames: + + {"type":"watch","id":"request-1","handle":"...","offset":"12_0","cursor":"123"} + {"type":"unwatch","id":"request-1"} + + `cursor` is the unwrapped value of the previous `electric-cursor` response + header. It may be `null` for a request without a previous cursor. + + An accepted watch receives an optional `ready` frame, followed by exactly one + terminal `wake` or `no_change` frame unless the client unwatches it. Wake + frames never contain shape data; the caller fetches that data through the + normal shape HTTP endpoint. A `no_change` frame contains the status, headers, + and JSON body needed to reproduce the normal empty live HTTP response. + """ + + alias Electric.Postgres.Lsn + alias Electric.Replication.LogOffset + alias Electric.Shapes.Api + alias Electric.Shapes.Api.Params + alias Electric.Shapes.Api.Response + + @protocol "electric.shape-multiplex.v1" + + @type watch_id :: binary() + @type wake_reason :: :changes | :rotation + @type server_frame :: + %{required(:type) => binary(), required(:id) => watch_id()} + | %{ + required(:type) => binary(), + required(:code) => binary(), + required(:message) => binary(), + required(:retryable) => boolean(), + optional(:id) => watch_id() + } + + @spec protocol() :: binary() + def protocol, do: @protocol + + @doc false + def available?(%Api{} = api, source, source_opts, availability_guard) do + source.active?(api, source_opts) and guard_available?(availability_guard) + rescue + _ -> false + catch + _, _ -> false + end + + @spec ready_frame(watch_id()) :: server_frame() + def ready_frame(id), do: %{type: "ready", id: id} + + @spec wake_frame(watch_id(), wake_reason()) :: server_frame() + def wake_frame(id, reason) when reason in [:changes, :rotation] do + %{type: "wake", id: id, reason: Atom.to_string(reason)} + end + + @spec error_frame(nil | watch_id(), binary(), binary(), boolean()) :: server_frame() + def error_frame(id, code, message, retryable) do + %{type: "error", code: code, message: message, retryable: retryable} + |> maybe_put_id(id) + end + + @doc """ + Builds the transport-neutral equivalent of an empty live shape response. + + Header generation deliberately goes through `Api.Response` so cache policy, + ETag format, and cursor advancement remain identical to the HTTP path. + """ + @spec no_change_frame( + watch_id(), + Api.t(), + Electric.shape_handle(), + LogOffset.t(), + binary() | nil + ) :: server_frame() + def no_change_frame(id, %Api{} = api, handle, %LogOffset{} = offset, cursor) do + global_last_seen_lsn = + case Electric.LsnTracker.get_last_processed_lsn(api.stack_id) do + nil -> offset.tx_offset + lsn -> Lsn.to_integer(lsn) + end + + body = [ + %{ + headers: %{ + control: "up-to-date", + global_last_seen_lsn: to_string(global_last_seen_lsn) + } + } + ] + + params = %Params{handle: handle, offset: offset, live: true} + + response = %Response{ + api: api, + handle: handle, + offset: offset, + params: params, + status: 200, + up_to_date: true, + no_changes: true, + body: body, + finalized?: true + } + + headers = Response.client_headers(response, %{"cursor" => cursor}) + + %{ + type: "no_change", + id: id, + response: %{status: 200, headers: headers, body: body} + } + end + + defp maybe_put_id(frame, nil), do: frame + defp maybe_put_id(frame, id), do: Map.put(frame, :id, id) + + defp guard_available?(nil), do: true + defp guard_available?(guard) when is_function(guard, 0), do: guard.() == :ok +end diff --git a/packages/sync-service/lib/electric/shapes/api/multiplex/source.ex b/packages/sync-service/lib/electric/shapes/api/multiplex/source.ex new file mode 100644 index 0000000000..b2d5748b2d --- /dev/null +++ b/packages/sync-service/lib/electric/shapes/api/multiplex/source.ex @@ -0,0 +1,39 @@ +defmodule Electric.Shapes.Api.Multiplex.Source do + @moduledoc false + + alias Electric.Shapes + alias Electric.Shapes.Api + + @callback active?(Api.t(), term()) :: boolean() + @callback lookup(Api.t(), Electric.shape_handle(), term()) :: + {:ok, Electric.Replication.LogOffset.t()} | :not_found + @callback subscribe(Api.t(), Electric.shape_handle(), reference(), term()) :: :ok + @callback unsubscribe(Api.t(), Electric.shape_handle(), term()) :: :ok + + def active?(%Api{stack_id: stack_id}, _opts) do + Electric.StatusMonitor.service_status(stack_id) == :active + end + + def lookup(%Api{stack_id: stack_id}, handle, _opts) do + with {:ok, latest_offset} <- Shapes.fetch_latest_offset(stack_id, handle) do + {:ok, latest_offset} + else + _ -> :not_found + end + end + + def subscribe(%Api{stack_id: stack_id}, handle, ref, _opts) do + ^ref = Electric.StackSupervisor.subscribe_to_shape_events(stack_id, handle, ref) + :ok + end + + def unsubscribe(%Api{stack_id: stack_id}, handle, _opts) do + registry = Electric.StackSupervisor.registry_name(stack_id) + + if GenServer.whereis(registry) != nil do + Registry.unregister(registry, handle) + end + + :ok + end +end diff --git a/packages/sync-service/lib/electric/shapes/api/multiplex/web_socket.ex b/packages/sync-service/lib/electric/shapes/api/multiplex/web_socket.ex new file mode 100644 index 0000000000..6eb0a1693b --- /dev/null +++ b/packages/sync-service/lib/electric/shapes/api/multiplex/web_socket.ex @@ -0,0 +1,530 @@ +defmodule Electric.Shapes.Api.Multiplex.WebSocket do + @moduledoc false + + @behaviour WebSock + + alias Electric.Replication.LogOffset + alias Electric.Shapes.Api.Multiplex + alias Electric.Shapes.Api.Multiplex.Source + + @status_check_interval 5_000 + @max_identifier_bytes 512 + @max_cursor_bytes 128 + + defmodule State do + @moduledoc false + + defstruct [ + :api, + :availability_guard, + :deadline_timer_at, + :deadline_timer_ref, + :deadline_timer_token, + :source_opts, + :status_timer_ref, + deadlines: :gb_trees.empty(), + handles: %{}, + refs: %{}, + source: Source, + status_check_interval: 5_000, + watches: %{} + ] + end + + @impl WebSock + def init(opts) do + state = %State{ + api: fetch_opt!(opts, :api), + availability_guard: Access.get(opts, :availability_guard), + source: Access.get(opts, :multiplex_source, Source), + source_opts: Access.get(opts, :multiplex_source_opts), + status_check_interval: + Access.get(opts, :multiplex_status_check_interval, @status_check_interval) + } + + if available?(state) do + {:ok, schedule_status_check(state)} + else + unavailable(state) + end + end + + @impl WebSock + def handle_in({payload, opcode: :text}, %State{} = state) do + if available?(state) do + case Jason.decode(payload) do + {:ok, %{"type" => "watch"} = frame} -> add_watch(frame, state) + {:ok, %{"type" => "unwatch"} = frame} -> unwatch(frame, state) + {:ok, _} -> push_error(state, nil, "invalid_frame", "Unknown frame type", false) + {:error, _} -> push_error(state, nil, "invalid_json", "Frame must be valid JSON", false) + end + else + unavailable(state) + end + end + + def handle_in({_payload, opcode: :binary}, %State{} = state) do + push_error(state, nil, "invalid_frame", "Only JSON text frames are supported", false) + end + + @impl WebSock + def handle_info( + {:multiplex_deadline, token}, + %State{deadline_timer_token: token} = state + ) do + state = %{ + state + | deadline_timer_at: nil, + deadline_timer_ref: nil, + deadline_timer_token: nil + } + + if available?(state) do + {expired, state} = pop_expired(state, monotonic_ms(), []) + + frames = + Enum.map(expired, fn watch -> + Multiplex.no_change_frame( + watch.id, + state.api, + watch.handle, + watch.offset, + watch.cursor + ) + end) + + state = sync_deadline_timer(state) + push_frames(frames, state) + else + unavailable(state) + end + end + + def handle_info({:multiplex_deadline, _stale_token}, %State{} = state), do: {:ok, state} + + def handle_info(:multiplex_check_availability, %State{} = state) do + state = %{state | status_timer_ref: nil} + + if available?(state) do + {:ok, schedule_status_check(state)} + else + unavailable(state) + end + end + + def handle_info({ref, :new_changes, %LogOffset{} = latest_offset}, %State{} = state) do + case Map.fetch(state.refs, ref) do + {:ok, handle} -> wake_changed_watches(state, handle, latest_offset) + :error -> {:ok, state} + end + end + + def handle_info({ref, :shape_rotation, _new_handle}, %State{} = state) do + wake_rotated_watches(state, ref) + end + + def handle_info({ref, :shape_rotation}, %State{} = state) do + wake_rotated_watches(state, ref) + end + + def handle_info(_message, %State{} = state), do: {:ok, state} + + @impl WebSock + def terminate(_reason, %State{} = state) do + cancel_timer(state.deadline_timer_ref) + cancel_timer(state.status_timer_ref) + + Enum.each(Map.keys(state.handles), fn handle -> + safe_unsubscribe(state, handle) + end) + + :ok + end + + defp add_watch(frame, state) do + already_subscribed? = + is_binary(frame["handle"]) and Map.has_key?(state.handles, frame["handle"]) + + with {:ok, watch} <- validate_watch(frame, state), + {:ok, latest_offset} <- lookup(state, watch.handle) do + case LogOffset.compare(watch.offset, latest_offset) do + :lt -> + push_frame(Multiplex.wake_frame(watch.id, :changes), state) + + :eq -> + arm_watch(watch, latest_offset, not already_subscribed?, state) + + :gt -> + push_frame(Multiplex.wake_frame(watch.id, :rotation), state) + end + else + {:error, id, code, message, retryable} -> + push_error(state, id, code, message, retryable) + + :not_found -> + id = valid_id_or_nil(frame["id"]) + push_error(state, id, "shape_not_found", "Shape handle does not exist", true) + end + end + + defp arm_watch(watch, initial_offset, recheck?, state) do + with {:ok, state} <- ensure_subscribed(state, watch.handle) do + watch = %{watch | deadline: monotonic_ms() + state.api.long_poll_timeout} + state = put_watch(state, watch) + + if recheck? do + recheck_armed_watch(watch, initial_offset, state) + else + state = sync_deadline_timer(state) + push_frame(Multiplex.ready_frame(watch.id), state) + end + else + {:error, _reason} -> + push_error( + state, + watch.id, + "subscription_failed", + "Unable to subscribe to shape changes", + true + ) + end + end + + # Subscribe before re-reading the head. Any change in the gap is now either + # visible here or queued as a registry event for this process. + defp recheck_armed_watch(watch, initial_offset, state) do + case lookup(state, watch.handle) do + {:ok, latest_offset} -> + cond do + LogOffset.compare(latest_offset, initial_offset) == :lt -> + state = state |> remove_watch(watch.id) |> sync_deadline_timer() + push_frame(Multiplex.wake_frame(watch.id, :rotation), state) + + LogOffset.compare(watch.offset, latest_offset) == :lt -> + state = state |> remove_watch(watch.id) |> sync_deadline_timer() + push_frame(Multiplex.wake_frame(watch.id, :changes), state) + + LogOffset.compare(watch.offset, latest_offset) == :eq -> + state = sync_deadline_timer(state) + push_frame(Multiplex.ready_frame(watch.id), state) + + true -> + state = state |> remove_watch(watch.id) |> sync_deadline_timer() + push_frame(Multiplex.wake_frame(watch.id, :rotation), state) + end + + :not_found -> + state = state |> remove_watch(watch.id) |> sync_deadline_timer() + push_frame(Multiplex.wake_frame(watch.id, :rotation), state) + end + end + + defp unwatch(%{"id" => id}, state) when is_binary(id) do + state = state |> remove_watch(id) |> sync_deadline_timer() + {:ok, state} + end + + defp unwatch(frame, state) do + push_error( + state, + valid_id_or_nil(frame["id"]), + "invalid_frame", + "Unwatch id must be a string", + false + ) + end + + defp validate_watch(frame, state) do + id = frame["id"] + handle = frame["handle"] + offset = frame["offset"] + + cond do + not valid_identifier?(id) -> + {:error, nil, "invalid_frame", "Watch id must be a non-empty string", false} + + Map.has_key?(state.watches, id) -> + {:error, id, "duplicate_id", "Watch id is already active", false} + + not valid_identifier?(handle) -> + {:error, id, "invalid_frame", "Shape handle must be a non-empty string", false} + + not Map.has_key?(frame, "cursor") -> + {:error, id, "invalid_frame", "Watch cursor is required", false} + + not valid_cursor?(frame["cursor"]) -> + {:error, id, "invalid_frame", "Watch cursor must be a string or null", false} + + not is_binary(offset) -> + {:error, id, "invalid_offset", "Watch offset must be a string", false} + + true -> + case LogOffset.from_string(offset) do + {:ok, %LogOffset{tx_offset: tx_offset, op_offset: op_offset} = parsed_offset} + when tx_offset >= 0 and + ((is_integer(op_offset) and op_offset >= 0) or + (tx_offset == 0 and op_offset == :infinity)) -> + {:ok, + %{ + id: id, + handle: handle, + offset: parsed_offset, + cursor: frame["cursor"], + deadline: nil + }} + + _ -> + {:error, id, "invalid_offset", "Watch offset is not a live shape offset", false} + end + end + end + + defp ensure_subscribed(%State{handles: handles} = state, handle) + when is_map_key(handles, handle) do + {:ok, state} + end + + defp ensure_subscribed(state, handle) do + ref = make_ref() + + case safe_subscribe(state, handle, ref) do + :ok -> + {:ok, + %{ + state + | handles: Map.put(state.handles, handle, %{ids: MapSet.new(), ref: ref}), + refs: Map.put(state.refs, ref, handle) + }} + + {:error, _reason} = error -> + error + end + end + + defp put_watch(state, watch) do + handles = + Map.update!(state.handles, watch.handle, fn subscription -> + %{subscription | ids: MapSet.put(subscription.ids, watch.id)} + end) + + %{ + state + | watches: Map.put(state.watches, watch.id, watch), + handles: handles, + deadlines: :gb_trees.insert({watch.deadline, watch.id}, true, state.deadlines) + } + end + + defp remove_watch(state, id) do + case Map.pop(state.watches, id) do + {nil, _watches} -> + state + + {watch, watches} -> + deadlines = :gb_trees.delete_any({watch.deadline, watch.id}, state.deadlines) + subscription = Map.fetch!(state.handles, watch.handle) + remaining_ids = MapSet.delete(subscription.ids, watch.id) + + if MapSet.size(remaining_ids) == 0 do + safe_unsubscribe(state, watch.handle) + + %{ + state + | watches: watches, + deadlines: deadlines, + handles: Map.delete(state.handles, watch.handle), + refs: Map.delete(state.refs, subscription.ref) + } + else + %{ + state + | watches: watches, + deadlines: deadlines, + handles: Map.put(state.handles, watch.handle, %{subscription | ids: remaining_ids}) + } + end + end + end + + defp wake_changed_watches(state, handle, latest_offset) do + ids = state.handles |> Map.fetch!(handle) |> Map.fetch!(:ids) + + ids_to_wake = + Enum.filter(ids, fn id -> + watch = Map.fetch!(state.watches, id) + LogOffset.compare(watch.offset, latest_offset) == :lt + end) + + state = + ids_to_wake + |> Enum.reduce(state, fn id, state -> remove_watch(state, id) end) + |> sync_deadline_timer() + + frames = Enum.map(ids_to_wake, &Multiplex.wake_frame(&1, :changes)) + push_frames(frames, state) + end + + defp wake_rotated_watches(state, ref) do + case Map.fetch(state.refs, ref) do + {:ok, handle} -> + ids = state.handles |> Map.fetch!(handle) |> Map.fetch!(:ids) |> Enum.to_list() + + state = + ids + |> Enum.reduce(state, fn id, state -> remove_watch(state, id) end) + |> sync_deadline_timer() + + frames = Enum.map(ids, &Multiplex.wake_frame(&1, :rotation)) + push_frames(frames, state) + + :error -> + {:ok, state} + end + end + + defp pop_expired(state, now, acc) do + if :gb_trees.is_empty(state.deadlines) do + {Enum.reverse(acc), state} + else + {{deadline, id}, _value} = :gb_trees.smallest(state.deadlines) + + if deadline <= now do + watch = Map.fetch!(state.watches, id) + state = remove_watch(state, id) + pop_expired(state, now, [watch | acc]) + else + {Enum.reverse(acc), state} + end + end + end + + defp sync_deadline_timer(state) do + next_deadline = + if :gb_trees.is_empty(state.deadlines) do + nil + else + {{deadline, _id}, _value} = :gb_trees.smallest(state.deadlines) + deadline + end + + cond do + next_deadline == state.deadline_timer_at -> + state + + is_nil(next_deadline) -> + cancel_timer(state.deadline_timer_ref) + + %{ + state + | deadline_timer_at: nil, + deadline_timer_ref: nil, + deadline_timer_token: nil + } + + true -> + cancel_timer(state.deadline_timer_ref) + token = make_ref() + delay = max(0, next_deadline - monotonic_ms()) + timer_ref = Process.send_after(self(), {:multiplex_deadline, token}, delay) + + %{ + state + | deadline_timer_at: next_deadline, + deadline_timer_ref: timer_ref, + deadline_timer_token: token + } + end + end + + defp schedule_status_check(%State{status_check_interval: interval} = state) do + ref = Process.send_after(self(), :multiplex_check_availability, interval) + %{state | status_timer_ref: ref} + end + + defp lookup(state, handle) do + state.source.lookup(state.api, handle, state.source_opts) + rescue + _ -> :not_found + catch + _, _ -> :not_found + end + + defp safe_subscribe(state, handle, ref) do + state.source.subscribe(state.api, handle, ref, state.source_opts) + rescue + error -> {:error, error} + catch + kind, reason -> {:error, {kind, reason}} + end + + defp safe_unsubscribe(state, handle) do + state.source.unsubscribe(state.api, handle, state.source_opts) + rescue + _ -> :ok + catch + _, _ -> :ok + end + + defp available?(state) do + Multiplex.available?( + state.api, + state.source, + state.source_opts, + state.availability_guard + ) + end + + defp unavailable(state) do + frame = + Multiplex.error_frame( + nil, + "inactive_instance", + "Multiplexing is only available on the active Electric instance", + true + ) + + {:stop, {:shutdown, :restart}, {1012, "inactive Electric instance"}, + [{:text, Jason.encode!(frame)}], state} + end + + defp push_error(state, id, code, message, retryable) do + push_frame(Multiplex.error_frame(id, code, message, retryable), state) + end + + defp push_frame(frame, state), do: {:push, {:text, Jason.encode!(frame)}, state} + + defp push_frames([], state), do: {:ok, state} + + defp push_frames(frames, state) do + messages = Enum.map(frames, &{:text, Jason.encode!(&1)}) + {:push, messages, state} + end + + defp valid_identifier?(value) do + is_binary(value) and byte_size(value) > 0 and byte_size(value) <= @max_identifier_bytes + end + + defp valid_id_or_nil(id), do: if(valid_identifier?(id), do: id, else: nil) + + defp valid_cursor?(nil), do: true + + defp valid_cursor?(cursor) do + is_binary(cursor) and byte_size(cursor) <= @max_cursor_bytes + end + + defp cancel_timer(nil), do: :ok + + defp cancel_timer(ref) do + Process.cancel_timer(ref) + :ok + end + + defp monotonic_ms, do: System.monotonic_time(:millisecond) + + defp fetch_opt!(opts, key) do + case Access.fetch(opts, key) do + {:ok, value} -> value + :error -> raise KeyError, key: key, term: opts + end + end +end diff --git a/packages/sync-service/lib/electric/shapes/api/response.ex b/packages/sync-service/lib/electric/shapes/api/response.ex index 6a99d06335..87e98125dd 100644 --- a/packages/sync-service/lib/electric/shapes/api/response.ex +++ b/packages/sync-service/lib/electric/shapes/api/response.ex @@ -229,6 +229,16 @@ defmodule Electric.Shapes.Api.Response do |> put_sse_headers(response) end + @doc false + @spec client_headers(t(), map()) :: %{binary() => binary()} + def client_headers(%__MODULE__{} = response, query_params \\ %{}) do + %Plug.Conn{method: "GET", query_params: query_params} + |> Plug.Conn.put_resp_content_type("application/json") + |> put_resp_headers(response) + |> Map.fetch!(:resp_headers) + |> Map.new() + end + defp put_shape_handle_header(conn, %__MODULE__{handle: nil}) do conn end diff --git a/packages/sync-service/mix.exs b/packages/sync-service/mix.exs index c06759cb94..7b67d67fec 100644 --- a/packages/sync-service/mix.exs +++ b/packages/sync-service/mix.exs @@ -106,6 +106,7 @@ defmodule Electric.MixProject do {:req, "~> 0.5"}, {:stream_split, "~> 0.1"}, {:tz, "~> 0.28"}, + {:websock_adapter, "~> 0.5.9"}, {:phoenix_live_dashboard, "~> 0.8"} ], dev_and_test_deps(), diff --git a/packages/sync-service/test/electric/shapes/api/multiplex_web_socket_test.exs b/packages/sync-service/test/electric/shapes/api/multiplex_web_socket_test.exs new file mode 100644 index 0000000000..3ff42d0fca --- /dev/null +++ b/packages/sync-service/test/electric/shapes/api/multiplex_web_socket_test.exs @@ -0,0 +1,410 @@ +defmodule Electric.Shapes.Api.Multiplex.WebSocketTest do + use ExUnit.Case, async: false + + import Plug.Conn + import Plug.Test + + alias Electric.Plug.Router + alias Electric.Plug.ShapeMultiplexPlug + alias Electric.Replication.LogOffset + alias Electric.Shapes.Api + alias Electric.Shapes.Api.Multiplex + alias Electric.Shapes.Api.Multiplex.WebSocket + + defmodule Source do + @behaviour Electric.Shapes.Api.Multiplex.Source + + @impl true + def active?(_api, agent), do: Agent.get(agent, & &1.active?) + + @impl true + def lookup(_api, handle, agent) do + Agent.get_and_update(agent, fn state -> + result = + case Map.fetch(state.heads, handle) do + {:ok, offset} -> {:ok, offset} + :error -> :not_found + end + + {result, %{state | lookups: [handle | state.lookups]}} + end) + end + + @impl true + def subscribe(_api, handle, ref, agent) do + Agent.update(agent, fn state -> + {new_head, advances} = Map.pop(state.advance_on_subscribe, handle) + heads = if new_head, do: Map.put(state.heads, handle, new_head), else: state.heads + + %{ + state + | advance_on_subscribe: advances, + heads: heads, + subscriptions: [{handle, ref} | state.subscriptions] + } + end) + + :ok + end + + @impl true + def unsubscribe(_api, handle, agent) do + Agent.update(agent, fn state -> + %{state | unsubscriptions: [handle | state.unsubscriptions]} + end) + + :ok + end + end + + setup do + stack_id = "multiplex-websocket-test" + :ok = Electric.LsnTracker.initialize(stack_id) + :ok = Electric.LsnTracker.set_last_processed_lsn(stack_id, 123) + + api = %Api{ + configured: true, + stack_id: stack_id, + long_poll_timeout: 20_000, + send_cache_headers?: true + } + + {:ok, source} = + start_supervised( + {Agent, + fn -> + %{ + active?: true, + advance_on_subscribe: %{}, + heads: %{"shape-1" => LogOffset.new(10, 0)}, + lookups: [], + subscriptions: [], + unsubscriptions: [] + } + end} + ) + + %{api: api, source: source} + end + + test "subscribe-then-head recheck wakes a watch that raced with registration", ctx do + Agent.update(ctx.source, fn state -> + %{state | advance_on_subscribe: %{"shape-1" => LogOffset.new(11, 0)}} + end) + + state = init_socket(ctx) + + assert {:push, {:text, payload}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", "40"), state) + + assert %{"type" => "wake", "id" => "request-1", "reason" => "changes"} = + Jason.decode!(payload) + + assert state.watches == %{} + assert state.handles == %{} + + source_state = Agent.get(ctx.source, & &1) + assert length(source_state.subscriptions) == 1 + assert length(source_state.lookups) == 2 + assert source_state.unsubscriptions == ["shape-1"] + end + + test "coalesces shape subscriptions and wakes every changed watch once", ctx do + state = init_socket(ctx) + + assert {:push, {:text, ready_1}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", nil), state) + + assert %{"type" => "ready", "id" => "request-1"} = Jason.decode!(ready_1) + + assert {:push, {:text, ready_2}, state} = + WebSocket.handle_in(watch_frame("request-2", "10_0", "20"), state) + + assert %{"type" => "ready", "id" => "request-2"} = Jason.decode!(ready_2) + + source_state = Agent.get(ctx.source, & &1) + assert length(source_state.subscriptions) == 1 + assert length(source_state.lookups) == 3 + + [subscription] = Map.values(state.handles) + + assert {:push, messages, state} = + WebSocket.handle_info( + {subscription.ref, :new_changes, LogOffset.new(11, 0)}, + state + ) + + frames = decode_messages(messages) + + assert MapSet.new(frames) == + MapSet.new([ + %{"type" => "wake", "id" => "request-1", "reason" => "changes"}, + %{"type" => "wake", "id" => "request-2", "reason" => "changes"} + ]) + + assert state.watches == %{} + assert state.handles == %{} + assert Agent.get(ctx.source, & &1.unsubscriptions) == ["shape-1"] + end + + test "deadline returns a normal no-change HTTP envelope and advances the cursor", ctx do + state = init_socket(ctx) + + assert {:push, {:text, _ready}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", "999999999999"), state) + + cancel_timer(state.deadline_timer_ref) + expired_at = System.monotonic_time(:millisecond) - 1 + watch = %{Map.fetch!(state.watches, "request-1") | deadline: expired_at} + token = make_ref() + + state = %{ + state + | watches: %{"request-1" => watch}, + deadlines: :gb_trees.insert({expired_at, "request-1"}, true, :gb_trees.empty()), + deadline_timer_at: expired_at, + deadline_timer_ref: nil, + deadline_timer_token: token + } + + assert {:push, messages, state} = + WebSocket.handle_info({:multiplex_deadline, token}, state) + + assert [%{"type" => "no_change"} = frame] = decode_messages(messages) + + assert %{ + "type" => "no_change", + "id" => "request-1", + "response" => %{ + "status" => 200, + "headers" => headers, + "body" => [ + %{ + "headers" => %{ + "control" => "up-to-date", + "global_last_seen_lsn" => "123" + } + } + ] + } + } = frame + + assert headers["electric-handle"] == "shape-1" + assert headers["electric-offset"] == "10_0" + assert headers["electric-has-data"] == "false" + assert headers["electric-up-to-date"] == "" + assert headers["cache-control"] == "public, max-age=5, stale-while-revalidate=5" + assert headers["content-type"] == "application/json; charset=utf-8" + assert is_binary(headers["etag"]) + assert String.to_integer(headers["electric-cursor"]) > 999_999_999_999 + assert state.watches == %{} + assert state.handles == %{} + assert length(Agent.get(ctx.source, & &1.lookups)) == 2 + end + + test "shape rotation wakes all watches without returning data", ctx do + state = init_socket(ctx) + + assert {:push, {:text, _ready}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", nil), state) + + [subscription] = Map.values(state.handles) + + assert {:push, messages, state} = + WebSocket.handle_info({subscription.ref, :shape_rotation}, state) + + assert [%{"type" => "wake", "id" => "request-1", "reason" => "rotation"}] = + decode_messages(messages) + + assert state.watches == %{} + assert state.handles == %{} + end + + test "a regressed head during registration wakes as a rotation", ctx do + Agent.update(ctx.source, fn state -> + %{state | advance_on_subscribe: %{"shape-1" => LogOffset.new(9, 0)}} + end) + + state = init_socket(ctx) + + assert {:push, {:text, payload}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", nil), state) + + assert %{"type" => "wake", "id" => "request-1", "reason" => "rotation"} = + Jason.decode!(payload) + + assert state.watches == %{} + assert state.handles == %{} + end + + test "a requested offset ahead of the observed head wakes as a rotation", ctx do + state = init_socket(ctx) + + assert {:push, {:text, payload}, state} = + WebSocket.handle_in(watch_frame("request-1", "11_0", nil), state) + + assert %{"type" => "wake", "id" => "request-1", "reason" => "rotation"} = + Jason.decode!(payload) + + assert state.watches == %{} + assert state.handles == %{} + assert Agent.get(ctx.source, & &1.subscriptions) == [] + end + + test "unwatch is idempotent and removes the underlying subscription", ctx do + state = init_socket(ctx) + + assert {:push, {:text, _ready}, state} = + WebSocket.handle_in(watch_frame("request-1", "10_0", nil), state) + + assert {:ok, state} = + WebSocket.handle_in( + {Jason.encode!(%{type: "unwatch", id: "request-1"}), opcode: :text}, + state + ) + + assert {:ok, state} = + WebSocket.handle_in( + {Jason.encode!(%{type: "unwatch", id: "request-1"}), opcode: :text}, + state + ) + + assert state.watches == %{} + assert Agent.get(ctx.source, & &1.unsubscriptions) == ["shape-1"] + end + + test "arms the canonical 0_inf quiet-shape offset", ctx do + Agent.update(ctx.source, fn state -> + %{state | heads: %{"shape-1" => LogOffset.last_before_real_offsets()}} + end) + + state = init_socket(ctx) + + assert {:push, {:text, payload}, state} = + WebSocket.handle_in(watch_frame("request-1", "0_inf", nil), state) + + assert %{"type" => "ready", "id" => "request-1"} = Jason.decode!(payload) + assert Map.has_key?(state.watches, "request-1") + end + + test "rejects other special and missing cursor watch values", ctx do + state = init_socket(ctx) + + assert {:push, {:text, payload}, state} = + WebSocket.handle_in(watch_frame("request-1", "1_inf", nil), state) + + assert %{"type" => "error", "code" => "invalid_offset", "retryable" => false} = + Jason.decode!(payload) + + frame = %{type: "watch", id: "request-2", handle: "shape-1", offset: "10_0"} + + assert {:push, {:text, payload}, _state} = + WebSocket.handle_in({Jason.encode!(frame), opcode: :text}, state) + + assert %{"type" => "error", "code" => "invalid_frame", "retryable" => false} = + Jason.decode!(payload) + end + + test "closes an established socket when active ownership is lost", ctx do + state = init_socket(ctx) + Agent.update(ctx.source, &%{&1 | active?: false}) + + assert {:stop, {:shutdown, :restart}, {1012, _reason}, [{:text, payload}], _state} = + WebSocket.handle_info(:multiplex_check_availability, state) + + assert %{"type" => "error", "code" => "inactive_instance", "retryable" => true} = + Jason.decode!(payload) + end + + describe "ShapeMultiplexPlug" do + test "upgrades with the selected subprotocol", ctx do + conn = websocket_conn() |> ShapeMultiplexPlug.call(plug_opts(ctx)) + + assert conn.state == :upgraded + assert get_resp_header(conn, "sec-websocket-protocol") == [Multiplex.protocol()] + end + + test "checks the embedding availability guard before upgrade", ctx do + opts = Keyword.put(plug_opts(ctx), :availability_guard, fn -> {:error, :not_owner} end) + conn = websocket_conn() |> ShapeMultiplexPlug.call(opts) + + assert conn.status == 503 + + assert %{"code" => "inactive_instance", "retryable" => true} = + Jason.decode!(conn.resp_body) + end + + test "requires the versioned subprotocol", ctx do + conn = + conn(:get, "/v1/shape/multiplex") + |> with_host_header() + |> put_req_header("connection", "upgrade") + |> put_req_header("upgrade", "websocket") + |> put_req_header("sec-websocket-key", Base.encode64(:crypto.strong_rand_bytes(16))) + |> put_req_header("sec-websocket-version", "13") + |> ShapeMultiplexPlug.call(plug_opts(ctx)) + + assert conn.status == 400 + assert %{"code" => "unsupported_subprotocol"} = Jason.decode!(conn.resp_body) + end + end + + test "standalone router authenticates the multiplex endpoint" do + opts = Router.init(secret: "source-secret", stack_id: "router-auth-test") + + conn = + conn(:get, "/v1/shape/multiplex") + |> Router.call(opts) + + assert conn.status == 401 + end + + defp init_socket(ctx) do + assert {:ok, state} = + WebSocket.init(%{ + api: ctx.api, + multiplex_source: Source, + multiplex_source_opts: ctx.source, + multiplex_status_check_interval: 60_000 + }) + + state + end + + defp watch_frame(id, offset, cursor) do + {Jason.encode!(%{type: "watch", id: id, handle: "shape-1", offset: offset, cursor: cursor}), + opcode: :text} + end + + defp websocket_conn do + conn(:get, "/v1/shape/multiplex") + |> with_host_header() + |> put_req_header("connection", "upgrade") + |> put_req_header("upgrade", "websocket") + |> put_req_header("sec-websocket-key", Base.encode64(:crypto.strong_rand_bytes(16))) + |> put_req_header("sec-websocket-version", "13") + |> put_req_header("sec-websocket-protocol", Multiplex.protocol()) + end + + defp plug_opts(ctx) do + [ + api: ctx.api, + availability_guard: fn -> :ok end, + multiplex_source: Source, + multiplex_source_opts: ctx.source, + subprotocol: Multiplex.protocol() + ] + end + + defp decode_messages({:text, payload}), do: [Jason.decode!(payload)] + + defp decode_messages(messages), + do: Enum.map(messages, fn {:text, payload} -> Jason.decode!(payload) end) + + defp cancel_timer(nil), do: :ok + defp cancel_timer(ref), do: Process.cancel_timer(ref) + + defp with_host_header(conn) do + %{conn | host: "example.test", req_headers: [{"host", "example.test"} | conn.req_headers]} + end +end