go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/consistenthash/consistent_hash_test.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
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.

*/

package consistenthash

import (
	"encoding/json"
	"fmt"
	"sort"
	"testing"

	"go.charczuk.com/sdk/assert"
	. "go.charczuk.com/sdk/assert"
)

func Test_ConsistentHash_typical(t *testing.T) {
	const bucketCount = 5
	const replicas = 32
	const itemCount = 150

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}
	var items []string
	for x := 0; x < itemCount; x++ {
		items = append(items, fmt.Sprintf("google-%d", x))
	}

	ch := New(OptReplicas(replicas))
	ch.AddBuckets(buckets...)

	ItsLen(t, ch.hashring, bucketCount*ch.Replicas())

	returnedBuckets := ch.Buckets()
	ItsLen(t, returnedBuckets, bucketCount)

	assignments := ch.Assignments(items...)

	worker0items := assignments["worker-0"]
	ItsNotEmpty(t, worker0items)

	worker1items := assignments["worker-1"]
	ItsNotEmpty(t, worker1items)

	worker2items := assignments["worker-2"]
	ItsNotEmpty(t, worker2items)

	worker3items := assignments["worker-3"]
	ItsNotEmpty(t, worker3items)

	worker4items := assignments["worker-4"]
	ItsNotEmpty(t, worker4items)

	// verify that all the consistent hash items exist in exactly one bucket assignment
	itemsExistInSets(t, items, worker0items, worker1items, worker2items, worker3items, worker4items)

	// verify that all the bucket assignments are disjoint, that is
	// none of the items in an assignment exist in another assignment
	setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items)

	// verify this is also the case through the `IsAssigned` method
	for _, item := range worker0items {
		ItsEqual(t, true, ch.IsAssigned("worker-0", item))
	}

	// verify the worker-0 items are not
	// assigned to any other nodes
	for _, item := range worker0items {
		ItsEqual(t, false, ch.IsAssigned("worker-1", item))
		ItsEqual(t, false, ch.IsAssigned("worker-2", item))
		ItsEqual(t, false, ch.IsAssigned("worker-3", item))
		ItsEqual(t, false, ch.IsAssigned("worker-4", item))
	}

	// verify the sets are relatively evenly sized
	// this is generally the most likely test to fail, because
	// the hash bucket assignments can be lumpy.
	setsAreEvenlySized(t, itemCount, worker0items, worker1items, worker2items, worker3items, worker4items)
}

func Test_ConsistentHash_AddBuckets(t *testing.T) {
	ch := new(ConsistentHash)

	res := ch.AddBuckets("worker-0", "worker-1", "worker-2")
	ItsEqual(t, true, res)
	res = ch.AddBuckets("worker-0", "worker-1", "worker-2")
	ItsEqual(t, false, res, "we should return false if _no_ new buckets were added")
	res = ch.AddBuckets("worker-0", "worker-1", "worker-2", "worker-3")
	ItsEqual(t, true, res, "we should return true if _any_ new buckets were added")
	buckets := ch.Buckets()
	sort.Strings(buckets)
	ItsLen(t, buckets, 4)
	ItsEqual(t, []string{"worker-0", "worker-1", "worker-2", "worker-3"}, buckets)
}

func Test_ConsistentHash_RemoveBucket(t *testing.T) {
	ch := new(ConsistentHash)

	res := ch.AddBuckets("worker-0", "worker-1", "worker-2")
	ItsEqual(t, true, res)
	res = ch.RemoveBucket("worker-3")
	ItsEqual(t, false, res, "we should return false if bucket not found")
	res = ch.RemoveBucket("worker-2")
	ItsEqual(t, true, res, "we should return true if bucket found")
	buckets := ch.Buckets()
	ItsLen(t, buckets, 2)
	ItsEqual(t, []string{"worker-0", "worker-1"}, buckets)
}

func Test_ConsistentHash_redistribute_addBuckets(t *testing.T) {
	const bucketCount = 5
	const replicas = 32
	const itemCount = 150
	const itemsPerBucket = itemCount / bucketCount

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}
	var items []string
	for x := 0; x < itemCount; x++ {
		items = append(items, fmt.Sprintf("google-%d", x))
	}

	ch := New(OptReplicas(replicas))
	ch.AddBuckets(buckets...)

	assignments := ch.Assignments(items...)

	oldWorker0items := assignments["worker-0"]
	ItsNotEmpty(t, oldWorker0items)

	oldWorker1items := assignments["worker-1"]
	ItsNotEmpty(t, oldWorker1items)

	oldWorker2items := assignments["worker-2"]
	ItsNotEmpty(t, oldWorker2items)

	oldWorker3items := assignments["worker-3"]
	ItsNotEmpty(t, oldWorker3items)

	oldWorker4items := assignments["worker-4"]
	ItsNotEmpty(t, oldWorker4items)

	// simulate adding a bucket
	ch.AddBuckets("worker-5")
	ItsEqual(t, bucketCount+1, len(ch.buckets))
	ItsLen(t, ch.hashring, (bucketCount+1)*ch.Replicas())

	newAssignments := ch.Assignments(items...)
	ItsEqual(t, bucketCount+1, len(newAssignments), "assignments length should mirror buckets")

	worker0items := newAssignments["worker-0"]
	ItsNotEmpty(t, worker0items)

	worker1items := newAssignments["worker-1"]
	ItsNotEmpty(t, worker1items)

	worker2items := newAssignments["worker-2"]
	ItsNotEmpty(t, worker2items)

	worker3items := newAssignments["worker-3"]
	ItsNotEmpty(t, worker3items)

	worker4items := newAssignments["worker-4"]
	ItsNotEmpty(t, worker4items)

	worker5items := newAssignments["worker-5"]
	ItsNotEmpty(t, worker5items)

	// verify that all the bucket assignments are disjoint, that is
	// none of the items in an assignment exist in another assignment
	setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items, worker5items)

	// verify that all the consistent hash items exist in a bucket assignment
	itemsExistInSets(t, items, worker0items, worker1items, worker2items, worker3items, worker4items, worker5items)

	// verify that we're moving items around consistently
	// that is, that we don't change the existing assignments
	// by much more than a single bucket's worth
	// of assignments
	itsConsistent(t, oldWorker0items, worker0items, itemsPerBucket)
	itsConsistent(t, oldWorker1items, worker1items, itemsPerBucket)
	itsConsistent(t, oldWorker2items, worker2items, itemsPerBucket)
	itsConsistent(t, oldWorker3items, worker3items, itemsPerBucket)
	itsConsistent(t, oldWorker4items, worker4items, itemsPerBucket)
}

func Test_ConsistentHash_redistribute_removeBucket(t *testing.T) {
	const bucketCount = 5
	const itemCount = 100
	const itemsPerBucket = itemCount / bucketCount
	const maxBucketDelta = 9

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}
	var items []string
	for x := 0; x < itemCount; x++ {
		items = append(items, fmt.Sprintf("google-%d", x))
	}

	ch := new(ConsistentHash)
	ch.AddBuckets(buckets...)

	assignments := ch.Assignments(items...)

	oldWorker0items := assignments["worker-0"]
	ItsNotEmpty(t, oldWorker0items)

	oldWorker1items := assignments["worker-1"]
	ItsNotEmpty(t, oldWorker1items)

	oldWorker2items := assignments["worker-2"]
	ItsNotEmpty(t, oldWorker2items)

	oldWorker3items := assignments["worker-3"]
	ItsNotEmpty(t, oldWorker3items)

	oldWorker4items := assignments["worker-4"]
	ItsNotEmpty(t, oldWorker4items)

	// the maximum number of items to move around is the
	// single bucket we're removing (plus a fudge factor)

	// simulate dropping a bucket (or node)
	ItsEqual(t, true, ch.RemoveBucket("worker-2"))
	ItsEqual(t, bucketCount-1, len(ch.buckets))
	ItsLen(t, ch.hashring, (bucketCount-1)*ch.Replicas())

	_, ok := ch.buckets["worker-2"]
	ItsEqual(t, false, ok)
	for _, bucket := range ch.hashring {
		ItsEqual(t, true, bucket.Bucket != "worker-2")
	}

	assignments = ch.Assignments(items...)
	ItsEqual(t, len(assignments), len(ch.buckets))

	worker0items := assignments["worker-0"]
	ItsNotEmpty(t, oldWorker0items)

	worker1items := assignments["worker-1"]
	ItsNotEmpty(t, oldWorker1items)

	worker2items := assignments["worker-2"]
	ItsNotEmpty(t, oldWorker2items)

	worker3items := assignments["worker-3"]
	ItsNotEmpty(t, oldWorker3items)

	worker4items := assignments["worker-4"]
	ItsNotEmpty(t, oldWorker4items)

	// verify that all the bucket assignments are disjoint, that is
	// none of the items in an assignment exist in another assignment
	setsAreDisjoint(t, worker0items, worker1items, worker2items, worker3items, worker4items)

	// verify that all the consistent hash items exist in a bucket assignment
	itemsExistInSets(t, items, worker0items, worker1items, worker3items, worker4items)

	// verify that we're moving items around consistently
	itsConsistent(t, oldWorker0items, worker0items, maxBucketDelta)
	itsConsistent(t, oldWorker1items, worker1items, maxBucketDelta)
	itsConsistent(t, oldWorker3items, worker3items, maxBucketDelta)
	itsConsistent(t, oldWorker4items, worker4items, maxBucketDelta)
}

func Test_ConsistentHash_notFound_removeBucket(t *testing.T) {
	const bucketCount = 5
	const itemCount = 100
	const itemsPerBucket = itemCount / bucketCount
	const maxBucketDelta = (itemsPerBucket / bucketCount) + 4

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}
	var items []string
	for x := 0; x < itemCount; x++ {
		items = append(items, fmt.Sprintf("google_%d", x))
	}

	ch := new(ConsistentHash)
	ch.AddBuckets(buckets...)

	oldAssignments := ch.Assignments(items...)
	ItsEqual(t, false, ch.RemoveBucket("not-worker-0"))
	newAssignments := ch.Assignments(items...)

	assignmentsAreEqual(t, oldAssignments, newAssignments)
}

func Test_ConsistentHash_String(t *testing.T) {
	const bucketCount = 5

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}

	ch := new(ConsistentHash)
	ch.AddBuckets(buckets...)

	ItsNotEmpty(t, []rune(ch.String()))
}

func Test_ConsistentHash_MarshalJSON(t *testing.T) {
	const bucketCount = 5

	var buckets []string
	for x := 0; x < bucketCount; x++ {
		buckets = append(buckets, fmt.Sprintf("worker-%d", x))
	}

	ch := new(ConsistentHash)
	ch.AddBuckets(buckets...)

	output, err := json.Marshal(ch)
	ItsEqual(t, nil, err)
	ItsNotEmpty(t, output)

	var verify []HashedBucket
	err = json.Unmarshal(output, &verify)
	ItsEqual(t, nil, err)
	ItsEqual(t, ch.hashring, verify)
}

func Test_ConsistentHash_insertionSort(t *testing.T) {
	ch := new(ConsistentHash)

	ch.insertionSort(HashedBucket{
		Hashcode: 1,
		Bucket:   "foo",
		Replica:  0,
	})
	assert.ItsLen(t, ch.hashring, 1)
	assert.ItsEqual(t, 1, ch.hashring[0].Hashcode)

	ch.insertionSort(HashedBucket{
		Hashcode: 10,
		Bucket:   "foo",
		Replica:  1,
	})
	assert.ItsLen(t, ch.hashring, 2)
	assert.ItsEqual(t, 1, ch.hashring[0].Hashcode)
	assert.ItsEqual(t, 10, ch.hashring[1].Hashcode)

	ch.insertionSort(HashedBucket{
		Hashcode: 5,
		Bucket:   "foo",
		Replica:  2,
	})
	assert.ItsLen(t, ch.hashring, 3)
	assert.ItsEqual(t, 1, ch.hashring[0].Hashcode)
	assert.ItsEqual(t, 5, ch.hashring[1].Hashcode)
	assert.ItsEqual(t, 10, ch.hashring[2].Hashcode)

	ch.insertionSort(HashedBucket{
		Hashcode: 7,
		Bucket:   "foo",
		Replica:  3,
	})
	assert.ItsLen(t, ch.hashring, 4)
	assert.ItsEqual(t, 1, ch.hashring[0].Hashcode)
	assert.ItsEqual(t, 5, ch.hashring[1].Hashcode)
	assert.ItsEqual(t, 7, ch.hashring[2].Hashcode)
	assert.ItsEqual(t, 10, ch.hashring[3].Hashcode)

	ch.insertionSort(HashedBucket{
		Hashcode: 11,
		Bucket:   "foo",
		Replica:  4,
	})
	assert.ItsLen(t, ch.hashring, 5)
	assert.ItsEqual(t, 1, ch.hashring[0].Hashcode)
	assert.ItsEqual(t, 5, ch.hashring[1].Hashcode)
	assert.ItsEqual(t, 7, ch.hashring[2].Hashcode)
	assert.ItsEqual(t, 10, ch.hashring[3].Hashcode)
	assert.ItsEqual(t, 11, ch.hashring[4].Hashcode)

	ch.insertionSort(HashedBucket{
		Hashcode: 0,
		Bucket:   "foo",
		Replica:  5,
	})
	assert.ItsLen(t, ch.hashring, 6)
	assert.ItsEqual(t, 0, ch.hashring[0].Hashcode)
	assert.ItsEqual(t, 1, ch.hashring[1].Hashcode)
	assert.ItsEqual(t, 5, ch.hashring[2].Hashcode)
	assert.ItsEqual(t, 7, ch.hashring[3].Hashcode)
	assert.ItsEqual(t, 10, ch.hashring[4].Hashcode)
	assert.ItsEqual(t, 11, ch.hashring[5].Hashcode)
}