Implements HRW.Weighted (#2)

Using
https://www.ietf.org/archive/id/draft-ietf-bess-weighted-hrw-00.html to
implement weighted HRW.

In practice this means you can now assign a weight to your nodes.

Author: joladev

README.md

@@ -11,7 +11,9 @@ The most common library in the Elixir community to use to solve that problem is
 
 For larger node sets, build a skeleton with `HRW.build` and pass it to `HRW.owner` to get O(log n) lookups. The skeleton is plain data — build it once, reuse it across calls.
 
-Additionally, there's `HRW.Bounded` for when you want to control the distribution of keys across nodes to limit skew. Consistent hashing and rendezvous hashing algorithms can easily result in uneven distribution for smaller node counts, and `HRW.Bounded` lets you control that, assuming that you have the whole key set up front.
+`HRW.owner` and `HRW.build` support an optional `scorer` option for alternative strategies. The available options are `%HRW{}` for the default algorithm, and `%HRW.Weighted{}` for when you want certain nodes to get a larger share of keys.
+
+For additional strategies, there's `HRW.Bounded` for when you want to control the distribution of keys across nodes to limit skew. Consistent hashing and rendezvous hashing algorithms can easily result in uneven distribution for smaller node counts, and `HRW.Bounded` lets you control that, assuming that you have the whole key set up front.
 
 ```elixir
 # HRW
@@ -28,6 +30,10 @@ skeleton = HRW.build(["server1", "server2", "server3"])
 HRW.owner("192.168.0.2", skeleton)
 #=> "server3"
 
+# HRW.Weighted
+HRW.owner("192.168.0.1", [{"server1", 1}, {"server2", 1}, {"server3", 10}], scorer: %HRW.Weighted{})
+#=> "server3"
+
 # HRW.Bounded
 HRW.Bounded.assignments(["a", "b", "c", "d"], ["x", "y"], epsilon: 0.0)
 #=> %{"a" => "x", "b" => "x", "c" => "y", "d" => "y"}
@@ -39,11 +45,11 @@ tl;dr HRW performs similarly to ExHashRing on smaller node lists, but falls behi
 
 Lookup latency on Apple M4 Pro / Elixir 1.19.5 / OTP 28.5, median per call:
 
-| nodes  | HRW.owner   | HRW.owner (skeleton) | ExHashRing.find_node |
-|-------:|------------:|-------------------:|---------------------:|
-|     10 |     292 ns  |      292 ns        |        333 ns        |
-|    100 |    2.67 µs  |      875 ns        |        375 ns        |
-|  1,000 |   25.54 µs  |     1.08 µs        |        380 ns        |
-| 10,000 |  253.58 µs  |     1.38 µs        |        420 ns        |
+| nodes  | HRW.owner | HRW.owner (weighted) | HRW.owner (skeleton) | HRW.owner (skeleton weighted) | ExHashRing.find_node |
+|-------:|----------:|---------------------:|---------------------:|------------------------------:|---------------------:|
+|     10 |    542 ns |              1.00 µs |               292 ns |                        667 ns |               333 ns |
+|    100 |   6.33 µs |             11.00 µs |               920 ns |                       1.71 µs |               380 ns |
+|  1,000 |  71.79 µs |               121 µs |              1.25 µs |                       2.25 µs |               380 ns |
+| 10,000 |    771 µs |              1.25 ms |              1.50 µs |                       2.92 µs |               420 ns |
 
 Reproduce with `elixir benches/bench.exs`.

benches/bench.exs

@@ -12,12 +12,14 @@ defmodule Bench do
       nodes = Enum.map(1..n, &"node-#{&1}")
       {:ok, ring} = Ring.start_link()
       Ring.set_nodes(ring, nodes, :infinity)
-      %{nodes: nodes, skeleton: HRW.build(nodes), ring: ring}
+      %{nodes: nodes, skeleton: HRW.build(nodes), ring: ring, weighted_nodes: Enum.map(1..n, &{"node-#{&1}", &1}), weighted_skeleton: HRW.build(Enum.map(1..n, &{"node-#{&1}", &1}), scorer: %HRW.Weighted{})}
     end
 
     Benchee.run(%{
       "HRW.owner" => fn %{nodes: nodes} -> HRW.owner("test", nodes) end,
+      "HRW.owner (weighted)" => fn %{weighted_nodes: weighted_nodes} -> HRW.owner("test", weighted_nodes, scorer: %HRW.Weighted{}) end,
       "HRW.owner (skeleton)" => fn %{skeleton: skeleton} -> HRW.owner("test", skeleton) end,
+      "HRW.owner (skeleton weighted)" => fn %{weighted_skeleton: weighted_skeleton} -> HRW.owner("test", weighted_skeleton) end,
       "ExHashRing.Ring.find_node" => fn %{ring: ring} -> Ring.find_node(ring, "test") end
     }, inputs: %{
       "A: 10" => setup.(10),

lib/hrw.ex

@@ -48,21 +48,26 @@ defmodule HRW do
   def owner(key, nodes_or_skeleton, opts \\ [])
 
   def owner(key, %HRW.Skeleton{} = skeleton, _opts) do
-    HRW.Skeleton.owner(key, skeleton)
+    result = HRW.Skeleton.owner(key, skeleton)
+    unwrap_node(result, skeleton.scorer)
   end
 
   def owner(key, nodes, opts) do
+    validate_nodes(nodes, Keyword.get(opts, :scorer, %HRW{}))
+
     nodes =
       nodes
       |> Enum.sort()
-      |> Enum.uniq()
+      |> Enum.dedup()
 
     if scorer = Keyword.get(opts, :scorer) do
       %mod{} = scorer
 
-      Enum.max_by(nodes, fn node ->
+      nodes
+      |> Enum.max_by(fn node ->
         mod.score(scorer, key, node)
       end)
+      |> unwrap_node(scorer)
     else
       Enum.max_by(nodes, fn node ->
         :erlang.phash2({key, node})
@@ -85,17 +90,20 @@ defmodule HRW do
   """
   @spec owners(term(), [term()], non_neg_integer(), keyword()) :: [term()]
   def owners(key, nodes, count, opts \\ []) do
+    validate_nodes(nodes, Keyword.get(opts, :scorer, %HRW{}))
+
     nodes =
       nodes
       |> Enum.sort()
-      |> Enum.uniq()
+      |> Enum.dedup()
 
     if scorer = Keyword.get(opts, :scorer) do
       %mod{} = scorer
 
       nodes
       |> Enum.sort_by(fn node -> mod.score(scorer, key, node) end, :desc)
       |> Enum.take(count)
+      |> Enum.map(&unwrap_node(&1, scorer))
     else
       nodes
       |> Enum.sort_by(fn node -> :erlang.phash2({key, node}) end, :desc)
@@ -121,6 +129,26 @@ defmodule HRW do
   """
   @spec build([term()], keyword()) :: HRW.Skeleton.t()
   def build(nodes, opts \\ []) do
+    validate_nodes(nodes, Keyword.get(opts, :scorer, %HRW{}))
     HRW.Skeleton.build(nodes, opts)
   end
+
+  defp unwrap_node({node, _weight}, %HRW.Weighted{}), do: node
+  defp unwrap_node(node, _scorer), do: node
+
+  defp validate_nodes(nodes, %HRW{}) when is_list(nodes), do: nil
+  defp validate_nodes(_nodes, %HRW{}), do: raise(ArgumentError, "nodes must be a list")
+
+  defp validate_nodes(nodes, %HRW.Weighted{}) when is_list(nodes) do
+    if not Enum.all?(nodes, fn
+         {_node, weight} when is_number(weight) and weight > 0 -> true
+         _ -> false
+       end) do
+      raise ArgumentError,
+            "HRW.Weighted requires a list of {node, weight} tuples with positive numeric weights"
+    end
+  end
+
+  defp validate_nodes(_nodes, %HRW.Weighted{}),
+    do: raise(ArgumentError, "nodes must be a list of tuples")
 end

lib/hrw/bounded.ex

@@ -46,13 +46,13 @@ defmodule HRW.Bounded do
 
     keys =
       keys
-      |> Enum.uniq()
       |> Enum.sort()
+      |> Enum.dedup()
 
     nodes =
       nodes
-      |> Enum.uniq()
       |> Enum.sort()
+      |> Enum.dedup()
 
     cap = ceil(length(keys) / length(nodes) * (1 + epsilon))
 

lib/hrw/scorer.ex

@@ -2,11 +2,11 @@ defmodule HRW.Scorer do
   @moduledoc """
   Behaviour for HRW scoring strategies. Each variant module (`HRW`, future
   `HRW.Weighted`, etc.) defines a struct holding its configuration and
-  implements `score/3` returning an integer score for a `(key, node)` pair.
+  implements `score/3` returning a score for a `(key, node)` pair.
 
   Pass an instance via the `:scorer` option to `HRW.owner/3`, `HRW.owners/4`,
   or `HRW.build/2`. The highest-scoring node wins.
   """
 
-  @callback score(scorer :: struct(), key :: term(), node :: term()) :: integer()
+  @callback score(scorer :: struct(), key :: term(), node :: term()) :: number()
 end

lib/hrw/skeleton.ex

@@ -35,6 +35,19 @@ defmodule HRW.Skeleton do
     fanout = Keyword.get_lazy(opts, :fanout, fn -> optimal_fanout(count) end)
     levels = if count > 1, do: ceil(:math.log(count) / :math.log(fanout)), else: 0
 
+    scorer =
+      case scorer do
+        %HRW.Weighted{} = weighted ->
+          %{
+            weighted
+            | branch_weights: compute_branch_weights(cluster_list, fanout, levels),
+              fanout: fanout
+          }
+
+        other ->
+          other
+      end
+
     %__MODULE__{
       clusters: clusters,
       fanout: fanout,
@@ -56,7 +69,9 @@ defmodule HRW.Skeleton do
   defp do_owner(key, %__MODULE__{scorer: %HRW{hash_fn: nil}} = skeleton, salt) do
     index =
       Enum.reduce(0..(skeleton.levels - 1), 0, fn level, acc ->
-        digit = Enum.max_by(0..(skeleton.fanout - 1), &:erlang.phash2({{key, salt, level}, &1}))
+        digit =
+          Enum.max_by(0..(skeleton.fanout - 1), &:erlang.phash2({{key, salt, level, acc}, &1}))
+
         acc * skeleton.fanout + digit
       end)
 
@@ -75,7 +90,9 @@ defmodule HRW.Skeleton do
   defp do_owner(key, %__MODULE__{scorer: %mod{} = scorer} = skeleton, salt) do
     index =
       Enum.reduce(0..(skeleton.levels - 1), 0, fn level, acc ->
-        digit = Enum.max_by(0..(skeleton.fanout - 1), &mod.score(scorer, {key, salt, level}, &1))
+        digit =
+          Enum.max_by(0..(skeleton.fanout - 1), &mod.score(scorer, {key, salt, level, acc}, &1))
+
         acc * skeleton.fanout + digit
       end)
 
@@ -91,8 +108,8 @@ defmodule HRW.Skeleton do
     # Deterministic ordering so the same node set always produces the same clusters.
     chunks =
       nodes
-      |> Enum.uniq()
       |> Enum.sort()
+      |> Enum.dedup()
       |> Enum.chunk_every(size)
 
     # An undersized last chunk would get the same routing probability as full chunks,
@@ -121,6 +138,36 @@ defmodule HRW.Skeleton do
       fanout * levels / (1 - overflow_prob)
     end)
   end
+
+  defp compute_branch_weights(_clusters, _fanout, 0), do: {}
+
+  defp compute_branch_weights(clusters, fanout, levels) do
+    cluster_weights =
+      clusters
+      |> Enum.with_index()
+      |> Enum.map(fn {nodes, index} ->
+        total =
+          Enum.sum(
+            Enum.map(nodes, fn
+              {_node, weight} -> weight
+              # plain nodes get implicit weight 1
+              _ -> 1
+            end)
+          )
+
+        {index, total}
+      end)
+
+    0..(levels - 1)
+    |> Enum.map(fn level ->
+      divisor = Integer.pow(fanout, levels - level - 1)
+
+      Enum.reduce(cluster_weights, %{}, fn {index, weight}, acc ->
+        Map.update(acc, div(index, divisor), weight, &(&1 + weight))
+      end)
+    end)
+    |> List.to_tuple()
+  end
 end
 
 defimpl Inspect, for: HRW.Skeleton do

lib/hrw/weighted.ex

@@ -0,0 +1,65 @@
+defmodule HRW.Weighted do
+  @moduledoc """
+  Weighted HRW implementation for ensuring certain nodes get a greater share of keys.
+
+  Inspired by https://www.ietf.org/archive/id/draft-ietf-bess-weighted-hrw-00.html
+
+  Works with both flat lists (`HRW.owner(key, [{"a", 1}, {"b", 4}], scorer: %HRW.Weighted{})`)
+  and skeletons (`HRW.build([...], scorer: %HRW.Weighted{})`).
+  """
+
+  @behaviour HRW.Scorer
+
+  import Bitwise
+
+  @phash2_max 0x7FFFFFF
+
+  defstruct [:hash_fn, :branch_weights, :fanout]
+
+  @type t :: %__MODULE__{
+          hash_fn: (term() -> integer()) | nil,
+          branch_weights: tuple() | nil,
+          fanout: pos_integer() | nil
+        }
+
+  @impl HRW.Scorer
+  def score(%__MODULE__{hash_fn: nil}, key, {node, weight}) do
+    hash = :erlang.phash2({key, node})
+    leaf_score(hash, weight)
+  end
+
+  def score(%__MODULE__{hash_fn: hash_fn}, key, {node, weight}) do
+    hash = hash_fn.({key, node})
+    leaf_score(hash, weight)
+  end
+
+  # Skeleton tree walk: key is {real_key, salt, level, prefix},
+  # entry is the candidate digit. Looks up branch weight from precomputed tree.
+  def score(
+        %__MODULE__{branch_weights: branch_weights, fanout: fanout},
+        {key, salt, level, prefix},
+        digit
+      ) do
+    candidate = prefix * fanout + digit
+
+    branch_weight =
+      branch_weights
+      |> elem(level)
+      |> Map.get(candidate, 0)
+
+    if branch_weight == 0 do
+      # Empty branch (no nodes under this prefix). Returning 0.0 ensures it never wins.
+      0.0
+    else
+      normalized_hash =
+        (:erlang.phash2({key, salt, level, prefix, digit}) &&& @phash2_max) / @phash2_max
+
+      :math.pow(normalized_hash, 1 / branch_weight)
+    end
+  end
+
+  defp leaf_score(hash, weight) do
+    normalized = (hash &&& @phash2_max) / @phash2_max
+    :math.pow(normalized, 1 / weight)
+  end
+end

test/hrw/skeleton_test.exs

@@ -5,6 +5,6 @@ defmodule HRW.SkeletonTest do
     nodes = Enum.map(1..12, &"server#{&1}")
     skeleton = HRW.build(nodes, cluster_size: 4)
 
-    assert HRW.owner("192.168.0.1", skeleton) == "server12"
+    assert HRW.owner("192.168.0.1", skeleton) == "server9"
   end
 end

test/hrw/weighted_test.exs

@@ -0,0 +1,63 @@
+defmodule HRWWeightedTest do
+  use ExUnit.Case, async: true
+  alias HRW.Weighted
+
+  test "higher weight nodes win more often in flat lists" do
+    # 80% of the time, the node with weight 4 wins vs weight 1
+    weights = [{"heavy", 4}, {"light", 1}]
+    keys = Enum.map(1..1_000, &"key-#{&1}")
+
+    winners =
+      keys
+      |> Enum.map(&HRW.owner(&1, weights, scorer: %Weighted{}))
+      |> Enum.frequencies()
+
+    assert winners["heavy"] > winners["light"]
+  end
+
+  test "weight ratio approximates win ratio at scale" do
+    keys = Enum.map(1..100, &"key-#{&1}")
+    weights = [{"a", 3}, {"b", 1}]
+
+    results =
+      keys
+      |> Enum.map(&HRW.owner(&1, weights, scorer: %Weighted{}))
+      |> Enum.frequencies()
+
+    # b wins ~1/4 of the time, a wins ~3/4 (roughly proportional to weight ratio 3:1)
+    total = results["a"] + results["b"]
+    ratio = results["a"] / total
+
+    # Expect a to win roughly 75% of the time; allow ±15% noise
+    assert ratio >= 0.60 and ratio <= 0.90
+  end
+
+  test "weighted skeleton routing matches weight ratios" do
+    nodes = [{"a", 3}, {"b", 2}, {"c", 1}]
+    skeleton = HRW.build(nodes, scorer: %Weighted{})
+
+    keys = Enum.map(1..500, &"key-#{&1}")
+
+    results =
+      keys
+      |> Enum.map(&HRW.owner(&1, skeleton))
+      |> Enum.frequencies()
+
+    total = results["a"] + results["b"] + results["c"]
+    a_ratio = results["a"] / total
+
+    # alpha has weight 3, total weight 6 → 50% of traffic
+    assert a_ratio >= 0.40 and a_ratio <= 0.60
+  end
+
+  test "weighted skeleton handles configurations with empty branches" do
+    nodes = Enum.map(1..100, fn i -> {"node#{i}", 1} end)
+    sk = HRW.build(nodes, scorer: %HRW.Weighted{}, fanout: 2, cluster_size: 15)
+    # capacity 8, 6 clusters, 2 empty slots
+
+    keys = Enum.map(1..10_000, fn i -> "key-#{i}" end)
+    results = Enum.map(keys, &HRW.owner(&1, sk))
+    assert Enum.all?(results, &is_binary/1)
+    assert length(Enum.uniq(results)) > 90
+  end
+end