shards_test.go

package kinesis

import (
	"fmt"
	"math"
	"sort"
	"testing"

	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/signalfx/golib/v3/pointer"
	"github.com/stretchr/testify/assert"
)

var xs42 = [][3]string{
	{"0", "8101961117165201511032728748375433605", "shardId-000000000528"},
	{"8101961117165201511032728748375433606", "16203922234330403022065457496750867211", "shardId-000000000532"},
	{"16203922234330403022065457496750867212", "24305883351495604533098186245126300817", "shardId-000000000535"},
	{"24305883351495604533098186245126300818", "32407844468660806044130914993501734423", "shardId-000000000538"},
	{"32407844468660806044130914993501734424", "40509805585826007555163643741877168029", "shardId-000000000539"},
	{"40509805585826007555163643741877168030", "48611766702991209066196372490252601635", "shardId-000000000543"},
	{"48611766702991209066196372490252601636", "56713727820156410577229101238628035241", "shardId-000000000546"},
	{"56713727820156410577229101238628035242", "64815688937321612088261829987003468847", "shardId-000000000549"},
	{"64815688937321612088261829987003468848", "72917650054486813599294558735378902453", "shardId-000000000550"},
	{"72917650054486813599294558735378902454", "81019611171652015110327287483754336059", "shardId-000000000554"},
	{"81019611171652015110327287483754336060", "89121572288817216621360016232129769666", "shardId-000000000557"},
	{"89121572288817216621360016232129769667", "97223533405982418132392744980505203272", "shardId-000000000560"},
	{"97223533405982418132392744980505203273", "105325494523147619643425473728880636878", "shardId-000000000561"},
	{"105325494523147619643425473728880636879", "113427455640312821154458202477256070484", "shardId-000000000565"},
	{"113427455640312821154458202477256070485", "121529416757478022665490931225631504090", "shardId-000000000568"},
	{"121529416757478022665490931225631504091", "129631377874643224176523659974006937696", "shardId-000000000571"},
	{"129631377874643224176523659974006937697", "137733338991808425687556388722382371302", "shardId-000000000572"},
	{"137733338991808425687556388722382371303", "145835300108973627198589117470757804908", "shardId-000000000576"},
	{"145835300108973627198589117470757804909", "153937261226138828709621846219133238514", "shardId-000000000579"},
	{"153937261226138828709621846219133238515", "162039222343304030220654574967508672120", "shardId-000000000582"},
	{"162039222343304030220654574967508672121", "170141183460469231731687303715884105727", "shardId-000000000581"},
	{"170141183460469231731687303715884105728", "178243144577634433242720032464259539333", "shardId-000000000583"},
	{"178243144577634433242720032464259539334", "186345105694799634753752761212634972939", "shardId-000000000587"},
	{"186345105694799634753752761212634972940", "194447066811964836264785489961010406545", "shardId-000000000590"},
	{"194447066811964836264785489961010406546", "202549027929130037775818218709385840151", "shardId-000000000593"},
	{"202549027929130037775818218709385840152", "210650989046295239286850947457761273757", "shardId-000000000594"},
	{"210650989046295239286850947457761273758", "218752950163460440797883676206136707363", "shardId-000000000598"},
	{"218752950163460440797883676206136707364", "226854911280625642308916404954512140969", "shardId-000000000601"},
	{"226854911280625642308916404954512140970", "234956872397790843819949133702887574575", "shardId-000000000604"},
	{"234956872397790843819949133702887574576", "243058833514956045330981862451263008181", "shardId-000000000605"},
	{"243058833514956045330981862451263008182", "251160794632121246842014591199638441787", "shardId-000000000609"},
	{"251160794632121246842014591199638441788", "259262755749286448353047319948013875394", "shardId-000000000612"},
	{"259262755749286448353047319948013875395", "267364716866451649864080048696389309000", "shardId-000000000615"},
	{"267364716866451649864080048696389309001", "275466677983616851375112777444764742606", "shardId-000000000616"},
	{"275466677983616851375112777444764742607", "283568639100782052886145506193140176212", "shardId-000000000620"},
	{"283568639100782052886145506193140176213", "291670600217947254397178234941515609818", "shardId-000000000623"},
	{"291670600217947254397178234941515609819", "299772561335112455908210963689891043424", "shardId-000000000626"},
	{"299772561335112455908210963689891043425", "307874522452277657419243692438266477030", "shardId-000000000627"},
	{"307874522452277657419243692438266477031", "315976483569442858930276421186641910636", "shardId-000000000631"},
	{"315976483569442858930276421186641910637", "324078444686608060441309149935017344242", "shardId-000000000634"},
	{"332180405803773261952341878683392777849", "340282366920938463463374607431768211455", "shardId-000000000636"},
	{"324078444686608060441309149935017344243", "332180405803773261952341878683392777848", "shardId-000000000637"},
}
var xs16 = [][3]string{
	{"0", "21267647932558653966460912964485513215", "shardId-000000000000"},
	{"21267647932558653966460912964485513216", "42535295865117307932921825928971026431", "shardId-000000000001"},
	{"42535295865117307932921825928971026432", "63802943797675961899382738893456539647", "shardId-000000000002"},
	{"63802943797675961899382738893456539648", "85070591730234615865843651857942052863", "shardId-000000000003"},
	{"85070591730234615865843651857942052864", "106338239662793269832304564822427566079", "shardId-000000000004"},
	{"106338239662793269832304564822427566080", "127605887595351923798765477786913079295", "shardId-000000000005"},
	{"127605887595351923798765477786913079296", "148873535527910577765226390751398592511", "shardId-000000000006"},
	{"148873535527910577765226390751398592512", "170141183460469231731687303715884105727", "shardId-000000000007"},
	{"170141183460469231731687303715884105728", "191408831393027885698148216680369618943", "shardId-000000000008"},
	{"191408831393027885698148216680369618944", "212676479325586539664609129644855132159", "shardId-000000000009"},
	{"212676479325586539664609129644855132160", "233944127258145193631070042609340645375", "shardId-000000000010"},
	{"233944127258145193631070042609340645376", "255211775190703847597530955573826158591", "shardId-000000000011"},
	{"255211775190703847597530955573826158592", "276479423123262501563991868538311671807", "shardId-000000000012"},
	{"276479423123262501563991868538311671808", "297747071055821155530452781502797185023", "shardId-000000000013"},
	{"297747071055821155530452781502797185024", "319014718988379809496913694467282698239", "shardId-000000000014"},
	{"319014718988379809496913694467282698240", "340282366920938463463374607431768211455", "shardId-000000000015"},
}

var xs4 = [][3]string{
	{"0", "85070591730234615865843651857942052863", "shardId-000000000000"},
	{"85070591730234615865843651857942052864", "170141183460469231731687303715884105727", "shardId-000000000001"},
	{"170141183460469231731687303715884105728", "255211775190703847597530955573826158591", "shardId-000000000002"},
	{"255211775190703847597530955573826158592", "340282366920938463463374607431768211455", "shardId-000000000003"},
}

const benchTestTraceID = "bd7a977555f6b982bd7a977555f6b982"

func BenchmarkV1Shards(b *testing.B) {
	ss := generateV1(xs16)
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		v, err := getShardProducerIndex(ss, benchTestTraceID)
		if err != nil || v != 13 {
			b.Fatal(err)
		}
	}
}
func BenchmarkV2Shards(b *testing.B) {
	shards := generateV2(xs16)
	var v int
	b.ResetTimer()
	b.ReportAllocs()
	var err error
	for i := 0; i < b.N; i++ {
		v, err = shards.getIndex(benchTestTraceID)
		if v != 13 || err != nil {
			b.Fatal("not 13")
		}
	}
}

func BenchmarkV2ShardsNoPower(b *testing.B) {
	shards := generateV2(xs16[1:])
	var v int
	b.ResetTimer()
	b.ReportAllocs()
	var err error
	for i := 0; i < b.N; i++ {
		v, err = shards.getIndex(benchTestTraceID)
		if v != 12 || err != nil {
			b.Fatal("not 12")
		}
	}
}

func generateV1(xs [][3]string) shards {
	ss := make(shards, len(xs))
	for i, x := range xs {
		s := Shard{
			shardID:         x[2],
			startingHashKey: toBigInt(x[0]),
			endingHashKey:   toBigInt(x[1]),
		}
		ss[i] = s
	}
	sort.Sort(ss)
	return ss
}

type testkin struct {
	xs [][3]string
}

func (t *testkin) ListShards(_ *kinesis.ListShardsInput) (*kinesis.ListShardsOutput, error) {
	ret := &kinesis.ListShardsOutput{}
	for _, x := range t.xs {
		s := &kinesis.Shard{
			ShardId: pointer.String(x[2]),
			HashKeyRange: &kinesis.HashKeyRange{
				EndingHashKey:   pointer.String(x[1]),
				StartingHashKey: pointer.String(x[0]),
			},
			SequenceNumberRange: &kinesis.SequenceNumberRange{},
		}
		ret.Shards = append(ret.Shards, s)
	}
	return ret, nil
}

func generateV2(xs [][3]string) *ShardInfo {
	ret, _ := getShardInfo(&testkin{xs: xs}, "")
	ret.shiftLen = uint(128 - math.Log2(float64(len(ret.shards))))
	ret.power = math.Ceil(math.Log2(float64(len(ret.shards)))) == math.Floor(math.Log2(float64(len(ret.shards))))

	return ret
}

func getShardProducerIndex(ss shards, traceID string) (int, error) {
	for i, s := range ss {
		key := partitionKeyToHashKey(traceID)
		ok, err := s.belongsToShardKey(key)
		if err != nil {
			return -1, err
		}
		if ok {
			return i, nil
		}
	}
	return -1, fmt.Errorf("no shard found for parition key %s", traceID)
}

func TestEquality(t *testing.T) {
	tests := []struct {
		name string
		test [][3]string
	}{
		{name: "42 vnodes", test: xs42},
		{name: "16 vnodes", test: xs16},
		{name: "16 vnodes", test: xs16},
		{name: "4 vnodes", test: xs4},
	}
	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			info := generateV2(test.test)
			shards := generateV1(test.test)
			for _, x := range []string{
				"07e8e0fe2b145563",
				"04b6b36f98010a35",
				"bd7a977555f6b982",
				"0000000000000000bd7a977555f6b982",
				"bd7a977555f6b982bd7a977555f6b982",
				"0",
			} {
				index, _ := getShardProducerIndex(shards, x)
				index2, _ := info.getIndex(x)
				assert.Equal(t, index, index2, x)
			}
		})
	}
}