diff --git a/tests/dragonfly/seeder/__init__.py b/tests/dragonfly/seeder/__init__.py index 231afe9ed..aeea29c40 100644 --- a/tests/dragonfly/seeder/__init__.py +++ b/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")) diff --git a/tests/dragonfly/seeder/script-generate.lua b/tests/dragonfly/seeder/script-generate.lua index 5f7ca33f9..d1f818425 100644 --- a/tests/dragonfly/seeder/script-generate.lua +++ b/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)] diff --git a/tests/dragonfly/seeder/script-genlib.lua b/tests/dragonfly/seeder/script-genlib.lua index 7cbc76c5e..937f26fe4 100644 --- a/tests/dragonfly/seeder/script-genlib.lua +++ b/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 diff --git a/tests/dragonfly/seeder_test.py b/tests/dragonfly/seeder_test.py index 61eba28ce..5086cc399 100644 --- a/tests/dragonfly/seeder_test.py +++ b/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"""