-
Notifications
You must be signed in to change notification settings - Fork 77
/
chash.lua
327 lines (241 loc) · 7.38 KB
/
chash.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
--
-- some of them borrow from https://github.com/cloudflare/lua-resty-json
--
local bit = require "bit"
local ffi = require 'ffi'
local utils = require "resty.balancer.utils"
local new_tab = utils.new_tab
local copy = utils.copy
local ffi_new = ffi.new
local C = ffi.C
local crc32 = ngx.crc32_short
local setmetatable = setmetatable
local floor = math.floor
local pairs = pairs
local tostring = tostring
local tonumber = tonumber
local bxor = bit.bxor
local error = error
local CHASH_OK = 0
ffi.cdef[[
typedef unsigned int uint32_t;
typedef struct {
uint32_t hash;
uint32_t id;
} chash_point_t;
void chash_point_init(chash_point_t *points, uint32_t base_hash, uint32_t start,
uint32_t num, uint32_t id);
int chash_point_sort(chash_point_t *points, uint32_t size);
int chash_point_add(chash_point_t *old_points, uint32_t old_length,
uint32_t base_hash, uint32_t from, uint32_t num, uint32_t id,
chash_point_t *new_points);
int chash_point_reduce(chash_point_t *old_points, uint32_t old_length,
uint32_t base_hash, uint32_t from, uint32_t num, uint32_t id);
void chash_point_delete(chash_point_t *old_points, uint32_t old_length,
uint32_t id);
]]
--
-- Find shared object file package.cpath, obviating the need of setting
-- LD_LIBRARY_PATH
-- Or we should add a little patch for ffi.load ?
--
local function load_shared_lib(so_name)
local string_gmatch = string.gmatch
local string_match = string.match
local io_open = io.open
local io_close = io.close
local cpath = package.cpath
local postfix = ".so"
if ffi.os == "OSX" then
postfix = ".dylib"
end
so_name = so_name .. postfix
for k, _ in string_gmatch(cpath, "[^;]+") do
local fpath = string_match(k, "(.*/)")
fpath = fpath .. so_name
-- Don't get me wrong, the only way to know if a file exist is trying
-- to open it.
local f = io_open(fpath)
if f ~= nil then
io_close(f)
return ffi.load(fpath)
end
end
end
local _M = {}
local mt = { __index = _M }
local clib = load_shared_lib("librestychash")
if not clib then
error("can not load librestychash")
end
local CONSISTENT_POINTS = 160 -- points per server
local pow32 = math.pow(2, 32)
local chash_point_t = ffi.typeof("chash_point_t[?]")
local function _precompute(nodes)
local n, total_weight = 0, 0
for id, weight in pairs(nodes) do
n = n + 1
total_weight = total_weight + weight
end
local newnodes = copy(nodes)
local ids = new_tab(n, 0)
local npoints = total_weight * CONSISTENT_POINTS
local points = ffi_new(chash_point_t, npoints)
local start, index = 0, 0
for id, weight in pairs(nodes) do
local num = weight * CONSISTENT_POINTS
local base_hash = bxor(crc32(tostring(id)), 0xffffffff)
index = index + 1
ids[index] = id
clib.chash_point_init(points, base_hash, start, num, index)
start = start + num
end
if clib.chash_point_sort(points, npoints) ~= CHASH_OK then
error("no memory")
end
return ids, points, npoints, newnodes
end
function _M.new(_, nodes)
local ids, points, npoints, newnodes = _precompute(nodes)
local self = {
nodes = newnodes, -- it's safer to copy one
ids = ids,
points = points,
npoints = npoints, -- points number
size = npoints,
}
return setmetatable(self, mt)
end
function _M.reinit(self, nodes)
self.ids, self.points, self.npoints, self.nodes = _precompute(nodes)
self.size = self.npoints
end
local function _delete(self, id)
local nodes = self.nodes
local ids = self.ids
local old_weight = nodes[id]
if not old_weight then
return
end
local index = 1
-- find the index: O(n)
while ids[index] ~= id do
index = index + 1
end
nodes[id] = nil
ids[index] = nil
clib.chash_point_delete(self.points, self.npoints, index)
self.npoints = self.npoints - CONSISTENT_POINTS * old_weight
end
_M.delete = _delete
local function _incr(self, id, weight)
local weight = tonumber(weight) or 1
local nodes = self.nodes
local ids = self.ids
local old_weight = nodes[id]
local index = 1
if old_weight then
-- find the index: O(n)
while ids[index] ~= id do
index = index + 1
end
else
old_weight = 0
index = #ids + 1
ids[index] = id
end
nodes[id] = old_weight + weight
local new_points = self.points
local new_npoints = self.npoints + weight * CONSISTENT_POINTS
if self.size < new_npoints then
new_points = ffi_new(chash_point_t, new_npoints)
end
local base_hash = bxor(crc32(tostring(id)), 0xffffffff)
local rc = clib.chash_point_add(self.points, self.npoints, base_hash,
old_weight * CONSISTENT_POINTS,
weight * CONSISTENT_POINTS,
index, new_points)
if rc ~= CHASH_OK then
error("no memory")
end
if self.size < new_npoints then
self.size = new_npoints
end
self.points = new_points
self.npoints = new_npoints
end
_M.incr = _incr
local function _decr(self, id, weight)
local weight = tonumber(weight) or 1
local nodes = self.nodes
local ids = self.ids
local old_weight = nodes[id]
if not old_weight then
return
end
if old_weight <= weight then
return _delete(self, id)
end
local index = 1
-- find the index: O(n)
while ids[index] ~= id do
index = index + 1
end
local base_hash = bxor(crc32(tostring(id)), 0xffffffff)
local from = (old_weight - weight) * CONSISTENT_POINTS
local num = CONSISTENT_POINTS * weight
local rc = clib.chash_point_reduce(self.points, self.npoints, base_hash,
from, num, index)
if rc ~= CHASH_OK then
error("no memory")
end
nodes[id] = old_weight - weight
self.npoints = self.npoints - CONSISTENT_POINTS * weight
end
_M.decr = _decr
function _M.set(self, id, new_weight)
local new_weight = tonumber(new_weight) or 0
local old_weight = self.nodes[id] or 0
if old_weight == new_weight then
return true
end
if old_weight < new_weight then
return _incr(self, id, new_weight - old_weight)
end
return _decr(self, id, old_weight - new_weight)
end
local function _find_id(points, npoints, hash)
local step = pow32 / npoints
local index = floor(hash / step)
local max_index = npoints - 1
-- it seems safer to do this
if index > max_index then
index = max_index
end
-- find the first points >= hash
if points[index].hash >= hash then
for i = index, 1, -1 do
if points[i - 1].hash < hash then
return points[i].id, i
end
end
return points[0].id, 0
end
for i = index + 1, max_index do
if hash <= points[i].hash then
return points[i].id, i
end
end
return points[0].id, 0
end
function _M.find(self, key)
local hash = crc32(tostring(key))
local id, index = _find_id(self.points, self.npoints, hash)
return self.ids[id], index
end
function _M.next(self, index)
local new_index = (index + 1) % self.npoints
local id = self.points[new_index].id
return self.ids[id], new_index
end
return _M