Share portsend bandwidth with a 64KB quantum

.github/workflows/ci.yml

@@ -71,5 +71,7 @@ jobs:
         # Wiring all of that up is tracked separately.
         mix test.turn
         mix test.nat
+        # Port multiplexer unit tests (no anvil, no WireGuard).
+        DIODE_MINIMAL_TEST=1 mix test --no-start test/network/mux_test.exs
         # Snapcraft recipe contract (unpinned libc6 / USN rebuilds); no anvil.
         DIODE_MINIMAL_TEST=1 mix test --no-start test/snap/snapcraft_test.exs

lib/network/mux.ex

@@ -0,0 +1,138 @@
+defmodule Network.Mux do
+  @moduledoc """
+  Round-robin queue of outbound frames for one edge socket.
+
+  Frames are grouped by partition (`{:port, ref}` for `portsend`, plus
+  the RPC partition). `pop/1` serves the current partition until it has
+  sent `quantum/0` bytes (64KB), then rotates. Weight is the payload
+  byte size. `coalesce/2` concatenates pops up to 64KB for one
+  `:ssl.send/2`.
+
+  The client scheduler is `DiodeClient.Mux`.
+  """
+
+  @coalesce_limit 64_000
+  @quantum 64_000
+
+  defstruct partitions: %{}, ring: [], turn: 0
+
+  def new, do: %__MODULE__{}
+
+  def coalesce_limit, do: @coalesce_limit
+
+  def quantum, do: @quantum
+
+  @doc "Queue `frame` on `partition`, remembering `meta` for the matching pop."
+  def enqueue(mux = %__MODULE__{}, partition, frame, meta) when is_binary(frame) do
+    partitions = mux.partitions
+    ring = if Map.has_key?(partitions, partition), do: mux.ring, else: mux.ring ++ [partition]
+
+    slot =
+      case Map.get(partitions, partition) do
+        nil -> {:queue.new(), :queue.new()}
+        existing -> existing
+      end
+
+    {frames, metas} = slot
+    slot = {:queue.in(frame, frames), :queue.in(meta, metas)}
+    %{mux | partitions: Map.put(partitions, partition, slot), ring: ring}
+  end
+
+  @doc """
+  Pop one frame from the partition whose turn it is.
+
+  Returns `{:empty, mux}` or `{:ok, mux, frame, meta}`.
+  """
+  def pop(mux = %__MODULE__{ring: []}), do: {:empty, %{mux | turn: 0}}
+
+  def pop(mux = %__MODULE__{ring: [partition | rest], turn: turn}) do
+    cond do
+      empty_slot?(mux, partition) ->
+        pop(drop_partition(mux, partition))
+
+      turn >= @quantum ->
+        pop(%{mux | ring: rest ++ [partition], turn: 0})
+
+      true ->
+        take_frame(mux, partition, rest)
+    end
+  end
+
+  @doc """
+  Pop and concatenate until `limit` bytes or the queue is empty.
+
+  Returns `{mux, binary, metas}`.
+  """
+  def coalesce(mux, limit \\ @coalesce_limit) do
+    do_coalesce(mux, limit, [], [])
+  end
+
+  def queued_frames(%__MODULE__{partitions: partitions}, partition) do
+    case Map.get(partitions, partition) do
+      nil -> []
+      {frames, _metas} -> :queue.to_list(frames)
+    end
+  end
+
+  def partition_count(%__MODULE__{partitions: partitions}), do: map_size(partitions)
+
+  defp take_frame(mux, partition, rest) do
+    {frames, metas} = Map.fetch!(mux.partitions, partition)
+    {:value, next} = :queue.peek(frames)
+
+    if mux.turn > 0 and mux.turn + byte_size(next) > @quantum do
+      pop(%{mux | ring: rest ++ [partition], turn: 0})
+    else
+      emit(mux, partition, rest, frames, metas)
+    end
+  end
+
+  defp emit(mux, partition, rest, frames, metas) do
+    {{:value, frame}, frames} = :queue.out(frames)
+    {{:value, meta}, metas} = :queue.out(metas)
+    partitions = Map.put(mux.partitions, partition, {frames, metas})
+    turn = mux.turn + byte_size(frame)
+    mux = %{mux | partitions: partitions, turn: turn}
+
+    mux =
+      cond do
+        empty_slot?(mux, partition) ->
+          drop_partition(%{mux | ring: [partition | rest]}, partition)
+
+        turn >= @quantum ->
+          %{mux | ring: rest ++ [partition], turn: 0}
+
+        true ->
+          mux
+      end
+
+    {:ok, mux, frame, meta}
+  end
+
+  defp drop_partition(mux, partition) do
+    %{
+      mux
+      | partitions: Map.delete(mux.partitions, partition),
+        ring: Enum.reject(mux.ring, &(&1 == partition)),
+        turn: 0
+    }
+  end
+
+  defp empty_slot?(mux, partition) do
+    case Map.get(mux.partitions, partition) do
+      nil -> true
+      {frames, _metas} -> :queue.is_empty(frames)
+    end
+  end
+
+  defp do_coalesce(mux, limit, frames, metas) do
+    size = IO.iodata_length(frames)
+
+    if partition_count(mux) > 0 and size < limit do
+      {:ok, mux, frame, meta} = pop(mux)
+      do_coalesce(mux, limit, [frames, frame], [meta | metas])
+    else
+      {mux, IO.iodata_to_binary(frames), Enum.reverse(metas)}
+    end
+  end
+end

lib/network/sender.ex

@@ -1,12 +1,14 @@
 defmodule Network.Sender do
-  @doc """
-    Quality-Of-Service aware network sender. The idea is to wrap
-    around a real socket and send on multiple "partitions". Small partitions
-    with fewer data in the queue are always preferred over the large partitions
+  @moduledoc """
+  Quality-Of-Service aware network sender. Wraps a socket and sends on
+  multiple partitions, rotating after each 64KB quantum.
+
+  Partition ordering lives in `Network.Mux`.
   """
   use GenServer
+  alias Network.Mux
   alias Network.Sender
-  defstruct [:partitions, :waiting, :relay]
+  defstruct [:mux, :waiting, :relay]
 
   def new(socket) do
     {:ok, pid} = GenServer.start_link(__MODULE__, [socket], hibernate_after: 5_000)
@@ -34,86 +36,38 @@ defmodule Network.Sender do
   end
 
   @impl true
-  def handle_call(
-        {:push_async, partition, data, trace},
-        _from,
-        state = %Sender{partitions: partitions, waiting: nil}
-      ) do
-    ps =
-      Map.update(partitions, partition, {[data], [{:trace, trace}]}, fn {q, rest} ->
-        {q ++ [data], rest ++ [{:trace, trace}]}
-      end)
-
-    {:reply, :ok, %Sender{state | partitions: ps}}
+  def handle_call({:push_async, partition, data, trace}, _from, state = %Sender{waiting: nil}) do
+    {:reply, :ok, enqueue(state, partition, data, {:trace, trace})}
   end
 
   @impl true
-  def handle_call(
-        {:push_async, _partition, data, trace},
-        _from,
-        state = %Sender{waiting: from}
-      ) do
-    GenServer.reply(from, data)
-    Network.EdgeV2.trace(trace)
-    {:reply, :ok, %Sender{state | waiting: nil}}
+  def handle_call({:push_async, partition, data, trace}, _from, state = %Sender{waiting: from})
+      when from != nil do
+    state = enqueue(state, partition, data, {:trace, trace})
+    deliver_waiting(state)
   end
 
   @impl true
-  def handle_call(
-        {:push, partition, data},
-        from,
-        state = %Sender{partitions: partitions, waiting: nil}
-      ) do
-    ps =
-      Map.update(partitions, partition, {[data], [from]}, fn {q, rest} ->
-        {q ++ [data], rest ++ [from]}
-      end)
-
-    {:noreply, %Sender{state | partitions: ps}}
+  def handle_call({:push, partition, data}, from, state = %Sender{waiting: nil}) do
+    {:noreply, enqueue(state, partition, data, from)}
   end
 
   @impl true
-  def handle_call(
-        {:push, _partition, data},
-        _from,
-        state = %Sender{waiting: from}
-      ) do
-    GenServer.reply(from, data)
-    {:reply, :ok, %Sender{state | waiting: nil}}
+  def handle_call({:push, partition, data}, from_push, state = %Sender{waiting: from})
+      when from != nil do
+    state = enqueue(state, partition, data, from_push)
+    deliver_waiting(state)
   end
 
   @impl true
-  def handle_call(:pop, _from, state = %Sender{partitions: partitions}) do
-    min =
-      Enum.min(
-        partitions,
-        fn {_ka, {data_a, _from_a}}, {_kb, {data_b, _from_b}} ->
-          :erts_debug.flat_size(data_a) < :erts_debug.flat_size(data_b)
-        end,
-        fn -> nil end
-      )
-
-    case min do
-      nil ->
-        {:reply, nil, state}
-
-      {partition, {[data | q], [from | rest]}} ->
-        ps =
-          if q == [] do
-            Map.delete(partitions, partition)
-          else
-            Map.put(partitions, partition, {q, rest})
-          end
-
-        state = %Sender{state | partitions: ps}
-
-        case from do
-          nil -> :nop
-          {:trace, trace} -> Network.EdgeV2.trace(trace)
-          from -> GenServer.reply(from, :ok)
-        end
-
-        {:reply, data, state}
+  def handle_call(:pop, _from, state = %Sender{}) do
+    case Mux.pop(state.mux) do
+      {:empty, mux} ->
+        {:reply, nil, %Sender{state | mux: mux}}
+
+      {:ok, mux, data, meta} ->
+        ack(meta)
+        {:reply, data, %Sender{state | mux: mux}}
     end
   end
 
@@ -122,9 +76,29 @@ defmodule Network.Sender do
     do_await(from, state)
   end
 
+  defp enqueue(state = %Sender{}, partition, data, meta) do
+    %Sender{state | mux: Mux.enqueue(state.mux, partition, data, meta)}
+  end
+
+  defp deliver_waiting(state = %Sender{waiting: from}) do
+    case Mux.pop(state.mux) do
+      {:empty, mux} ->
+        {:reply, :ok, %Sender{state | mux: mux}}
+
+      {:ok, mux, data, meta} ->
+        ack(meta)
+        GenServer.reply(from, data)
+        {:reply, :ok, %Sender{state | mux: mux, waiting: nil}}
+    end
+  end
+
+  defp ack(nil), do: :ok
+  defp ack({:trace, trace}), do: Network.EdgeV2.trace(trace)
+  defp ack(from), do: GenServer.reply(from, :ok)
+
   # Coalescing data frames into 64kb at least when available
-  defp do_await(data \\ "", from, state = %Sender{partitions: partitions, waiting: nil}) do
-    if map_size(partitions) > 0 and byte_size(data) < 64_000 do
+  defp do_await(data \\ "", from, state = %Sender{waiting: nil}) do
+    if Mux.partition_count(state.mux) > 0 and byte_size(data) < Mux.coalesce_limit() do
       {:reply, new_data, state} = handle_call(:pop, from, state)
       do_await(data <> new_data, from, state)
     else
@@ -149,7 +123,7 @@ defmodule Network.Sender do
   def init([socket]) do
     q = self()
     relay = spawn_link(__MODULE__, :relayer_loop, [q, socket])
-    {:ok, %Sender{partitions: %{}, waiting: nil, relay: relay}}
+    {:ok, %Sender{mux: Mux.new(), waiting: nil, relay: relay}}
   end
 
   def relayer_loop(q, socket) do