From f7a57c637e54896eeccd8eeed92442f124ada2c9 Mon Sep 17 00:00:00 2001
From: Yu Ning <78631860+Chillax-0v0@users.noreply.github.com>
Date: Wed, 21 Feb 2024 16:52:21 +0800
Subject: [PATCH] feat(s3stream/wal/benchmark): make iops configurable (#940)

Signed-off-by: Ning Yu <ningyu@automq.com>
---
 .../com/automq/stream/s3/wal/benchmark/WriteBench.java    | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java
index 1a07a994d..201eaf1ec 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/benchmark/WriteBench.java
@@ -54,6 +54,9 @@ public WriteBench(Config config) throws IOException {
         if (config.depth != null) {
             builder.ioThreadNums(config.depth);
         }
+        if (config.iops != null) {
+            builder.writeRateLimit(config.iops);
+        }
         this.log = builder.build();
         this.log.start();
         recoverAndReset(this.log);
@@ -194,6 +197,7 @@ static class Config {
         final String path;
         final Long capacity;
         final Integer depth;
+        final Integer iops;
 
         // following fields are benchmark configuration
         final Integer threads;
@@ -205,6 +209,7 @@ static class Config {
             this.path = ns.getString("path");
             this.capacity = ns.getLong("capacity");
             this.depth = ns.getInt("depth");
+            this.iops = ns.getInt("iops");
             this.threads = ns.getInt("threads");
             this.throughputBytes = ns.getInt("throughput");
             this.recordSizeBytes = ns.getInt("recordSize");
@@ -227,6 +232,9 @@ static ArgumentParser parser() {
             parser.addArgument("-d", "--depth")
                 .type(Integer.class)
                 .help("IO depth of the WAL");
+            parser.addArgument("--iops")
+                .type(Integer.class)
+                .help("IOPS of the WAL");
             parser.addArgument("--threads")
                 .type(Integer.class)
                 .setDefault(1)