feat(bench): Add sink bench tool (#14064)

risingwavelabs · Jan 16, 2024 · 84c9a08 · 84c9a08
1 parent b1cdb98
commit 84c9a08
Show file tree

Hide file tree

Showing 12 changed files with 832 additions and 14 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/...-node/risingwave-sink-cassandra/src/main/java/com/risingwave/connector/CassandraSink.java b/...-node/risingwave-sink-cassandra/src/main/java/com/risingwave/connector/CassandraSink.java
@@ -34,6 +34,8 @@
 
 public class CassandraSink extends SinkWriterBase {
     private static final Logger LOG = LoggerFactory.getLogger(CassandraSink.class);
+    private static final Integer MAX_BATCH_SIZE = 1024 * 16;
+
     private final CqlSession session;
     private final List<SinkRow> updateRowCache = new ArrayList<>(1);
     private final HashMap<String, PreparedStatement> stmtMap;
@@ -122,6 +124,7 @@ private void write_append_only(Iterator<SinkRow> rows) {
                             .withDescription("Unknown operation: " + op)
                             .asRuntimeException();
             }
+            tryCommit();
         }
     }
 
@@ -155,6 +158,13 @@ private void write_upsert(Iterator<SinkRow> rows) {
                             .withDescription("Unknown operation: " + op)
                             .asRuntimeException();
             }
+            tryCommit();
+        }
+    }
+
+    private void tryCommit() {
+        if (batchBuilder.getStatementsCount() >= MAX_BATCH_SIZE) {
+            sync();
         }
     }
 

diff --git a/java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/EsSink.java b/java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/EsSink.java
@@ -248,7 +248,7 @@ public BulkListener(RequestTracker requestTracker) {
         /** This method is called just before bulk is executed. */
         @Override
         public void beforeBulk(long executionId, BulkRequest request) {
-            LOG.info("Sending bulk of {} actions to Elasticsearch.", request.numberOfActions());
+            LOG.debug("Sending bulk of {} actions to Elasticsearch.", request.numberOfActions());
         }
 
         /** This method is called after bulk execution. */
@@ -262,7 +262,7 @@ public void afterBulk(long executionId, BulkRequest request, BulkResponse respon
                 this.requestTracker.addErrResult(errMessage);
             } else {
                 this.requestTracker.addOkResult(request.numberOfActions());
-                LOG.info("Sent bulk of {} actions to Elasticsearch.", request.numberOfActions());
+                LOG.debug("Sent bulk of {} actions to Elasticsearch.", request.numberOfActions());
             }
         }
 

diff --git a/src/bench/Cargo.toml b/src/bench/Cargo.toml
@@ -18,6 +18,7 @@ bytes = "1"
 bytesize = { version = "1", features = ["serde"] }
 clap = { version = "4", features = ["derive"] }
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
+futures-async-stream = { workspace = true }
 hdrhistogram = "7"
 itertools = "0.12"
 libc = "0.2"
@@ -26,9 +27,13 @@ parking_lot = "0.12"
 prometheus = { version = "0.13", features = ["process"] }
 rand = "0.8"
 risingwave_common = { workspace = true }
+risingwave_connector = { workspace = true }
+risingwave_pb = { workspace = true }
 risingwave_rt = { workspace = true, optional = true }
 risingwave_storage = { workspace = true }
+risingwave_stream = { workspace = true }
 serde = { version = "1", features = ["derive"] }
+serde_yaml = "0.9"
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "fs",
     "rt",
@@ -53,6 +58,10 @@ nix = { version = "0.27", features = ["fs", "mman"] }
 name = "s3-bench"
 path = "s3_bench/main.rs"
 
+[[bin]]
+name = "sink-bench"
+path = "sink_bench/main.rs"
+
 [features]
 bpf = ["bcc", "risingwave_storage/bpf"]
 trace = ["opentelemetry", "risingwave_rt", "tracing/release_max_level_trace"]