fix: too large arguments (#595)

valkey-io · Aug 11, 2024 · 9c25a9e · 9c25a9e
1 parent ae64b78
commit 9c25a9e
Show file tree

Hide file tree

Showing 4 changed files with 132 additions and 34 deletions.
diff --git a/valkeyprob/bloomfilter.go b/valkeyprob/bloomfilter.go
@@ -21,21 +21,12 @@ local numElements = tonumber(#ARGV) - 1
 local filterKey = KEYS[1]
 local counterKey = KEYS[2]
 
-local bitfieldArgs = { filterKey }
-for i=2, numElements+1 do
-    table.insert(bitfieldArgs, 'SET')
-    table.insert(bitfieldArgs, 'u1')
-    table.insert(bitfieldArgs, ARGV[i])
-    table.insert(bitfieldArgs, '1')
-end
-
-local bitset = redis.call('BITFIELD', unpack(bitfieldArgs))
-
 local counter = 0
 local oneBits = 0
-for i=1, #bitset do
-	oneBits = oneBits + bitset[i]
+for i=1, numElements do
+	local bitset = redis.call('BITFIELD', filterKey, 'SET', 'u1', ARGV[i+1], '1')
 
+	oneBits = oneBits + bitset[1]
 	if i % hashIterations == 0 then
 		if oneBits ~= hashIterations then
 			counter = counter + 1
@@ -53,22 +44,13 @@ local hashIterations = tonumber(ARGV[1])
 local numElements = tonumber(#ARGV) - 1
 local filterKey = KEYS[1]
 
-local bitfieldArgs = { filterKey }
-for i=2, numElements+1 do
-	local index = tonumber(ARGV[i])
-
-	table.insert(bitfieldArgs, 'GET')
-	table.insert(bitfieldArgs, 'u1')
-	table.insert(bitfieldArgs, index)
-end
-
-local bitset = redis.call('BITFIELD', unpack(bitfieldArgs))
-
 local result = {}
 local oneBits = 0
-for i=1, #bitset do
-	oneBits = oneBits + bitset[i]
+for i=1, numElements do
+	local index = tonumber(ARGV[i+1])
+	local bitset = redis.call('BITFIELD', filterKey, 'GET', 'u1', index)
 
+	oneBits = oneBits + bitset[1]
 	if i % hashIterations == 0 then
 		table.insert(result, oneBits == hashIterations)
 

diff --git a/valkeyprob/bloomfilter_test.go b/valkeyprob/bloomfilter_test.go
@@ -390,6 +390,43 @@ func TestBloomFilterAddMulti(t *testing.T) {
 			t.Error("Count is not 3")
 		}
 	})
+
+	t.Run("add very large number of items", func(t *testing.T) {
+		client, flushAllAndClose, err := setup()
+		if err != nil {
+			t.Error(err)
+		}
+		defer func() {
+			err := flushAllAndClose()
+			if err != nil {
+				t.Error(err)
+			}
+		}()
+
+		bf, err := NewBloomFilter(client, "test", 10000000, 0.1)
+		if err != nil {
+			t.Error(err)
+		}
+
+		// Above `LUAI_MAXCSTACK`(8000) limit
+		keys := make([]string, 8001)
+		for i := 0; i < 8001; i++ {
+			keys[i] = strconv.Itoa(i)
+		}
+
+		err = bf.AddMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+
+		count, err := bf.Count(context.Background())
+		if err != nil {
+			t.Error(err)
+		}
+		if count != 8001 {
+			t.Error("Count is not 1000")
+		}
+	})
 }
 
 func TestBloomFilterAddMultiError(t *testing.T) {
@@ -588,6 +625,45 @@ func TestBloomFilterExistsMulti(t *testing.T) {
 			t.Error("Exists is not empty")
 		}
 	})
+
+	t.Run("exists very large number of items", func(t *testing.T) {
+		client, flushAllAndClose, err := setup()
+		if err != nil {
+			t.Error(err)
+		}
+		defer func() {
+			err := flushAllAndClose()
+			if err != nil {
+				t.Error(err)
+			}
+		}()
+
+		bf, err := NewBloomFilter(client, "test", 10000000, 0.1)
+		if err != nil {
+			t.Error(err)
+		}
+
+		// Above `LUAI_MAXCSTACK`(8000) limit
+		keys := make([]string, 8001)
+		for i := 0; i < 8001; i++ {
+			keys[i] = strconv.Itoa(i)
+		}
+
+		err = bf.AddMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+
+		exists, err := bf.ExistsMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+		for _, e := range exists {
+			if !e {
+				t.Error("Key test does not exist")
+			}
+		}
+	})
 }
 
 func TestBloomFilterExistsMultiError(t *testing.T) {

diff --git a/valkeyprob/countingbloomfilter.go b/valkeyprob/countingbloomfilter.go
@@ -44,21 +44,16 @@ local hashIterations = tonumber(ARGV[#ARGV])
 local filterKey = KEYS[1]
 local counterKey = KEYS[2]
 
-local hmgetArgs = {}
-for i=1, numElements do
-    table.insert(hmgetArgs, ARGV[i])
-end
-
-local counts = redis.call('HMGET', filterKey, unpack(hmgetArgs))
 local indexCounter = {}
-for i=1, #counts do
+for i=1, numElements do
 	local index = ARGV[i]
+	local count = redis.call('HGET', filterKey, index)
 
 	if (not indexCounter[index]) then
-		if (not counts[i]) then
+		if (not count) then
 			indexCounter[index] = 0
 		else
-			indexCounter[index] = tonumber(counts[i])
+			indexCounter[index] = tonumber(count)
 		end
 	end
 end

diff --git a/valkeyprob/countingbloomfilter_test.go b/valkeyprob/countingbloomfilter_test.go
@@ -3,6 +3,7 @@ package valkeyprob
 import (
 	"context"
 	"errors"
+	"fmt"
 	"math/rand"
 	"strconv"
 	"testing"
@@ -1095,6 +1096,50 @@ func TestCountingBloomFilterRemoveMulti(t *testing.T) {
 			}
 		}
 	})
+
+	t.Run("remove very large items", func(t *testing.T) {
+		client, flushAllAndClose, err := setup()
+		if err != nil {
+			t.Error(err)
+		}
+		defer func() {
+			err := flushAllAndClose()
+			if err != nil {
+				t.Error(err)
+			}
+		}()
+
+		// Above `LUAI_MAXCSTACK`(8000) limit
+		keys := make([]string, 8001)
+		for i := 0; i < 8001; i++ {
+			keys[i] = fmt.Sprintf("%d", i)
+		}
+
+		bf, err := NewCountingBloomFilter(client, "test", 10000, 0.05)
+		if err != nil {
+			t.Error(err)
+		}
+
+		err = bf.AddMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+
+		err = bf.RemoveMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+
+		exists, err := bf.ExistsMulti(context.Background(), keys)
+		if err != nil {
+			t.Error(err)
+		}
+		for _, e := range exists {
+			if e {
+				t.Error("Key exists")
+			}
+		}
+	})
 }
 
 func TestCountingBloomFilterDelete(t *testing.T) {