From df2906672baa04d7383e3f93162814aefd60e86e Mon Sep 17 00:00:00 2001 From: proost Date: Tue, 30 Jul 2024 15:13:49 +0900 Subject: [PATCH] fix: too large arguments --- rueidisprob/bloomfilter.go | 32 +++-------- rueidisprob/bloomfilter_test.go | 76 +++++++++++++++++++++++++ rueidisprob/countingbloomfilter.go | 16 ++---- rueidisprob/countingbloomfilter_test.go | 48 +++++++++++++++- 4 files changed, 136 insertions(+), 36 deletions(-) diff --git a/rueidisprob/bloomfilter.go b/rueidisprob/bloomfilter.go index 3896d555..2a30146e 100644 --- a/rueidisprob/bloomfilter.go +++ b/rueidisprob/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/rueidisprob/bloomfilter_test.go b/rueidisprob/bloomfilter_test.go index cc59c8ac..cd1d341e 100644 --- a/rueidisprob/bloomfilter_test.go +++ b/rueidisprob/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/rueidisprob/countingbloomfilter.go b/rueidisprob/countingbloomfilter.go index f9500ffc..4e02b934 100644 --- a/rueidisprob/countingbloomfilter.go +++ b/rueidisprob/countingbloomfilter.go @@ -3,9 +3,10 @@ package rueidisprob import ( "context" "errors" - "github.com/redis/rueidis" "math" "strconv" + + "github.com/redis/rueidis" ) var ( @@ -43,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/rueidisprob/countingbloomfilter_test.go b/rueidisprob/countingbloomfilter_test.go index 0888e4c8..b8937cf3 100644 --- a/rueidisprob/countingbloomfilter_test.go +++ b/rueidisprob/countingbloomfilter_test.go @@ -3,10 +3,12 @@ package rueidisprob import ( "context" "errors" - "github.com/redis/rueidis" + "fmt" "math/rand" "strconv" "testing" + + "github.com/redis/rueidis" ) func TestNewCountingBloomFilter(t *testing.T) { @@ -1094,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) {