Improve IncrementalIndex concurrency scalability

benchmarks/src/main/java/org/apache/druid/benchmark/indexing/IndexIngestionMultithreadedBenchmark.java

@@ -0,0 +1,158 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.benchmark.indexing;
+
+import org.apache.druid.benchmark.datagen.BenchmarkDataGenerator;
+import org.apache.druid.benchmark.datagen.BenchmarkSchemaInfo;
+import org.apache.druid.benchmark.datagen.BenchmarkSchemas;
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.hll.HyperLogLogHash;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.aggregation.hyperloglog.HyperUniquesSerde;
+import org.apache.druid.segment.incremental.IncrementalIndex;
+import org.apache.druid.segment.incremental.IncrementalIndexSchema;
+import org.apache.druid.segment.serde.ComplexMetrics;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+import java.util.ArrayList;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 10)
+@Measurement(iterations = 25)
+public class IndexIngestionMultithreadedBenchmark
+{
+  @Param({"75000"})
+  private int rowsPerSegment;
+
+  @Param({"basic"})
+  private String schema;
+
+  @Param({"true", "false"})
+  private boolean rollup;
+
+  private static final Logger log = new Logger(IndexIngestionBenchmark.class);
+  private static final int RNG_SEED = 9999;
+  private IncrementalIndex incIndex;
+  private ArrayList<InputRow> rows;
+  private BenchmarkSchemaInfo schemaInfo;
+  AtomicInteger threadIdAllocator = new AtomicInteger(0);
+
+  @Setup
+  public void setup()
+  {
+    ComplexMetrics.registerSerde("hyperUnique", new HyperUniquesSerde(HyperLogLogHash.getDefault()));
+
+    rows = new ArrayList<InputRow>();
+    schemaInfo = BenchmarkSchemas.SCHEMA_MAP.get(schema);
+
+    BenchmarkDataGenerator gen = new BenchmarkDataGenerator(
+        schemaInfo.getColumnSchemas(),
+        RNG_SEED,
+        schemaInfo.getDataInterval(),
+        rowsPerSegment
+    );
+
+    for (int i = 0; i < rowsPerSegment; i++) {
+      InputRow row = gen.nextRow();
+      if (i % 10000 == 0) {
+        log.info(i + " rows generated.");
+      }
+
+      rows.add(row);
+    }
+  }
+
+  @State(Scope.Thread)
+  public static class ThreadState
+  {
+    int rowOffset = 0;
+
+    @Setup
+    public void setup(IndexIngestionMultithreadedBenchmark globalState)
+    {
+      rowOffset = globalState.threadIdAllocator.getAndIncrement();
+    }
+  }
+
+  @Setup(Level.Invocation)
+  public void setup2()
+  {
+    incIndex = makeIncIndex();
+  }
+
+  private IncrementalIndex makeIncIndex()
+  {
+    return new IncrementalIndex.Builder()
+      .setIndexSchema(
+          new IncrementalIndexSchema.Builder()
+              .withMetrics(schemaInfo.getAggsArray())
+              .withRollup(rollup)
+              .build()
+       )
+       .setReportParseExceptions(false)
+       .setMaxRowCount(rowsPerSegment * 2)
+       .buildOnheap();
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.SingleShotTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public void addRows(Blackhole blackhole, ThreadState threadState) throws Exception
+  {
+    int threads = threadIdAllocator.get();
+    for (int i = threadState.rowOffset; i < rowsPerSegment; i += threads) {
+      InputRow row = rows.get(i);
+      int rv = incIndex.add(row).getRowCount();
+      blackhole.consume(rv);
+    }
+  }
+
+  public static void main(String[] args) throws RunnerException
+  {
+    Options opt = new OptionsBuilder()
+        .include(IndexIngestionMultithreadedBenchmark.class.getSimpleName())
+        .forks(1)
+        .threads(4)
+        .param("schema", "basic")
+        .param("rollup", "false")
+        .build();
+
+    new Runner(opt).run();
+  }
+}