chore: collection size (#3844)

tests/dragonfly/seeder/__init__.py

@@ -79,6 +79,7 @@ class StaticSeeder(SeederBase):
         data_size=100,
         variance=5,
         samples=10,
+        collection_size=None,
         types: typing.Optional[typing.List[str]] = None,
     ):
         SeederBase.__init__(self, types)
@@ -87,6 +88,11 @@ class StaticSeeder(SeederBase):
         self.variance = variance
         self.samples = samples
 
+        if collection_size is None:
+            self.collection_size = data_size ** (1 / 3)
+        else:
+            self.collection_size = collection_size
+
     async def run(self, client: aioredis.Redis):
         """Run with specified options until key_target is met"""
         samples = [
@@ -105,9 +111,9 @@ class StaticSeeder(SeederBase):
             dsize = random.uniform(self.data_size / self.variance, self.data_size * self.variance)
             csize = 1
         else:
-            dsize = self.data_size ** (1 / 3)
-            dsize = random.uniform(dsize / self.variance, dsize * self.variance)
-            csize = int(self.data_size // dsize) + 1
+            csize = self.collection_size
+            csize = math.ceil(random.uniform(csize / self.variance, csize * self.variance))
+            dsize = self.data_size // csize
 
         args = ["DEBUG", "POPULATE", key_target, prefix, math.ceil(dsize)]
         args += ["RAND", "TYPE", dtype, "ELEMENTS", csize]
@@ -124,14 +130,26 @@ class Seeder(SeederBase):
 
     units: typing.List[Unit]
 
-    def __init__(self, units=10, key_target=10_000, data_size=100):
-        SeederBase.__init__(self)
+    def __init__(
+        self,
+        units=10,
+        key_target=10_000,
+        data_size=100,
+        collection_size=None,
+        types: typing.Optional[typing.List[str]] = None,
+    ):
+        SeederBase.__init__(self, types)
         self.key_target = key_target
         self.data_size = data_size
+        if collection_size is None:
+            self.collection_size = math.ceil(data_size ** (1 / 3))
+        else:
+            self.collection_size = collection_size
+
         self.units = [
             Seeder.Unit(
                 prefix=f"k-s{self.uid}u{i}-",
-                type=Seeder.DEFAULT_TYPES[i % len(Seeder.DEFAULT_TYPES)],
+                type=self.types[i % len(self.types)],
                 counter=0,
                 stop_key=f"_s{self.uid}u{i}-stop",
             )
@@ -147,6 +165,7 @@ class Seeder(SeederBase):
             target_ops if target_ops is not None else 0,
             target_deviation if target_deviation is not None else -1,
             self.data_size,
+            self.collection_size,
         ]
 
         sha = await client.script_load(Seeder._load_script("generate"))

tests/dragonfly/seeder/script-generate.lua

@@ -17,12 +17,13 @@ local key_target = tonumber(ARGV[5])
 local total_ops = tonumber(ARGV[6])
 local min_dev = tonumber(ARGV[7])
 local data_size = tonumber(ARGV[8])
+local collection_size = tonumber(ARGV[9])
 
 -- collect all keys belonging to this script
 -- assumes exclusive ownership
 local keys = LU_collect_keys(prefix, type)
 
-LG_funcs.init(data_size)
+LG_funcs.init(data_size, collection_size)
 local addfunc = LG_funcs['add_' .. string.lower(type)]
 local modfunc = LG_funcs['mod_' .. string.lower(type)]
 

tests/dragonfly/seeder/script-genlib.lua

@@ -1,9 +1,9 @@
 local LG_funcs = {}
 
-function LG_funcs.init(dsize)
+function LG_funcs.init(dsize, csize)
     LG_funcs.dsize = dsize
-    LG_funcs.csize = math.floor(dsize ^ (2/3))
-    LG_funcs.esize = math.ceil(dsize ^ (1/3))
+    LG_funcs.csize = csize
+    LG_funcs.esize = math.ceil(dsize / csize)
 end
 
 -- strings

tests/dragonfly/seeder_test.py

@@ -14,6 +14,26 @@ async def test_static_seeder(async_client: aioredis.Redis):
     assert abs(await async_client.dbsize() - 10_000) <= 50
 
 
+@dfly_args({"proactor_threads": 4})
+async def test_static_collection_size(async_client: aioredis.Redis):
+    async def check_list():
+        keys = await async_client.keys()
+        assert (await async_client.llen(keys[0])) == 1
+        assert len(await async_client.lpop(keys[0])) == 10_000
+
+    s = StaticSeeder(
+        key_target=10, data_size=10_000, variance=1, samples=1, collection_size=1, types=["LIST"]
+    )
+    await s.run(async_client)
+    await check_list()
+
+    await async_client.flushall()
+
+    s = Seeder(units=1, key_target=10, data_size=10_000, collection_size=1, types=["LIST"])
+    await s.run(async_client)
+    await check_list()
+
+
 @dfly_args({"proactor_threads": 4})
 async def test_seeder_key_target(async_client: aioredis.Redis):
     """Ensure seeder reaches its key targets"""