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  }