Support multi-threaded writes in pull based ingestion

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

Diff for: CHANGELOG.md

@@ -38,6 +38,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Star Tree] [Search] Add query changes to support unsigned-long in star tree ([#17275](https://github.com/opensearch-project/OpenSearch/pull/17275))
 - Add `ApproximateMatchAllQuery` that targets match_all queries and approximates sorts ([#17772](https://github.com/opensearch-project/OpenSearch/pull/17772))
 - Add TermsQuery support to Search GRPC endpoint ([#17888](https://github.com/opensearch-project/OpenSearch/pull/17888))
+- Add multi-threaded writer support in pull-based ingestion ([#17912](https://github.com/opensearch-project/OpenSearch/pull/17912))
 
 ### Changed
 - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912))

Diff for: plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/IngestFromKafkaIT.java

@@ -200,4 +200,31 @@ public void testUpdateWithoutIDField() throws Exception {
             return 30 == (Integer) response.getHits().getHits()[0].getSourceAsMap().get("age");
         });
     }
+
+    public void testMultiThreadedWrites() throws Exception {
+        // create index with 5 writer threads
+        createIndexWithDefaultSettings(indexName, 1, 0, 5);
+        ensureGreen(indexName);
+
+        // Step 1: Produce messages
+        for (int i = 0; i < 1000; i++) {
+            produceData(Integer.toString(i), "name" + i, "25");
+        }
+
+        waitForState(() -> {
+            SearchResponse searchableDocsResponse = client().prepareSearch(indexName).setSize(2000).setPreference("_only_local").get();
+            return searchableDocsResponse.getHits().getTotalHits().value() == 1000;
+        });
+
+        // Step 2: Produce an update message and validate
+        for (int i = 0; i < 1000; i++) {
+            produceData(Integer.toString(i), "name" + i, "30");
+        }
+
+        waitForState(() -> {
+            RangeQueryBuilder query = new RangeQueryBuilder("age").gte(28);
+            SearchResponse response = client().prepareSearch(indexName).setQuery(query).get();
+            return response.getHits().getTotalHits().value() == 1000;
+        });
+    }
 }

Diff for: plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/KafkaIngestionBaseIT.java

@@ -164,10 +164,10 @@ protected ResumeIngestionResponse resumeIngestion(String indexName) throws Execu
     }
 
     protected void createIndexWithDefaultSettings(int numShards, int numReplicas) {
-        createIndexWithDefaultSettings(indexName, numShards, numReplicas);
+        createIndexWithDefaultSettings(indexName, numShards, numReplicas, 1);
     }
 
-    protected void createIndexWithDefaultSettings(String indexName, int numShards, int numReplicas) {
+    protected void createIndexWithDefaultSettings(String indexName, int numShards, int numReplicas, int numProcessorThreads) {
         createIndex(
             indexName,
             Settings.builder()
@@ -178,6 +178,7 @@ protected void createIndexWithDefaultSettings(String indexName, int numShards, i
                 .put("ingestion_source.param.topic", topicName)
                 .put("ingestion_source.param.bootstrap_servers", kafka.getBootstrapServers())
                 .put("index.replication.type", "SEGMENT")
+                .put("ingestion_source.num_processor_threads", numProcessorThreads)
                 // set custom kafka consumer properties
                 .put("ingestion_source.param.fetch.min.bytes", 30000)
                 .put("ingestion_source.param.enable.auto.commit", false)

Diff for: plugins/ingestion-kafka/src/internalClusterTest/java/org/opensearch/plugin/kafka/RemoteStoreKafkaIT.java

@@ -246,8 +246,8 @@ public void testPaginatedGetIngestionState() throws ExecutionException, Interrup
         internalCluster().startClusterManagerOnlyNode();
         internalCluster().startDataOnlyNode();
         internalCluster().startDataOnlyNode();
-        createIndexWithDefaultSettings("index1", 5, 0);
-        createIndexWithDefaultSettings("index2", 5, 0);
+        createIndexWithDefaultSettings("index1", 5, 0, 1);
+        createIndexWithDefaultSettings("index2", 5, 0, 1);
         ensureGreen("index1");
         ensureGreen("index2");
 

Diff for: server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java

@@ -831,6 +831,18 @@ public Iterator<Setting<?>> settings() {
         Property.Dynamic
     );
 
+    /**
+     * Defines the number of processor threads that will write to the lucene index.
+     */
+    public static final String SETTING_INGESTION_SOURCE_NUM_PROCESSOR_THREADS = "index.ingestion_source.num_processor_threads";
+    public static final Setting<Integer> INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING = Setting.intSetting(
+        SETTING_INGESTION_SOURCE_NUM_PROCESSOR_THREADS,
+        1,
+        1,
+        Setting.Property.IndexScope,
+        Setting.Property.Final
+    );
+
     public static final Setting.AffixSetting<Object> INGESTION_SOURCE_PARAMS_SETTING = Setting.prefixKeySetting(
         "index.ingestion_source.param.",
         key -> new Setting<>(key, "", (value) -> {
@@ -1073,11 +1085,14 @@ public IngestionSource getIngestionSource() {
             final Map<String, Object> ingestionSourceParams = INGESTION_SOURCE_PARAMS_SETTING.getAsMap(settings);
             final long maxPollSize = INGESTION_SOURCE_MAX_POLL_SIZE.get(settings);
             final int pollTimeout = INGESTION_SOURCE_POLL_TIMEOUT.get(settings);
+            final int numProcessorThreads = INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING.get(settings);
+
             return new IngestionSource.Builder(ingestionSourceType).setParams(ingestionSourceParams)
                 .setPointerInitReset(pointerInitReset)
                 .setErrorStrategy(errorStrategy)
                 .setMaxPollSize(maxPollSize)
                 .setPollTimeout(pollTimeout)
+                .setNumProcessorThreads(numProcessorThreads)
                 .build();
         }
         return null;

Diff for: server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java

@@ -18,6 +18,7 @@
 import java.util.Objects;
 
 import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_MAX_POLL_SIZE;
+import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING;
 import static org.opensearch.cluster.metadata.IndexMetadata.INGESTION_SOURCE_POLL_TIMEOUT;
 
 /**
@@ -31,21 +32,24 @@ public class IngestionSource {
     private final Map<String, Object> params;
     private final long maxPollSize;
     private final int pollTimeout;
+    private int numProcessorThreads;
 
     private IngestionSource(
         String type,
         PointerInitReset pointerInitReset,
         IngestionErrorStrategy.ErrorStrategy errorStrategy,
         Map<String, Object> params,
         long maxPollSize,
-        int pollTimeout
+        int pollTimeout,
+        int numProcessorThreads
     ) {
         this.type = type;
         this.pointerInitReset = pointerInitReset;
         this.params = params;
         this.errorStrategy = errorStrategy;
         this.maxPollSize = maxPollSize;
         this.pollTimeout = pollTimeout;
+        this.numProcessorThreads = numProcessorThreads;
     }
 
     public String getType() {
@@ -72,6 +76,10 @@ public int getPollTimeout() {
         return pollTimeout;
     }
 
+    public int getNumProcessorThreads() {
+        return numProcessorThreads;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -82,12 +90,13 @@ public boolean equals(Object o) {
             && Objects.equals(errorStrategy, ingestionSource.errorStrategy)
             && Objects.equals(params, ingestionSource.params)
             && Objects.equals(maxPollSize, ingestionSource.maxPollSize)
-            && Objects.equals(pollTimeout, ingestionSource.pollTimeout);
+            && Objects.equals(pollTimeout, ingestionSource.pollTimeout)
+            && Objects.equals(numProcessorThreads, ingestionSource.numProcessorThreads);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(type, pointerInitReset, params, errorStrategy, maxPollSize, pollTimeout);
+        return Objects.hash(type, pointerInitReset, params, errorStrategy, maxPollSize, pollTimeout, numProcessorThreads);
     }
 
     @Override
@@ -108,6 +117,8 @@ public String toString() {
             + maxPollSize
             + ", pollTimeout="
             + pollTimeout
+            + ", numProcessorThreads="
+            + numProcessorThreads
             + '}';
     }
 
@@ -163,6 +174,7 @@ public static class Builder {
         private Map<String, Object> params;
         private long maxPollSize = INGESTION_SOURCE_MAX_POLL_SIZE.getDefault(Settings.EMPTY);
         private int pollTimeout = INGESTION_SOURCE_POLL_TIMEOUT.getDefault(Settings.EMPTY);
+        private int numProcessorThreads = INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING.getDefault(Settings.EMPTY);
 
         public Builder(String type) {
             this.type = type;
@@ -206,8 +218,13 @@ public Builder setPollTimeout(int pollTimeout) {
             return this;
         }
 
+        public Builder setNumProcessorThreads(int numProcessorThreads) {
+            this.numProcessorThreads = numProcessorThreads;
+            return this;
+        }
+
         public IngestionSource build() {
-            return new IngestionSource(type, pointerInitReset, errorStrategy, params, maxPollSize, pollTimeout);
+            return new IngestionSource(type, pointerInitReset, errorStrategy, params, maxPollSize, pollTimeout, numProcessorThreads);
         }
 
     }

Diff for: server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java

@@ -269,6 +269,9 @@ public final class IndexScopedSettings extends AbstractScopedSettings {
                 IndexMetadata.INGESTION_SOURCE_POINTER_INIT_RESET_VALUE_SETTING,
                 IndexMetadata.INGESTION_SOURCE_PARAMS_SETTING,
                 IndexMetadata.INGESTION_SOURCE_ERROR_STRATEGY_SETTING,
+                IndexMetadata.INGESTION_SOURCE_MAX_POLL_SIZE,
+                IndexMetadata.INGESTION_SOURCE_POLL_TIMEOUT,
+                IndexMetadata.INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING,
 
                 // validate that built-in similarities don't get redefined
                 Setting.groupSetting("index.similarity.", (s) -> {

Diff for: server/src/main/java/org/opensearch/index/engine/IngestionEngine.java

@@ -123,7 +123,8 @@ public void start() {
             ingestionErrorStrategy,
             initialPollerState,
             ingestionSource.getMaxPollSize(),
-            ingestionSource.getPollTimeout()
+            ingestionSource.getPollTimeout(),
+            ingestionSource.getNumProcessorThreads()
         );
         streamPoller.start();
     }

Diff for: server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java

@@ -17,12 +17,11 @@
 import org.opensearch.index.Message;
 import org.opensearch.index.engine.IngestionEngine;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
@@ -47,8 +46,6 @@ public class DefaultStreamPoller implements StreamPoller {
 
     private ExecutorService consumerThread;
 
-    private ExecutorService processorThread;
-
     // start of the batch, inclusive
     private IngestionShardPointer initialBatchStartPointer;
     private boolean includeBatchStartPointer = false;
@@ -58,16 +55,14 @@ public class DefaultStreamPoller implements StreamPoller {
 
     private Set<IngestionShardPointer> persistedPointers;
 
-    private BlockingQueue<IngestionShardConsumer.ReadResult<? extends IngestionShardPointer, ? extends Message>> blockingQueue;
-
-    private MessageProcessorRunnable processorRunnable;
-
     private final CounterMetric totalPolledCount = new CounterMetric();
 
     // A pointer to the max persisted pointer for optimizing the check
     @Nullable
     private IngestionShardPointer maxPersistedPointer;
 
+    private PartitionedBlockingQueueContainer blockingQueueContainer;
+
     public DefaultStreamPoller(
         IngestionShardPointer startPointer,
         Set<IngestionShardPointer> persistedPointers,
@@ -78,13 +73,14 @@ public DefaultStreamPoller(
         IngestionErrorStrategy errorStrategy,
         State initialState,
         long maxPollSize,
-        int pollTimeout
+        int pollTimeout,
+        int numProcessorThreads
     ) {
         this(
             startPointer,
             persistedPointers,
             consumer,
-            new MessageProcessorRunnable(new ArrayBlockingQueue<>(100), ingestionEngine, errorStrategy),
+            new PartitionedBlockingQueueContainer(numProcessorThreads, consumer.getShardId(), ingestionEngine, errorStrategy),
             resetState,
             resetValue,
             errorStrategy,
@@ -96,7 +92,7 @@ public DefaultStreamPoller(
         IngestionShardPointer startPointer,
         Set<IngestionShardPointer> persistedPointers,
         IngestionShardConsumer consumer,
-        MessageProcessorRunnable processorRunnable,
+        PartitionedBlockingQueueContainer blockingQueueContainer,
         ResetState resetState,
         String resetValue,
         IngestionErrorStrategy errorStrategy,
@@ -111,22 +107,13 @@ public DefaultStreamPoller(
         if (!this.persistedPointers.isEmpty()) {
             maxPersistedPointer = this.persistedPointers.stream().max(IngestionShardPointer::compareTo).get();
         }
-        this.processorRunnable = processorRunnable;
-        blockingQueue = processorRunnable.getBlockingQueue();
+        this.blockingQueueContainer = blockingQueueContainer;
         this.consumerThread = Executors.newSingleThreadExecutor(
             r -> new Thread(
                 r,
                 String.format(Locale.ROOT, "stream-poller-consumer-%d-%d", consumer.getShardId(), System.currentTimeMillis())
             )
         );
-
-        // TODO: allow multiple threads for processing the messages in parallel
-        this.processorThread = Executors.newSingleThreadExecutor(
-            r -> new Thread(
-                r,
-                String.format(Locale.ROOT, "stream-poller-processor-%d-%d", consumer.getShardId(), System.currentTimeMillis())
-            )
-        );
         this.errorStrategy = errorStrategy;
     }
 
@@ -144,7 +131,7 @@ public void start() {
         // when we start, we need to include the batch start pointer in the read for the first read
         includeBatchStartPointer = true;
         consumerThread.submit(this::startPoll);
-        processorThread.submit(processorRunnable);
+        blockingQueueContainer.startProcessorThreads();
     }
 
     /**
@@ -242,7 +229,7 @@ private void processRecords(List<IngestionShardConsumer.ReadResult<? extends Ing
                     continue;
                 }
                 totalPolledCount.inc();
-                blockingQueue.put(result);
+                blockingQueueContainer.add(result);
 
                 logger.debug(
                     "Put message {} with pointer {} to the blocking queue",
@@ -322,10 +309,8 @@ public void close() {
                 logger.error("Error in closing the poller of shard {}: {}", consumer.getShardId(), e);
             }
         }
-        blockingQueue.clear();
         consumerThread.shutdown();
-        // interrupts the processor
-        processorThread.shutdownNow();
+        blockingQueueContainer.close();
         logger.info("closed the poller of shard {}", consumer.getShardId());
     }
 
@@ -343,19 +328,23 @@ public boolean isClosed() {
      * Returns the batch start pointer from where the poller can resume in case of shard recovery. The poller and
      * processor are decoupled in this implementation, and hence the latest pointer tracked by the processor acts as the
      * recovery/start point. In case the processor has not started tracking, then the initial batchStartPointer used by
-     * the poller acts as the start point.
+     * the poller acts as the start point. If multiple processor threads are used, the minimum shard pointer across
+     * processors indicates the start point.
      */
     @Override
     public IngestionShardPointer getBatchStartPointer() {
-        IngestionShardPointer currentShardPointer = processorRunnable.getCurrentShardPointer();
-        return currentShardPointer == null ? initialBatchStartPointer : currentShardPointer;
+        return blockingQueueContainer.getCurrentShardPointers()
+            .stream()
+            .filter(Objects::nonNull)
+            .min(Comparator.naturalOrder())
+            .orElseGet(() -> initialBatchStartPointer);
     }
 
     @Override
     public PollingIngestStats getStats() {
         PollingIngestStats.Builder builder = new PollingIngestStats.Builder();
         builder.setTotalPolledCount(totalPolledCount.count());
-        builder.setTotalProcessedCount(processorRunnable.getStats().count());
+        builder.setTotalProcessedCount(blockingQueueContainer.getTotalProcessedCount());
         return builder.build();
     }
 
@@ -371,6 +360,6 @@ public IngestionErrorStrategy getErrorStrategy() {
     @Override
     public void updateErrorStrategy(IngestionErrorStrategy errorStrategy) {
         this.errorStrategy = errorStrategy;
-        processorRunnable.setErrorStrategy(errorStrategy);
+        blockingQueueContainer.updateErrorStrategy(errorStrategy);
     }
 }