Improve table resizing (#19)

mtrudel · web-flow · commit 708c25d570c0 · 2024-12-06T08:06:11.000+01:00
diff --git a/lib/hpax.ex b/lib/hpax.ex
@@ -82,7 +82,20 @@ defmodule HPAX do
   end
 
   @doc """
-  Resizes the given table to the given size.
+  Resizes the given table to the given maximum size.
+
+  This is intended for use where the overlying protocol has signaled a change to the table's
+  maximum size, such as when an HTTP/2 `SETTINGS` frame is received.
+
+  If the indicated size is less than the table's current size, entries
+  will be evicted as needed to fit within the specified size, and the table's
+  maximum size will be decreased to the specified value. A flag will also be
+  set which will enqueue a "dynamic table size update" command to be prefixed
+  to the next block encoded with this table, per
+  [RFC9113§4.3.1](https://www.rfc-editor.org/rfc/rfc9113.html#section-4.3.1).
+
+  If the indicated size is greater than or equal to the table's current max size, no entries are evicted
+  and the table's maximum size changes to the specified value.
 
   ## Examples
 
@@ -91,7 +104,7 @@ defmodule HPAX do
 
   """
   @spec resize(table(), non_neg_integer()) :: table()
-  defdelegate resize(table, new_size), to: Table
+  defdelegate resize(table, new_max_size), to: Table
 
   @doc """
   Decodes a header block fragment (HBF) through a given table.
@@ -114,12 +127,12 @@ defmodule HPAX do
   # Dynamic resizes must occur only at the start of a block
   # https://datatracker.ietf.org/doc/html/rfc7541#section-4.2
   def decode(<<0b001::3, rest::bitstring>>, %Table{} = table) do
-    {new_size, rest} = decode_integer(rest, 5)
+    {new_max_size, rest} = decode_integer(rest, 5)
 
-    # Dynamic resizes must be less than max table size
+    # Dynamic resizes must be less than protocol max table size
     # https://datatracker.ietf.org/doc/html/rfc7541#section-6.3
-    if new_size <= table.max_table_size do
-      decode(rest, Table.resize(table, new_size))
+    if new_max_size <= table.protocol_max_table_size do
+      decode(rest, Table.dynamic_resize(table, new_max_size))
     else
       {:error, :protocol_error}
     end
@@ -149,7 +162,9 @@ defmodule HPAX do
         when header: {action, header_name(), header_value()},
              action: :store | :store_name | :no_store | :never_store
   def encode(headers, %Table{} = table) when is_list(headers) do
-    encode_headers(headers, table, _acc = [])
+    {table, pending_resizes} = Table.pop_pending_resizes(table)
+    acc = Enum.map(pending_resizes, &[<<0b001::3, Types.encode_integer(&1, 5)::bitstring>>])
+    encode_headers(headers, table, acc)
   end
 
   @doc """
diff --git a/lib/hpax/table.ex b/lib/hpax/table.ex
@@ -3,21 +3,25 @@ defmodule HPAX.Table do
 
   @enforce_keys [:max_table_size, :huffman_encoding]
   defstruct [
+    :protocol_max_table_size,
     :max_table_size,
     :huffman_encoding,
     entries: [],
     size: 0,
-    length: 0
+    length: 0,
+    pending_minimum_resize: nil
   ]
 
   @type huffman_encoding() :: :always | :never
 
   @type t() :: %__MODULE__{
+          protocol_max_table_size: non_neg_integer(),
           max_table_size: non_neg_integer(),
           huffman_encoding: huffman_encoding(),
           entries: [{binary(), binary()}],
           size: non_neg_integer(),
-          length: non_neg_integer()
+          length: non_neg_integer(),
+          pending_minimum_resize: non_neg_integer() | nil
         }
 
   @static_table [
@@ -94,10 +98,14 @@ defmodule HPAX.Table do
   http://httpwg.org/specs/rfc7541.html#maximum.table.size.
   """
   @spec new(non_neg_integer(), huffman_encoding()) :: t()
-  def new(max_table_size, huffman_encoding)
-      when is_integer(max_table_size) and max_table_size >= 0 and
+  def new(protocol_max_table_size, huffman_encoding)
+      when is_integer(protocol_max_table_size) and protocol_max_table_size >= 0 and
              huffman_encoding in [:always, :never] do
-    %__MODULE__{max_table_size: max_table_size, huffman_encoding: huffman_encoding}
+    %__MODULE__{
+      protocol_max_table_size: protocol_max_table_size,
+      max_table_size: protocol_max_table_size,
+      huffman_encoding: huffman_encoding
+    }
   end
 
   @doc """
@@ -124,7 +132,7 @@ defmodule HPAX.Table do
 
       size + entry_size > max_table_size ->
         table
-        |> resize(max_table_size - entry_size)
+        |> evict_to_size(max_table_size - entry_size)
         |> add_header(name, value, entry_size)
 
       true ->
@@ -242,15 +250,78 @@ defmodule HPAX.Table do
   end
 
   @doc """
-  Resizes the table.
+  Changes the table's protocol negotiated maximum size, possibly evicting entries as needed to satisfy.
+
+  If the indicated size is less than the table's current max size, entries
+  will be evicted as needed to fit within the specified size, and the table's
+  maximum size will be decreased to the specified value. An will also be
+  set which will enqueue a 'dynamic table size update' command to be prefixed
+  to the next block encoded with this table, per RFC9113§4.3.1.
+
+  If the indicated size is greater than or equal to the table's current max size, no entries are evicted
+  and the table's maximum size changes to the specified value.
 
-  If the existing entries do not fit in the new table size the oldest entries are evicted.
+  In all cases, the table's `:protocol_max_table_size` is updated accordingly
   """
   @spec resize(t(), non_neg_integer()) :: t()
-  def resize(%__MODULE__{entries: entries, size: size} = table, new_size) do
-    {new_entries_reversed, new_size} = evict_towards_size(Enum.reverse(entries), size, new_size)
+  def resize(%__MODULE__{max_table_size: max_table_size} = table, new_protocol_max_table_size)
+      when new_protocol_max_table_size >= max_table_size do
+    %__MODULE__{
+      table
+      | protocol_max_table_size: new_protocol_max_table_size,
+        max_table_size: new_protocol_max_table_size
+    }
+  end
+
+  def resize(%__MODULE__{} = table, new_protocol_max_table_size) do
+    pending_minimum_resize =
+      case table.pending_minimum_resize do
+        nil -> new_protocol_max_table_size
+        current -> min(current, new_protocol_max_table_size)
+      end
+
+    %__MODULE__{
+      evict_to_size(table, new_protocol_max_table_size)
+      | protocol_max_table_size: new_protocol_max_table_size,
+        max_table_size: new_protocol_max_table_size,
+        pending_minimum_resize: pending_minimum_resize
+    }
+  end
+
+  def dynamic_resize(%__MODULE__{} = table, new_max_table_size) do
+    %__MODULE__{
+      evict_to_size(table, new_max_table_size)
+      | max_table_size: new_max_table_size
+    }
+  end
+
+  @doc """
+  Returns (and clears) any pending resize events on the table which will need to be signalled to
+  the decoder via dynamic table size update messages. Intended to be called at the start of any
+  block encode to prepend such dynamic table size update(s) as needed. The value of
+  `pending_minimum_resize` indicates the smallest maximum size of this table which has not yet
+  been signalled to the decoder, and is always included in the list returned if it is set.
+  Additionally, if the current max table size is larger than this value, it is also included int
+  the list, per https://www.rfc-editor.org/rfc/rfc7541#section-4.2
+  """
+  def pop_pending_resizes(%{pending_minimum_resize: nil} = table), do: {table, []}
+
+  def pop_pending_resizes(table) do
+    pending_resizes =
+      if table.max_table_size > table.pending_minimum_resize,
+        do: [table.pending_minimum_resize, table.max_table_size],
+        else: [table.pending_minimum_resize]
+
+    {%__MODULE__{table | pending_minimum_resize: nil}, pending_resizes}
+  end
+
+  # Removes records as necessary to have the total size of entries within the table be less than
+  # or equal to the specified value. Does not change the table's max size.
+  defp evict_to_size(%__MODULE__{entries: entries, size: size} = table, new_size) do
+    {new_entries_reversed, new_size} =
+      evict_towards_size(Enum.reverse(entries), size, new_size)
 
-    %{
+    %__MODULE__{
       table
       | entries: Enum.reverse(new_entries_reversed),
         size: new_size,
diff --git a/test/hpax/table_test.exs b/test/hpax/table_test.exs
@@ -26,7 +26,7 @@ defmodule HPAX.TableTest do
     assert {:name, _} = Table.lookup_by_header(table, ":my-header", nil)
   end
 
-  test "resizing" do
+  test "LRU eviction" do
     dynamic_table_start = length(Table.__static_table__()) + 1
 
     # This fits two headers that have name and value of 4 bytes (4 + 4 + 32, twice).
@@ -42,10 +42,6 @@ defmodule HPAX.TableTest do
     assert Table.lookup_by_index(table, dynamic_table_start) == {:ok, {"cccc", "CCCC"}}
     assert Table.lookup_by_index(table, dynamic_table_start + 1) == {:ok, {"bbbb", "BBBB"}}
     assert Table.lookup_by_index(table, dynamic_table_start + 2) == :error
-
-    # If we resize so that no headers fit, all headers are removed.
-    table = Table.resize(table, 30)
-    assert Table.lookup_by_index(table, dynamic_table_start) == :error
   end
 
   describe "looking headers up by index" do
@@ -73,4 +69,36 @@ defmodule HPAX.TableTest do
       assert {:full, 62} = Table.lookup_by_header(table, name, value)
     end
   end
+
+  describe "resizing" do
+    test "increasing the protocol max table size" do
+      table = Table.new(4096, :never)
+      table = Table.add(table, "aaaa", "AAAA")
+      table = Table.resize(table, 8192)
+      assert table.size == 40
+      assert table.max_table_size == 8192
+      assert table.protocol_max_table_size == 8192
+    end
+
+    test "decreasing the protocol max table size not below the max table size" do
+      table = Table.new(4096, :never)
+      table = Table.add(table, "aaaa", "AAAA")
+      table = Table.add(table, "bbbb", "BBBB")
+      table = Table.dynamic_resize(table, 2048)
+      table = Table.resize(table, 6000)
+      assert table.size == 40
+      assert table.max_table_size == 6000
+      assert table.protocol_max_table_size == 6000
+    end
+
+    test "decreasing the protocol max table size below the max table size" do
+      table = Table.new(4096, :never)
+      table = Table.add(table, "aaaa", "AAAA")
+      table = Table.add(table, "bbbb", "BBBB")
+      table = Table.resize(table, 40)
+      assert table.size == 40
+      assert table.max_table_size == 40
+      assert table.protocol_max_table_size == 40
+    end
+  end
 end
diff --git a/test/hpax_test.exs b/test/hpax_test.exs
@@ -107,28 +107,62 @@ defmodule HPAXTest do
     end
   end
 
+  property "encode/3 prepends dynamic resizes at the start of a block" do
+    enc_table = HPAX.new(20_000)
+    # Start with a non-empty decode table
+    dec_table = HPAX.new(20_000)
+
+    # Put a record in both to prime the pump. The table sizes should match
+    {encoded, enc_table} = HPAX.encode([{:store, "bogus", "BOGUS"}], enc_table)
+    encoded = IO.iodata_to_binary(encoded)
+    assert {:ok, _decoded, dec_table} = HPAX.decode(encoded, dec_table)
+    assert dec_table.size == enc_table.size
+    assert enc_table.max_table_size == 20_000
+    assert dec_table.max_table_size == 20_000
+
+    # Encode a record after resizing the table. We expect a dynamic resize to be
+    # encoded and the for two table sizes to be identical after decoding
+    enc_table = HPAX.resize(enc_table, 0)
+    enc_table = HPAX.resize(enc_table, 1234)
+    {encoded, enc_table} = HPAX.encode([{:store, "lame", "LAME"}], enc_table)
+    encoded = IO.iodata_to_binary(encoded)
+
+    # Ensure that we see two resizes in order
+    assert <<0b001::3, rest::bitstring>> = encoded
+    assert {:ok, 0, rest} = HPAX.Types.decode_integer(rest, 5)
+    assert <<0b001::3, rest::bitstring>> = rest
+    assert {:ok, 1234, _rest} = HPAX.Types.decode_integer(rest, 5)
+
+    # Finally, ensure that the decoder makes proper sense of this encoding
+    assert {:ok, _decoded, dec_table} = HPAX.decode(encoded, dec_table)
+    assert dec_table.size == enc_table.size
+    assert enc_table.max_table_size == 1234
+    assert dec_table.max_table_size == 1234
+  end
+
   # https://datatracker.ietf.org/doc/html/rfc7541#section-4.2
   property "decode/2 accepts dynamic resizes at the start of a block" do
     enc_table = HPAX.new(20_000)
     # Start with a non-empty decode table
     dec_table = HPAX.new(20_000)
-    {encoded, _enc_table} = HPAX.encode([{:store, "bogus", "BOGUS"}], dec_table)
+
+    # Put a record in both to prime the pump. The table sizes should match
+    {encoded, enc_table} = HPAX.encode([{:store, "bogus", "BOGUS"}], enc_table)
     encoded = IO.iodata_to_binary(encoded)
     assert {:ok, _decoded, dec_table} = HPAX.decode(encoded, dec_table)
-    assert dec_table.size > 0
-
-    check all headers_to_encode <- list_of(header_with_store(), min_length: 1) do
-      assert {encoded, enc_table} = HPAX.encode(headers_to_encode, enc_table)
-      encoded = IO.iodata_to_binary(encoded)
-      assert {:ok, _decoded, new_dec_table} = HPAX.decode(encoded, dec_table)
-      assert new_dec_table.size > enc_table.size
-
-      # Now prepend a table zeroing to the beginning and ensure that we are exactly
-      # the same size as the encode table
-      encoded = <<0b001::3, 0::5>> <> encoded
-      assert {:ok, _decoded, new_dec_table} = HPAX.decode(encoded, dec_table)
-      assert new_dec_table.size == enc_table.size
-    end
+    assert dec_table.size == enc_table.size
+    assert enc_table.max_table_size == 20_000
+    assert dec_table.max_table_size == 20_000
+
+    # Encode a record but prepend a resize to it. The decode side will now be
+    # smaller since it only contains the newly added record
+    old_enc_table_size = enc_table.size
+    {encoded, _enc_table} = HPAX.encode([{:store, "lame", "LAME"}], dec_table)
+    encoded = <<0b001::3, 0::5>> <> IO.iodata_to_binary(encoded)
+    assert {:ok, _decoded, dec_table} = HPAX.decode(encoded, dec_table)
+    assert dec_table.size == enc_table.size - old_enc_table_size
+    assert enc_table.max_table_size == 20_000
+    assert dec_table.max_table_size == 0
   end
 
   # https://datatracker.ietf.org/doc/html/rfc7541#section-4.2
