feat: simple bench for ingestion (#69)

Author: fengjiachun

ingester-common/src/main/java/io/greptime/common/SPI.java

@@ -32,5 +32,11 @@
 
     String name() default "";
 
+    /**
+     * The priority of the SPI implementation.
+     * If multiple SPI implementations are found, the ones with higher priority will be placed first.
+     *
+     * @return the priority of the SPI implementation
+     */
     int priority() default 0;
 }

ingester-example/src/main/java/io/greptime/BulkWriteApiQuickStart.java

@@ -80,15 +80,15 @@ public static void main(String[] args) throws Exception {
                 .addField("field_json", DataType.Json)
                 .build();
 
-        Config config = Config.newBuilder()
+        Config cfg = Config.newBuilder()
                 .allocatorInitReservation(0)
                 .allocatorMaxAllocation(1024 * 1024 * 1024)
                 .timeoutMsPerMessage(10000)
                 .maxRequestsInFlight(8)
                 .build();
         Context ctx = Context.newDefault().withCompression(Compression.None);
 
-        try (BulkStreamWriter bulkStreamWriter = greptimeDB.bulkStreamWriter(schema, config, ctx)) {
+        try (BulkStreamWriter bulkStreamWriter = greptimeDB.bulkStreamWriter(schema, cfg, ctx)) {
 
             // Write 100 times, each time write 100000 rows
             for (int i = 0; i < 100; i++) {

ingester-example/src/main/java/io/greptime/bench/BulkWriteBenchmark.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright 2023 Greptime Team
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.greptime.bench;
+
+import io.greptime.BulkStreamWriter;
+import io.greptime.BulkWrite;
+import io.greptime.GreptimeDB;
+import io.greptime.common.util.ServiceLoader;
+import io.greptime.common.util.SystemPropertyUtil;
+import io.greptime.models.Table;
+import io.greptime.models.TableSchema;
+import io.greptime.rpc.Compression;
+import io.greptime.rpc.Context;
+import java.util.Iterator;
+import java.util.concurrent.CompletableFuture;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * BulkWriteBenchmark is a benchmark for the bulk write API of GreptimeDB.
+ *
+ * Env:
+ * - db_endpoint: the endpoint of the GreptimeDB server
+ * - db_name: the name of the database
+ * - batch_size_per_request: the batch size per request
+ * - zstd_compression: whether to use zstd compression
+ */
+public class BulkWriteBenchmark {
+
+    private static final Logger LOG = LoggerFactory.getLogger(BulkWriteBenchmark.class);
+
+    public static void main(String[] args) throws Exception {
+        String endpoint = SystemPropertyUtil.get("db_endpoint", "127.0.0.1:4001");
+        String dbName = SystemPropertyUtil.get("db_name", "public");
+        boolean zstdCompression = SystemPropertyUtil.getBool("zstd_compression", false);
+        int batchSize = SystemPropertyUtil.getInt("batch_size_per_request", 100 * 1024);
+        LOG.info("Connect to db: {}, endpoint: {}", dbName, endpoint);
+        LOG.info("Using zstd compression: {}", zstdCompression);
+        LOG.info("Batch size: {}", batchSize);
+
+        GreptimeDB greptimeDB = DBConnector.connectTo(new String[] {endpoint}, dbName);
+        TableDataProvider tableDataProvider =
+                ServiceLoader.load(TableDataProvider.class).first();
+        tableDataProvider.init();
+        TableSchema tableSchema = tableDataProvider.tableSchema();
+
+        BulkWrite.Config cfg = BulkWrite.Config.newBuilder()
+                .allocatorInitReservation(0)
+                .allocatorMaxAllocation(4 * 1024 * 1024 * 1024)
+                .timeoutMsPerMessage(10000)
+                .maxRequestsInFlight(8)
+                .build();
+        Compression compression = zstdCompression ? Compression.Zstd : Compression.None;
+        Context ctx = Context.newDefault().withCompression(compression);
+
+        LOG.info("Start writing data");
+        try (BulkStreamWriter writer = greptimeDB.bulkStreamWriter(tableSchema, cfg, ctx)) {
+            Iterator<Object[]> rows = tableDataProvider.rows();
+
+            long start = System.nanoTime();
+            for (; ; ) {
+                Table.TableBufferRoot table = writer.tableBufferRoot();
+                for (int i = 0; i < batchSize; i++) {
+                    if (!rows.hasNext()) {
+                        break;
+                    }
+                    table.addRow(rows.next());
+                }
+                // Complete the table; adding rows is no longer permitted.
+                table.complete();
+
+                // Write the table data to the server
+                CompletableFuture<Integer> future = writer.writeNext();
+                future.whenComplete((r, t) -> {
+                    if (t != null) {
+                        LOG.error("Error writing data", t);
+                    } else {
+                        LOG.info("Wrote rows: {}", r);
+                    }
+                });
+
+                if (!rows.hasNext()) {
+                    break;
+                }
+            }
+
+            writer.completed();
+
+            LOG.info("Completed writing data, time cost: {}s", (System.nanoTime() - start) / 1000000000);
+        }
+
+        greptimeDB.shutdownGracefully();
+    }
+}

ingester-example/src/main/java/io/greptime/bench/DBConnector.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2023 Greptime Team
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.greptime.bench;
+
+import io.greptime.GreptimeDB;
+import io.greptime.common.util.SerializingExecutor;
+import io.greptime.options.GreptimeOptions;
+
+/**
+ * DBConnector is a helper class to connect to a GreptimeDB instance.
+ */
+public class DBConnector {
+
+    public static GreptimeDB connectTo(String[] endpoints, String dbname) {
+        GreptimeOptions opts = GreptimeOptions.newBuilder(endpoints, dbname)
+                .asyncPool(new SerializingExecutor("bench_async_pool"))
+                .writeMaxRetries(0)
+                .defaultStreamMaxWritePointsPerSecond(Integer.MAX_VALUE)
+                .useZeroCopyWriteInBulkWrite(true)
+                .build();
+        return GreptimeDB.create(opts);
+    }
+}