perf: saving SendToReplicas results with a bitmap (#55)

Signed-off-by: Rueian <rueiancsie@gmail.com>

Author: rueian

bitmap.go

@@ -0,0 +1,31 @@
+package valkey
+
+type bitmap struct {
+	bits [3]uint64 // 24 bytes
+	exts []uint64  // 24 bytes
+}
+
+func (b *bitmap) Init(n int) {
+	if n > 192 {
+		b.exts = make([]uint64, (n-192+63)/64)
+	}
+}
+
+func (b *bitmap) Set(i int) {
+	if i < 192 {
+		b.bits[i/64] |= 1 << (i % 64)
+	} else {
+		b.exts[(i-192)/64] |= 1 << (i % 64)
+	}
+}
+
+func (b *bitmap) Get(i int) bool {
+	if i < 192 {
+		return b.bits[i/64]&(1<<(i%64)) != 0
+	}
+	return b.exts[(i-192)/64]&(1<<(i%64)) != 0
+}
+
+func (b *bitmap) Len() int {
+	return len(b.exts)*64 + 192
+}

bitmap_test.go

@@ -0,0 +1,25 @@
+package valkey
+
+import "testing"
+
+func TestBitmap(t *testing.T) {
+	var bm bitmap
+	bm.Init(2953)
+	if bm.Len() < 2953 {
+		t.Errorf("bitmap length = %d, expected at least %d", bm.Len(), 2953)
+	}
+	for i := 0; i < bm.Len(); i++ {
+		bm.Set(i)
+		for j := 0; j < bm.Len(); j++ {
+			if j <= i {
+				if !bm.Get(j) {
+					t.Errorf("bitmap[%d] should be set after iteration %d, but is not", j, i)
+				}
+			} else {
+				if bm.Get(j) {
+					t.Errorf("bitmap[%d] should not be set after iteration %d, but is", j, i)
+				}
+			}
+		}
+	}
+}

cluster.go

@@ -574,24 +574,19 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 	count := conncountp.Get(len(c.conns), len(c.conns))
 
 	if !init && c.rslots != nil && c.opt.SendToReplicas != nil {
-		var destination []conn
-		var stackDestination [32]conn
-		if len(multi) <= len(stackDestination) {
-			destination = stackDestination[:len(multi)]
-		} else {
-			destination = make([]conn, len(multi))
-		}
+		var bm bitmap
+		bm.Init(len(multi))
 		for i, cmd := range multi {
 			var cc conn
 			if c.opt.SendToReplicas(cmd) {
+				bm.Set(i)
 				cc = c.rslots[cmd.Slot()]
 			} else {
 				cc = c.pslots[cmd.Slot()]
 			}
 			if cc == nil {
 				return nil, false
 			}
-			destination[i] = cc
 			count.m[cc]++
 		}
 
@@ -602,7 +597,12 @@ func (c *clusterClient) _pickMulti(multi []Completed) (retries *connretry, init
 		conncountp.Put(count)
 
 		for i, cmd := range multi {
-			cc := destination[i]
+			var cc conn
+			if bm.Get(i) {
+				cc = c.rslots[cmd.Slot()]
+			} else {
+				cc = c.pslots[cmd.Slot()]
+			}
 			re := retries.m[cc]
 			re.commands = append(re.commands, cmd)
 			re.cIndexes = append(re.cIndexes, i)