From ae84717255cb196013c1b0fe0a625b2b443e7dc4 Mon Sep 17 00:00:00 2001
From: Matthew Wild <mwild1@gmail.com>
Date: Fri, 2 Dec 2022 20:32:36 +0000
Subject: util.hashring: Support associating arbitrary data with nodes

In this API, a 'node' is always a simple text string. Sometimes the caller may
have a more complex structure representing a node, but the hash ring is really
only concerned with the node's name.

This API change allows :add_nodes() to take a table of `node_name = value`
pairs, as well as the simple array of node names previously accepted.

The 'value' of the selected node is returned as a new second result from
:get_node().

If no value is passed when a node is added, it defaults to `true` (as before,
but this was never previously exposed).
---
 spec/util_hashring_spec.lua |  7 +++++++
 util/hashring.lua           | 32 +++++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/spec/util_hashring_spec.lua b/spec/util_hashring_spec.lua
index 05f053e4..4f6ec3a3 100644
--- a/spec/util_hashring_spec.lua
+++ b/spec/util_hashring_spec.lua
@@ -83,4 +83,11 @@ describe("util.hashring", function ()
 		end
 	end);
 
+	it("should support values associated with nodes", function ()
+		local r = hashring.new(128, sha256);
+		r:add_node("node1", { a = 1 });
+		local node, value = r:get_node("foo");
+		assert.is_equal("node1", node);
+		assert.same({ a = 1 }, value);
+	end);
 end);
diff --git a/util/hashring.lua b/util/hashring.lua
index d4555669..5e71654b 100644
--- a/util/hashring.lua
+++ b/util/hashring.lua
@@ -1,3 +1,5 @@
+local it = require "util.iterators";
+
 local function generate_ring(nodes, num_replicas, hash)
 	local new_ring = {};
 	for _, node_name in ipairs(nodes) do
@@ -28,18 +30,22 @@ local function new(num_replicas, hash_function)
 	return setmetatable({ nodes = {}, num_replicas = num_replicas, hash = hash_function }, hashring_mt);
 end;
 
-function hashring_methods:add_node(name)
+function hashring_methods:add_node(name, value)
 	self.ring = nil;
-	self.nodes[name] = true;
+	self.nodes[name] = value == nil and true or value;
 	table.insert(self.nodes, name);
 	return true;
 end
 
 function hashring_methods:add_nodes(nodes)
 	self.ring = nil;
-	for _, node_name in ipairs(nodes) do
-		if not self.nodes[node_name] then
-			self.nodes[node_name] = true;
+	local iter = pairs;
+	if nodes[1] then -- simple array?
+		iter = it.values;
+	end
+	for node_name, node_value in iter(nodes) do
+		if self.nodes[node_name] == nil then
+			self.nodes[node_name] = node_value == nil and true or node_value;
 			table.insert(self.nodes, node_name);
 		end
 	end
@@ -48,7 +54,7 @@ end
 
 function hashring_methods:remove_node(node_name)
 	self.ring = nil;
-	if self.nodes[node_name] then
+	if self.nodes[node_name] ~= nil then
 		for i, stored_node_name in ipairs(self.nodes) do
 			if node_name == stored_node_name then
 				self.nodes[node_name] = nil;
@@ -69,18 +75,26 @@ end
 
 function hashring_methods:clone()
 	local clone_hashring = new(self.num_replicas, self.hash);
-	clone_hashring:add_nodes(self.nodes);
+	for node_name, node_value in pairs(self.nodes) do
+		clone_hashring.nodes[node_name] = node_value;
+	end
+	clone_hashring.ring = nil;
 	return clone_hashring;
 end
 
 function hashring_methods:get_node(key)
+	local node;
 	local key_hash = self.hash(key);
 	for _, replica_hash in ipairs(self.ring) do
 		if key_hash < replica_hash then
-			return self.ring[replica_hash];
+			node = self.ring[replica_hash];
+			break;
 		end
 	end
-	return self.ring[self.ring[1]];
+	if not node then
+		node = self.ring[self.ring[1]];
+	end
+	return node, self.nodes[node];
 end
 
 return {
-- 
cgit v1.2.3