go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/consistenthash/consistent_hash_test.go (about) 1 /* 2 3 Copyright (c) 2023 - Present. Will Charczuk. All rights reserved. 4 Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository. 5 6 */ 7 8 package consistenthash 9 10 import ( 11 "encoding/json" 12 "fmt" 13 "sort" 14 "testing" 15 16 "go.charczuk.com/sdk/assert" 17 . "go.charczuk.com/sdk/assert" 18 ) 19 20 func Test_ConsistentHash_typical(t *testing.T) { 21 const bucketCount = 5 22 const replicas = 32 23 const itemCount = 150 24 25 var buckets []string 26 for x := 0; x < bucketCount; x++ { 27 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 28 } 29 var items []string 30 for x := 0; x < itemCount; x++ { 31 items = append(items, fmt.Sprintf("google-%d", x)) 32 } 33 34 ch := New(OptReplicas(replicas)) 35 ch.AddBuckets(buckets...) 36 37 ItsLen(t, ch.hashring, bucketCount*ch.Replicas()) 38 39 returnedBuckets := ch.Buckets() 40 ItsLen(t, returnedBuckets, bucketCount) 41 42 assignments := ch.Assignments(items...) 43 44 worker0items := assignments["worker-0"] 45 ItsNotEmpty(t, worker0items) 46 47 worker1items := assignments["worker-1"] 48 ItsNotEmpty(t, worker1items) 49 50 worker2items := assignments["worker-2"] 51 ItsNotEmpty(t, worker2items) 52 53 worker3items := assignments["worker-3"] 54 ItsNotEmpty(t, worker3items) 55 56 worker4items := assignments["worker-4"] 57 ItsNotEmpty(t, worker4items) 58 59 // verify that all the consistent hash items exist in exactly one bucket assignment 60 itemsExistInSets(t, items, worker0items, worker1items, worker2items, worker3items, worker4items) 61 62 // verify that all the bucket assignments are disjoint, that is 63 // none of the items in an assignment exist in another assignment 64 setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items) 65 66 // verify this is also the case through the `IsAssigned` method 67 for _, item := range worker0items { 68 ItsEqual(t, true, ch.IsAssigned("worker-0", item)) 69 } 70 71 // verify the worker-0 items are not 72 // assigned to any other nodes 73 for _, item := range worker0items { 74 ItsEqual(t, false, ch.IsAssigned("worker-1", item)) 75 ItsEqual(t, false, ch.IsAssigned("worker-2", item)) 76 ItsEqual(t, false, ch.IsAssigned("worker-3", item)) 77 ItsEqual(t, false, ch.IsAssigned("worker-4", item)) 78 } 79 80 // verify the sets are relatively evenly sized 81 // this is generally the most likely test to fail, because 82 // the hash bucket assignments can be lumpy. 83 setsAreEvenlySized(t, itemCount, worker0items, worker1items, worker2items, worker3items, worker4items) 84 } 85 86 func Test_ConsistentHash_AddBuckets(t *testing.T) { 87 ch := new(ConsistentHash) 88 89 res := ch.AddBuckets("worker-0", "worker-1", "worker-2") 90 ItsEqual(t, true, res) 91 res = ch.AddBuckets("worker-0", "worker-1", "worker-2") 92 ItsEqual(t, false, res, "we should return false if _no_ new buckets were added") 93 res = ch.AddBuckets("worker-0", "worker-1", "worker-2", "worker-3") 94 ItsEqual(t, true, res, "we should return true if _any_ new buckets were added") 95 buckets := ch.Buckets() 96 sort.Strings(buckets) 97 ItsLen(t, buckets, 4) 98 ItsEqual(t, []string{"worker-0", "worker-1", "worker-2", "worker-3"}, buckets) 99 } 100 101 func Test_ConsistentHash_RemoveBucket(t *testing.T) { 102 ch := new(ConsistentHash) 103 104 res := ch.AddBuckets("worker-0", "worker-1", "worker-2") 105 ItsEqual(t, true, res) 106 res = ch.RemoveBucket("worker-3") 107 ItsEqual(t, false, res, "we should return false if bucket not found") 108 res = ch.RemoveBucket("worker-2") 109 ItsEqual(t, true, res, "we should return true if bucket found") 110 buckets := ch.Buckets() 111 ItsLen(t, buckets, 2) 112 ItsEqual(t, []string{"worker-0", "worker-1"}, buckets) 113 } 114 115 func Test_ConsistentHash_redistribute_addBuckets(t *testing.T) { 116 const bucketCount = 5 117 const replicas = 32 118 const itemCount = 150 119 const itemsPerBucket = itemCount / bucketCount 120 121 var buckets []string 122 for x := 0; x < bucketCount; x++ { 123 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 124 } 125 var items []string 126 for x := 0; x < itemCount; x++ { 127 items = append(items, fmt.Sprintf("google-%d", x)) 128 } 129 130 ch := New(OptReplicas(replicas)) 131 ch.AddBuckets(buckets...) 132 133 assignments := ch.Assignments(items...) 134 135 oldWorker0items := assignments["worker-0"] 136 ItsNotEmpty(t, oldWorker0items) 137 138 oldWorker1items := assignments["worker-1"] 139 ItsNotEmpty(t, oldWorker1items) 140 141 oldWorker2items := assignments["worker-2"] 142 ItsNotEmpty(t, oldWorker2items) 143 144 oldWorker3items := assignments["worker-3"] 145 ItsNotEmpty(t, oldWorker3items) 146 147 oldWorker4items := assignments["worker-4"] 148 ItsNotEmpty(t, oldWorker4items) 149 150 // simulate adding a bucket 151 ch.AddBuckets("worker-5") 152 ItsEqual(t, bucketCount+1, len(ch.buckets)) 153 ItsLen(t, ch.hashring, (bucketCount+1)*ch.Replicas()) 154 155 newAssignments := ch.Assignments(items...) 156 ItsEqual(t, bucketCount+1, len(newAssignments), "assignments length should mirror buckets") 157 158 worker0items := newAssignments["worker-0"] 159 ItsNotEmpty(t, worker0items) 160 161 worker1items := newAssignments["worker-1"] 162 ItsNotEmpty(t, worker1items) 163 164 worker2items := newAssignments["worker-2"] 165 ItsNotEmpty(t, worker2items) 166 167 worker3items := newAssignments["worker-3"] 168 ItsNotEmpty(t, worker3items) 169 170 worker4items := newAssignments["worker-4"] 171 ItsNotEmpty(t, worker4items) 172 173 worker5items := newAssignments["worker-5"] 174 ItsNotEmpty(t, worker5items) 175 176 // verify that all the bucket assignments are disjoint, that is 177 // none of the items in an assignment exist in another assignment 178 setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items, worker5items) 179 180 // verify that all the consistent hash items exist in a bucket assignment 181 itemsExistInSets(t, items, worker0items, worker1items, worker2items, worker3items, worker4items, worker5items) 182 183 // verify that we're moving items around consistently 184 // that is, that we don't change the existing assignments 185 // by much more than a single bucket's worth 186 // of assignments 187 itsConsistent(t, oldWorker0items, worker0items, itemsPerBucket) 188 itsConsistent(t, oldWorker1items, worker1items, itemsPerBucket) 189 itsConsistent(t, oldWorker2items, worker2items, itemsPerBucket) 190 itsConsistent(t, oldWorker3items, worker3items, itemsPerBucket) 191 itsConsistent(t, oldWorker4items, worker4items, itemsPerBucket) 192 } 193 194 func Test_ConsistentHash_redistribute_removeBucket(t *testing.T) { 195 const bucketCount = 5 196 const itemCount = 100 197 const itemsPerBucket = itemCount / bucketCount 198 const maxBucketDelta = 9 199 200 var buckets []string 201 for x := 0; x < bucketCount; x++ { 202 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 203 } 204 var items []string 205 for x := 0; x < itemCount; x++ { 206 items = append(items, fmt.Sprintf("google-%d", x)) 207 } 208 209 ch := new(ConsistentHash) 210 ch.AddBuckets(buckets...) 211 212 assignments := ch.Assignments(items...) 213 214 oldWorker0items := assignments["worker-0"] 215 ItsNotEmpty(t, oldWorker0items) 216 217 oldWorker1items := assignments["worker-1"] 218 ItsNotEmpty(t, oldWorker1items) 219 220 oldWorker2items := assignments["worker-2"] 221 ItsNotEmpty(t, oldWorker2items) 222 223 oldWorker3items := assignments["worker-3"] 224 ItsNotEmpty(t, oldWorker3items) 225 226 oldWorker4items := assignments["worker-4"] 227 ItsNotEmpty(t, oldWorker4items) 228 229 // the maximum number of items to move around is the 230 // single bucket we're removing (plus a fudge factor) 231 232 // simulate dropping a bucket (or node) 233 ItsEqual(t, true, ch.RemoveBucket("worker-2")) 234 ItsEqual(t, bucketCount-1, len(ch.buckets)) 235 ItsLen(t, ch.hashring, (bucketCount-1)*ch.Replicas()) 236 237 _, ok := ch.buckets["worker-2"] 238 ItsEqual(t, false, ok) 239 for _, bucket := range ch.hashring { 240 ItsEqual(t, true, bucket.Bucket != "worker-2") 241 } 242 243 assignments = ch.Assignments(items...) 244 ItsEqual(t, len(assignments), len(ch.buckets)) 245 246 worker0items := assignments["worker-0"] 247 ItsNotEmpty(t, oldWorker0items) 248 249 worker1items := assignments["worker-1"] 250 ItsNotEmpty(t, oldWorker1items) 251 252 worker2items := assignments["worker-2"] 253 ItsNotEmpty(t, oldWorker2items) 254 255 worker3items := assignments["worker-3"] 256 ItsNotEmpty(t, oldWorker3items) 257 258 worker4items := assignments["worker-4"] 259 ItsNotEmpty(t, oldWorker4items) 260 261 // verify that all the bucket assignments are disjoint, that is 262 // none of the items in an assignment exist in another assignment 263 setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items) 264 265 // verify that all the consistent hash items exist in a bucket assignment 266 itemsExistInSets(t, items, worker0items, worker1items, worker3items, worker4items) 267 268 // verify that we're moving items around consistently 269 itsConsistent(t, oldWorker0items, worker0items, maxBucketDelta) 270 itsConsistent(t, oldWorker1items, worker1items, maxBucketDelta) 271 itsConsistent(t, oldWorker3items, worker3items, maxBucketDelta) 272 itsConsistent(t, oldWorker4items, worker4items, maxBucketDelta) 273 } 274 275 func Test_ConsistentHash_notFound_removeBucket(t *testing.T) { 276 const bucketCount = 5 277 const itemCount = 100 278 const itemsPerBucket = itemCount / bucketCount 279 const maxBucketDelta = (itemsPerBucket / bucketCount) + 4 280 281 var buckets []string 282 for x := 0; x < bucketCount; x++ { 283 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 284 } 285 var items []string 286 for x := 0; x < itemCount; x++ { 287 items = append(items, fmt.Sprintf("google_%d", x)) 288 } 289 290 ch := new(ConsistentHash) 291 ch.AddBuckets(buckets...) 292 293 oldAssignments := ch.Assignments(items...) 294 ItsEqual(t, false, ch.RemoveBucket("not-worker-0")) 295 newAssignments := ch.Assignments(items...) 296 297 assignmentsAreEqual(t, oldAssignments, newAssignments) 298 } 299 300 func Test_ConsistentHash_String(t *testing.T) { 301 const bucketCount = 5 302 303 var buckets []string 304 for x := 0; x < bucketCount; x++ { 305 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 306 } 307 308 ch := new(ConsistentHash) 309 ch.AddBuckets(buckets...) 310 311 ItsNotEmpty(t, []rune(ch.String())) 312 } 313 314 func Test_ConsistentHash_MarshalJSON(t *testing.T) { 315 const bucketCount = 5 316 317 var buckets []string 318 for x := 0; x < bucketCount; x++ { 319 buckets = append(buckets, fmt.Sprintf("worker-%d", x)) 320 } 321 322 ch := new(ConsistentHash) 323 ch.AddBuckets(buckets...) 324 325 output, err := json.Marshal(ch) 326 ItsEqual(t, nil, err) 327 ItsNotEmpty(t, output) 328 329 var verify []HashedBucket 330 err = json.Unmarshal(output, &verify) 331 ItsEqual(t, nil, err) 332 ItsEqual(t, ch.hashring, verify) 333 } 334 335 func Test_ConsistentHash_insertionSort(t *testing.T) { 336 ch := new(ConsistentHash) 337 338 ch.insertionSort(HashedBucket{ 339 Hashcode: 1, 340 Bucket: "foo", 341 Replica: 0, 342 }) 343 assert.ItsLen(t, ch.hashring, 1) 344 assert.ItsEqual(t, 1, ch.hashring[0].Hashcode) 345 346 ch.insertionSort(HashedBucket{ 347 Hashcode: 10, 348 Bucket: "foo", 349 Replica: 1, 350 }) 351 assert.ItsLen(t, ch.hashring, 2) 352 assert.ItsEqual(t, 1, ch.hashring[0].Hashcode) 353 assert.ItsEqual(t, 10, ch.hashring[1].Hashcode) 354 355 ch.insertionSort(HashedBucket{ 356 Hashcode: 5, 357 Bucket: "foo", 358 Replica: 2, 359 }) 360 assert.ItsLen(t, ch.hashring, 3) 361 assert.ItsEqual(t, 1, ch.hashring[0].Hashcode) 362 assert.ItsEqual(t, 5, ch.hashring[1].Hashcode) 363 assert.ItsEqual(t, 10, ch.hashring[2].Hashcode) 364 365 ch.insertionSort(HashedBucket{ 366 Hashcode: 7, 367 Bucket: "foo", 368 Replica: 3, 369 }) 370 assert.ItsLen(t, ch.hashring, 4) 371 assert.ItsEqual(t, 1, ch.hashring[0].Hashcode) 372 assert.ItsEqual(t, 5, ch.hashring[1].Hashcode) 373 assert.ItsEqual(t, 7, ch.hashring[2].Hashcode) 374 assert.ItsEqual(t, 10, ch.hashring[3].Hashcode) 375 376 ch.insertionSort(HashedBucket{ 377 Hashcode: 11, 378 Bucket: "foo", 379 Replica: 4, 380 }) 381 assert.ItsLen(t, ch.hashring, 5) 382 assert.ItsEqual(t, 1, ch.hashring[0].Hashcode) 383 assert.ItsEqual(t, 5, ch.hashring[1].Hashcode) 384 assert.ItsEqual(t, 7, ch.hashring[2].Hashcode) 385 assert.ItsEqual(t, 10, ch.hashring[3].Hashcode) 386 assert.ItsEqual(t, 11, ch.hashring[4].Hashcode) 387 388 ch.insertionSort(HashedBucket{ 389 Hashcode: 0, 390 Bucket: "foo", 391 Replica: 5, 392 }) 393 assert.ItsLen(t, ch.hashring, 6) 394 assert.ItsEqual(t, 0, ch.hashring[0].Hashcode) 395 assert.ItsEqual(t, 1, ch.hashring[1].Hashcode) 396 assert.ItsEqual(t, 5, ch.hashring[2].Hashcode) 397 assert.ItsEqual(t, 7, ch.hashring[3].Hashcode) 398 assert.ItsEqual(t, 10, ch.hashring[4].Hashcode) 399 assert.ItsEqual(t, 11, ch.hashring[5].Hashcode) 400 }