From a24812a60ab5a04df12bee10cff1ddb194592eca Mon Sep 17 00:00:00 2001 From: d0zingcat Date: Wed, 24 Jul 2024 12:18:56 +0800 Subject: [PATCH] add samples args Signed-off-by: d0zingcat --- rks/main.py | 4 +++- rks/redis_utils.py | 16 ++++++++++------ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/rks/main.py b/rks/main.py index 6d2d006..8741335 100644 --- a/rks/main.py +++ b/rks/main.py @@ -12,6 +12,8 @@ def main(): parser.add_argument('--port', required=True, help='Redis port') parser.add_argument('--username', default=None, help='Redis username') parser.add_argument('--password', default=None, help='Redis password') + parser.add_argument('--samples', default=5, + help='Redis memory usage samples') parser.add_argument('--cluster', action='store_true', help='Enable cluster mode') parser.add_argument('--batch_size', type=int, default=1000, @@ -91,7 +93,7 @@ def main(): if db.startswith('db'): db_num = int(db[2:]) get_redis_keys(r, args.batch_size, - db_num, args.pretty_format) + db_num, args.pretty_format, args.samples) r.close() diff --git a/rks/redis_utils.py b/rks/redis_utils.py index 6d52296..fc2ff02 100644 --- a/rks/redis_utils.py +++ b/rks/redis_utils.py @@ -4,7 +4,7 @@ from .display import analyze_redis_keys from .analysis import update_topk_heap, update_statistics -def get_redis_keys(r, batch_size, db_num, use_pretty): +def get_redis_keys(r, batch_size, db_num, use_pretty, samples): r.execute_command('SELECT', db_num) @@ -25,7 +25,7 @@ def get_redis_keys(r, batch_size, db_num, use_pretty): for i, key in ipairs(scanResult[2]) do local key_type = redis.call('TYPE', key)['ok'] - local memory = redis.call('MEMORY', 'USAGE', key) + local memory = redis.call('MEMORY', 'USAGE', key, 'SAMPLES', SAMPLES_PLACEHOLDER) local ttl = redis.call('TTL', key) table.insert(result, {key, key_type, memory, ttl}) @@ -35,6 +35,7 @@ def get_redis_keys(r, batch_size, db_num, use_pretty): """ script = script.replace("BATCH_SIZE_PLACEHOLDER", str(batch_size)) + script = script.replace("SAMPLES_PLACEHOLDER", str(samples)) cursor = b'0' while True: @@ -78,7 +79,8 @@ def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty): masters = [] for node in nodes: if 'master' in node['flags']: - master_dict = {'master': node['id'], 'slots': node['slots'], 'slaves': []} + master_dict = {'master': node['id'], + 'slots': node['slots'], 'slaves': []} masters.append(master_dict.copy()) for node in nodes: @@ -93,9 +95,11 @@ def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty): for master in masters: if slave_flag is False: - node = next((node for node in nodes if node['id'] == master['master']), None) + node = next( + (node for node in nodes if node['id'] == master['master']), None) else: - node = next((node for node in nodes if node['id'] == master['slaves'][0]), None) + node = next( + (node for node in nodes if node['id'] == master['slaves'][0]), None) r = redis.Redis(host=node['host'], port=node['port']) r.execute_command('READONLY') @@ -144,4 +148,4 @@ def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty): r.close() - return analyze_redis_keys(min_heap, prefix_statistics_map, total_key_size, total_key_count, key_count_by_type, 0, use_pretty) \ No newline at end of file + return analyze_redis_keys(min_heap, prefix_statistics_map, total_key_size, total_key_count, key_count_by_type, 0, use_pretty)