util/hashring.lua

   1 local function generate_ring(nodes, num_replicas, hash)
   2         local new_ring = {};
   3         for _, node_name in ipairs(nodes) do
   4                 for replica = 1, num_replicas do
   5                         local replica_hash = hash(node_name..":"..replica);
   6                         new_ring[replica_hash] = node_name;
   7                         table.insert(new_ring, replica_hash);
   8                 end
   9         end
  10         table.sort(new_ring);
  11         return new_ring;
  12 end
  13
  14 local hashring_methods = {};
  15 local hashring_mt = {
  16         __index = function (self, k)
  17                 -- Automatically build self.ring if it's missing
  18                 if k == "ring" then
  19                         local ring = generate_ring(self.nodes, self.num_replicas, self.hash);
  20                         rawset(self, "ring", ring);
  21                         return ring;
  22                 end
  23                 return rawget(hashring_methods, k);
  24         end
  25 };
  26
  27 local function new(num_replicas, hash_function)
  28         return setmetatable({ nodes = {}, num_replicas = num_replicas, hash = hash_function }, hashring_mt);
  29 end;
  30
  31 function hashring_methods:add_node(name)
  32         self.ring = nil;
  33         self.nodes[name] = true;
  34         table.insert(self.nodes, name);
  35         return true;
  36 end
  37
  38 function hashring_methods:add_nodes(nodes)
  39         self.ring = nil;
  40         for _, node_name in ipairs(nodes) do
  41                 if not self.nodes[node_name] then
  42                         self.nodes[node_name] = true;
  43                         table.insert(self.nodes, node_name);
  44                 end
  45         end
  46         return true;
  47 end
  48
  49 function hashring_methods:remove_node(node_name)
  50         self.ring = nil;
  51         if self.nodes[node_name] then
  52                 for i, stored_node_name in ipairs(self.nodes) do
  53                         if node_name == stored_node_name then
  54                                 self.nodes[node_name] = nil;
  55                                 table.remove(self.nodes, i);
  56                                 return true;
  57                         end
  58                 end
  59         end
  60         return false;
  61 end
  62
  63 function hashring_methods:remove_nodes(nodes)
  64         self.ring = nil;
  65         for _, node_name in ipairs(nodes) do
  66                 self:remove_node(node_name);
  67         end
  68 end
  69
  70 function hashring_methods:clone()
  71         local clone_hashring = new(self.num_replicas, self.hash);
  72         clone_hashring:add_nodes(self.nodes);
  73         return clone_hashring;
  74 end
  75
  76 function hashring_methods:get_node(key)
  77         local key_hash = self.hash(key);
  78         for _, replica_hash in ipairs(self.ring) do
  79                 if key_hash < replica_hash then
  80                         return self.ring[replica_hash];
  81                 end
  82         end
  83         return self.ring[self.ring[1]];
  84 end
  85
  86 return {
  87         new = new;
  88 }