Support multiple processor threads in pull-based ingestion

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

Author: varunbharadwaj

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Security Manager Replacement] Create initial Java Agent to intercept Socket::connect calls ([#17724](https://github.com/opensearch-project/OpenSearch/pull/17724))
 - Add ingestion management APIs for pause, resume and get ingestion state ([#17631](https://github.com/opensearch-project/OpenSearch/pull/17631))
 - [Security Manager Replacement] Enhance Java Agent to intercept System::exit ([#17746](https://github.com/opensearch-project/OpenSearch/pull/17746))
+- Support multi-threaded writes in pull-based ingestion  ([#17771](https://github.com/opensearch-project/OpenSearch/pull/17771))
 
 ### Changed
 - Migrate BC libs to their FIPS counterparts ([#14912](https://github.com/opensearch-project/OpenSearch/pull/14912))

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

@@ -135,4 +135,37 @@ public void testCloseIndex() throws Exception {
         ensureGreen(indexName);
         client().admin().indices().close(Requests.closeIndexRequest(indexName)).get();
     }
+
+    public void testKafkaIngestionWithMultipleProcessorThreads() {
+        for (int i = 1; i <= 100; i++) {
+            produceData(String.valueOf(i), "name" + i, "25");
+        }
+
+        createIndex(
+            indexName,
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                .put("ingestion_source.type", "kafka")
+                .put("ingestion_source.pointer.init.reset", "earliest")
+                .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", 3)
+                .build(),
+            "{\"properties\":{\"name\":{\"type\": \"text\"},\"age\":{\"type\": \"integer\"}}}}"
+        );
+
+        RangeQueryBuilder query = new RangeQueryBuilder("age").gte(21);
+        await().atMost(60, TimeUnit.SECONDS).untilAsserted(() -> {
+            refresh(indexName);
+            SearchResponse response = client().prepareSearch(indexName).setQuery(query).get();
+            assertThat(response.getHits().getTotalHits().value(), is(100L));
+            PollingIngestStats stats = client().admin().indices().prepareStats(indexName).get().getIndex(indexName).getShards()[0]
+                .getPollingIngestStats();
+            assertNotNull(stats);
+            assertThat(stats.getMessageProcessorStats().getTotalProcessedCount(), is(100L));
+            assertThat(stats.getConsumerStats().getTotalPolledCount(), is(100L));
+        });
+    }
 }

server/src/main/java/org/opensearch/cluster/metadata/IndexMetadata.java

@@ -784,6 +784,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) -> {
@@ -1023,8 +1035,9 @@ public IngestionSource getIngestionSource() {
             );
 
             final IngestionErrorStrategy.ErrorStrategy errorStrategy = INGESTION_SOURCE_ERROR_STRATEGY_SETTING.get(settings);
+            final int numProcessorThreads = INGESTION_SOURCE_NUM_PROCESSOR_THREADS_SETTING.get(settings);
             final Map<String, Object> ingestionSourceParams = INGESTION_SOURCE_PARAMS_SETTING.getAsMap(settings);
-            return new IngestionSource(ingestionSourceType, pointerInitReset, errorStrategy, ingestionSourceParams);
+            return new IngestionSource(ingestionSourceType, pointerInitReset, errorStrategy, numProcessorThreads, ingestionSourceParams);
         }
         return null;
     }

server/src/main/java/org/opensearch/cluster/metadata/IngestionSource.java

@@ -23,18 +23,21 @@ public class IngestionSource {
     private String type;
     private PointerInitReset pointerInitReset;
     private IngestionErrorStrategy.ErrorStrategy errorStrategy;
+    private int numMessageProcessorThreads;
     private Map<String, Object> params;
 
     public IngestionSource(
         String type,
         PointerInitReset pointerInitReset,
         IngestionErrorStrategy.ErrorStrategy errorStrategy,
+        int numMessageProcessorThreads,
         Map<String, Object> params
     ) {
         this.type = type;
         this.pointerInitReset = pointerInitReset;
         this.params = params;
         this.errorStrategy = errorStrategy;
+        this.numMessageProcessorThreads = numMessageProcessorThreads;
     }
 
     public String getType() {
@@ -53,6 +56,10 @@ public Map<String, Object> params() {
         return params;
     }
 
+    public int getNumMessageProcessorThreads() {
+        return numMessageProcessorThreads;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -61,12 +68,13 @@ public boolean equals(Object o) {
         return Objects.equals(type, ingestionSource.type)
             && Objects.equals(pointerInitReset, ingestionSource.pointerInitReset)
             && Objects.equals(errorStrategy, ingestionSource.errorStrategy)
-            && Objects.equals(params, ingestionSource.params);
+            && Objects.equals(params, ingestionSource.params)
+            && Objects.equals(numMessageProcessorThreads, ingestionSource.numMessageProcessorThreads);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(type, pointerInitReset, params, errorStrategy);
+        return Objects.hash(type, pointerInitReset, params, errorStrategy, numMessageProcessorThreads);
     }
 
     @Override
@@ -81,6 +89,9 @@ public String toString() {
             + ",error_strategy='"
             + errorStrategy
             + '\''
+            + ",numMessageProcessorThreads='"
+            + numMessageProcessorThreads
+            + '\''
             + ", params="
             + params
             + '}';

server/src/main/java/org/opensearch/common/settings/IndexScopedSettings.java

@@ -267,6 +267,7 @@ 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_NUM_PROCESSOR_THREADS_SETTING,
 
                 // validate that built-in similarities don't get redefined
                 Setting.groupSetting("index.similarity.", (s) -> {

server/src/main/java/org/opensearch/index/engine/IngestionEngine.java

@@ -117,7 +117,8 @@ public void start() {
             resetState,
             resetValue,
             ingestionErrorStrategy,
-            initialPollerState
+            initialPollerState,
+            ingestionSource.getNumMessageProcessorThreads()
         );
         streamPoller.start();
     }

server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java

@@ -21,8 +21,6 @@
 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;
 
@@ -48,8 +46,6 @@ public class DefaultStreamPoller implements StreamPoller {
 
     private ExecutorService consumerThread;
 
-    private ExecutorService processorThread;
-
     // start of the batch, inclusive
     private IngestionShardPointer batchStartPointer;
     private boolean includeBatchStartPointer = false;
@@ -59,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,
@@ -77,13 +71,14 @@ public DefaultStreamPoller(
         ResetState resetState,
         String resetValue,
         IngestionErrorStrategy errorStrategy,
-        State initialState
+        State initialState,
+        int numProcessorThreads
     ) {
         this(
             startPointer,
             persistedPointers,
             consumer,
-            new MessageProcessorRunnable(new ArrayBlockingQueue<>(100), ingestionEngine, errorStrategy),
+            new PartitionedBlockingQueueContainer(numProcessorThreads, consumer.getShardId(), ingestionEngine, errorStrategy),
             resetState,
             resetValue,
             errorStrategy,
@@ -95,7 +90,7 @@ public DefaultStreamPoller(
         IngestionShardPointer startPointer,
         Set<IngestionShardPointer> persistedPointers,
         IngestionShardConsumer consumer,
-        MessageProcessorRunnable processorRunnable,
+        PartitionedBlockingQueueContainer blockingQueueContainer,
         ResetState resetState,
         String resetValue,
         IngestionErrorStrategy errorStrategy,
@@ -110,22 +105,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;
     }
 
@@ -143,7 +129,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();
     }
 
     /**
@@ -234,7 +220,7 @@ protected void startPoll() {
                         continue;
                     }
                     totalPolledCount.inc();
-                    blockingQueue.put(result);
+                    blockingQueueContainer.add(result);
 
                     logger.debug(
                         "Put message {} with pointer {} to the blocking queue",
@@ -311,10 +297,10 @@ 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());
     }
 
@@ -337,7 +323,7 @@ public IngestionShardPointer getBatchStartPointer() {
     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();
     }
 
@@ -353,6 +339,6 @@ public IngestionErrorStrategy getErrorStrategy() {
     @Override
     public void updateErrorStrategy(IngestionErrorStrategy errorStrategy) {
         this.errorStrategy = errorStrategy;
-        processorRunnable.setErrorStrategy(errorStrategy);
+        blockingQueueContainer.updateErrorStrategy(errorStrategy);
     }
 }

server/src/main/java/org/opensearch/indices/pollingingest/IngestionUtils.java

@@ -0,0 +1,33 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.indices.pollingingest;
+
+import org.opensearch.common.UUIDs;
+import org.opensearch.common.xcontent.XContentHelper;
+import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.xcontent.MediaTypeRegistry;
+
+import java.util.Map;
+
+/**
+ * Holds common utilities for streaming ingestion.
+ */
+public class IngestionUtils {
+
+    public static Map<String, Object> getParsedPayloadMap(byte[] payload) {
+        BytesReference payloadBR = new BytesArray(payload);
+        Map<String, Object> payloadMap = XContentHelper.convertToMap(payloadBR, false, MediaTypeRegistry.xContentType(payloadBR)).v2();
+        return payloadMap;
+    }
+
+    public static String generateID() {
+        return UUIDs.base64UUID();
+    }
+}