From 9c25a9e35f2029f86c83b08ecd38b2b3f9a87182 Mon Sep 17 00:00:00 2001 From: Hyeonho Kim Date: Tue, 30 Jul 2024 21:57:20 +0900 Subject: [PATCH] fix: too large arguments (#595) --- valkeyprob/bloomfilter.go | 32 +++-------- valkeyprob/bloomfilter_test.go | 76 ++++++++++++++++++++++++++ valkeyprob/countingbloomfilter.go | 13 ++--- valkeyprob/countingbloomfilter_test.go | 45 +++++++++++++++ 4 files changed, 132 insertions(+), 34 deletions(-) diff --git a/valkeyprob/bloomfilter.go b/valkeyprob/bloomfilter.go index 0fcd1f9..2669a85 100644 --- 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 index 03dd0cd..1531bfe 100644 --- 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 index a6cbc72..a0a0ffd 100644 --- 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 index 933240b..87e0f4f 100644 --- 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) {