|
35 | 35 | import java.util.Collections; |
36 | 36 | import java.util.Iterator; |
37 | 37 | import java.util.concurrent.CompletableFuture; |
| 38 | +import java.util.concurrent.Semaphore; |
| 39 | +import java.util.concurrent.atomic.AtomicLong; |
38 | 40 | import org.slf4j.Logger; |
39 | 41 | import org.slf4j.LoggerFactory; |
40 | 42 |
|
41 | 43 | /** |
42 | 44 | * BatchingWriteBenchmark is a benchmark for the batching write API of GreptimeDB. |
43 | 45 | * |
44 | 46 | * Env: |
45 | | - * - batch_size_per_request: the batch size per request |
46 | 47 | * - zstd_compression: whether to use zstd compression |
47 | | - * - max_points_per_second: the max number of points that can be written per second, exceeding which may cause blockage |
| 48 | + * - batch_size_per_request: the batch size per request |
| 49 | + * - concurrency: the number of concurrent writers |
48 | 50 | */ |
49 | 51 | public class BatchingWriteBenchmark { |
50 | 52 |
|
51 | 53 | private static final Logger LOG = LoggerFactory.getLogger(BatchingWriteBenchmark.class); |
52 | 54 |
|
53 | 55 | public static void main(String[] args) throws Exception { |
54 | | - boolean zstdCompression = SystemPropertyUtil.getBool("zstd_compression", true); |
| 56 | + boolean zstdCompression = SystemPropertyUtil.getBool("zstd_compression", false); |
55 | 57 | int batchSize = SystemPropertyUtil.getInt("batch_size_per_request", 64 * 1024); |
| 58 | + int concurrency = SystemPropertyUtil.getInt("concurrency", 4); |
56 | 59 |
|
57 | 60 | LOG.info("Using zstd compression: {}", zstdCompression); |
58 | 61 | LOG.info("Batch size: {}", batchSize); |
| 62 | + LOG.info("Concurrency: {}", concurrency); |
| 63 | + |
| 64 | + Compression compression = zstdCompression ? Compression.Zstd : Compression.None; |
| 65 | + Context ctx = Context.newDefault().withCompression(compression); |
59 | 66 |
|
60 | 67 | // Start a metrics exporter |
61 | 68 | MetricsExporter metricsExporter = new MetricsExporter(MetricsUtil.metricRegistry()); |
62 | 69 | metricsExporter.init(ExporterOptions.newDefault()); |
63 | | - GreptimeDB greptimeDB = DBConnector.connect(); |
64 | 70 |
|
65 | | - Compression compression = zstdCompression ? Compression.Zstd : Compression.None; |
66 | | - Context ctx = Context.newDefault().withCompression(compression); |
| 71 | + GreptimeDB greptimeDB = DBConnector.connect(); |
67 | 72 |
|
| 73 | + Semaphore semaphore = new Semaphore(concurrency); |
68 | 74 | TableDataProvider tableDataProvider = |
69 | 75 | ServiceLoader.load(TableDataProvider.class).first(); |
70 | 76 | LOG.info("Table data provider: {}", tableDataProvider.getClass().getName()); |
71 | 77 | tableDataProvider.init(); |
72 | 78 | TableSchema tableSchema = tableDataProvider.tableSchema(); |
73 | | - Iterator<Object[]> rows = tableDataProvider.rows(); |
74 | | - |
75 | | - LOG.info("Start writing data"); |
76 | | - long start = System.nanoTime(); |
77 | | - do { |
78 | | - Table table = Table.from(tableSchema); |
79 | | - for (int i = 0; i < batchSize; i++) { |
80 | | - if (!rows.hasNext()) { |
81 | | - break; |
| 79 | + AtomicLong totalRowsWritten = new AtomicLong(0); |
| 80 | + |
| 81 | + try { |
| 82 | + Iterator<Object[]> rows = tableDataProvider.rows(); |
| 83 | + |
| 84 | + LOG.info( |
| 85 | + "Start writing data, table: {}, row count: {}", |
| 86 | + tableSchema.getTableName(), |
| 87 | + tableDataProvider.rowCount()); |
| 88 | + |
| 89 | + long start = System.nanoTime(); |
| 90 | + do { |
| 91 | + Table table = Table.from(tableSchema); |
| 92 | + for (int j = 0; j < batchSize; j++) { |
| 93 | + if (!rows.hasNext()) { |
| 94 | + break; |
| 95 | + } |
| 96 | + table.addRow(rows.next()); |
82 | 97 | } |
83 | | - table.addRow(rows.next()); |
84 | | - } |
85 | | - LOG.info("Table bytes used: {}", table.bytesUsed()); |
86 | | - // Complete the table; adding rows is no longer permitted. |
87 | | - table.complete(); |
88 | | - long fStart = System.nanoTime(); |
89 | | - // Write the table data to the server |
90 | | - CompletableFuture<Result<WriteOk, Err>> future = |
91 | | - greptimeDB.write(Collections.singletonList(table), WriteOp.Insert, ctx); |
92 | | - // Wait for the writing to complete |
93 | | - int numRows = future.get().mapOr(0, WriteOk::getSuccess); |
94 | | - long costMs = (System.nanoTime() - fStart) / 1000000; |
95 | | - LOG.info("Write rows: {}, time cost: {}ms", numRows, costMs); |
96 | | - |
97 | | - } while (rows.hasNext()); |
98 | | - |
99 | | - LOG.info("Completed writing data, time cost: {}s", (System.nanoTime() - start) / 1000000000); |
100 | | - |
101 | | - greptimeDB.shutdownGracefully(); |
102 | | - tableDataProvider.close(); |
103 | | - metricsExporter.shutdownGracefully(); |
| 98 | + |
| 99 | + // Complete the table; adding rows is no longer permitted. |
| 100 | + table.complete(); |
| 101 | + |
| 102 | + semaphore.acquire(); |
| 103 | + |
| 104 | + // Write the table data to the server |
| 105 | + CompletableFuture<Result<WriteOk, Err>> future = |
| 106 | + greptimeDB.write(Collections.singletonList(table), WriteOp.Insert, ctx); |
| 107 | + long fStart = System.nanoTime(); |
| 108 | + future.whenComplete((result, error) -> { |
| 109 | + semaphore.release(); |
| 110 | + |
| 111 | + long costMs = (System.nanoTime() - fStart) / 1000000; |
| 112 | + if (error != null) { |
| 113 | + LOG.error("Error writing data", error); |
| 114 | + return; |
| 115 | + } |
| 116 | + |
| 117 | + int numRows = result.mapOr(0, writeOk -> writeOk.getSuccess()); |
| 118 | + long totalRows = totalRowsWritten.addAndGet(numRows); |
| 119 | + long totalElapsedSec = (System.nanoTime() - start) / 1000000000; |
| 120 | + long writeRatePerSecond = totalElapsedSec > 0 ? totalRows / totalElapsedSec : 0; |
| 121 | + LOG.info( |
| 122 | + "Wrote rows: {}, time cost: {}ms, total rows: {}, total elapsed: {}s, write rate: {} rows/sec", |
| 123 | + numRows, |
| 124 | + costMs, |
| 125 | + totalRows, |
| 126 | + totalElapsedSec, |
| 127 | + writeRatePerSecond); |
| 128 | + }); |
| 129 | + } while (rows.hasNext()); |
| 130 | + |
| 131 | + // Wait for all the requests to complete |
| 132 | + semaphore.acquire(concurrency); |
| 133 | + |
| 134 | + LOG.info("Completed writing data, time cost: {}s", (System.nanoTime() - start) / 1000000000); |
| 135 | + |
| 136 | + } finally { |
| 137 | + greptimeDB.shutdownGracefully(); |
| 138 | + metricsExporter.shutdownGracefully(); |
| 139 | + } |
104 | 140 | } |
105 | 141 | } |
0 commit comments