Wire IngestService from Node to IngestionEngine for pull-based ingestion (opensearch-project#20757)

* Wire IngestService from Node to IngestionEngine for pull-based ingestion

Signed-off-by: Rishab Nahata <rishabnahata07@gmail.com>

Author: imRishN

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

@@ -40,6 +40,7 @@
 import org.opensearch.indices.pollingingest.IngestionSettings;
 import org.opensearch.indices.pollingingest.PollingIngestStats;
 import org.opensearch.indices.pollingingest.StreamPoller;
+import org.opensearch.ingest.IngestService;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -60,11 +61,17 @@ public class IngestionEngine extends InternalEngine {
     private StreamPoller streamPoller;
     private final IngestionConsumerFactory ingestionConsumerFactory;
     private final DocumentMapperForType documentMapperForType;
+    private final IngestService ingestService;
     private volatile IngestionShardPointer lastCommittedBatchStartPointer;
 
     public IngestionEngine(EngineConfig engineConfig, IngestionConsumerFactory ingestionConsumerFactory) {
+        this(engineConfig, ingestionConsumerFactory, null);
+    }
+
+    public IngestionEngine(EngineConfig engineConfig, IngestionConsumerFactory ingestionConsumerFactory, IngestService ingestService) {
         super(engineConfig);
         this.ingestionConsumerFactory = Objects.requireNonNull(ingestionConsumerFactory);
+        this.ingestService = ingestService;
         this.documentMapperForType = engineConfig.getDocumentMapperForTypeSupplier().get();
         registerDynamicIndexSettingsHandlers();
     }

server/src/main/java/org/opensearch/indices/IndicesService.java

@@ -164,6 +164,7 @@
 import org.opensearch.indices.replication.checkpoint.ReferencedSegmentsPublisher;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.indices.replication.common.ReplicationType;
+import org.opensearch.ingest.IngestService;
 import org.opensearch.node.Node;
 import org.opensearch.node.remotestore.RemoteStoreNodeAttribute;
 import org.opensearch.plugins.IndexStorePlugin;
@@ -395,6 +396,7 @@ public class IndicesService extends AbstractLifecycleComponent
     private final Map<String, IndexStorePlugin.DirectoryFactory> directoryFactories;
     private final Map<String, IndexStorePlugin.CompositeDirectoryFactory> compositeDirectoryFactories;
     private final Map<String, IngestionConsumerFactory> ingestionConsumerFactories;
+    private final Supplier<IngestService> ingestServiceSupplier;
     private final Map<String, IndexStorePlugin.RecoveryStateFactory> recoveryStateFactories;
     private final Map<String, IndexStorePlugin.StoreFactory> storeFactories;
     final AbstractRefCounted indicesRefCount; // pkg-private for testing
@@ -458,6 +460,7 @@ public IndicesService(
         SearchRequestStats searchRequestStats,
         @Nullable RemoteStoreStatsTrackerFactory remoteStoreStatsTrackerFactory,
         Map<String, IngestionConsumerFactory> ingestionConsumerFactories,
+        Supplier<IngestService> ingestServiceSupplier,
         RecoverySettings recoverySettings,
         CacheService cacheService,
         RemoteStoreSettings remoteStoreSettings,
@@ -519,6 +522,7 @@ public void onRemoval(ShardId shardId, String fieldName, boolean wasEvicted, lon
         this.recoveryStateFactories = recoveryStateFactories;
         this.storeFactories = storeFactories;
         this.ingestionConsumerFactories = ingestionConsumerFactories;
+        this.ingestServiceSupplier = ingestServiceSupplier;
         // doClose() is called when shutting down a node, yet there might still be ongoing requests
         // that we need to wait for before closing some resources such as the caches. In order to
         // avoid closing these resources while ongoing requests are still being processed, we use a
@@ -666,6 +670,7 @@ public IndicesService(
             searchRequestStats,
             remoteStoreStatsTrackerFactory,
             ingestionConsumerFactories,
+            () -> null,
             recoverySettings,
             cacheService,
             remoteStoreSettings,
@@ -1164,7 +1169,7 @@ private EngineFactory getEngineFactory(final IndexSettings idxSettings) {
         // streaming ingestion
         if (indexMetadata != null && indexMetadata.useIngestionSource()) {
             IngestionConsumerFactory ingestionConsumerFactory = getIngestionConsumerFactory(idxSettings);
-            return new IngestionEngineFactory(ingestionConsumerFactory);
+            return new IngestionEngineFactory(ingestionConsumerFactory, ingestServiceSupplier);
         }
 
         final List<Optional<EngineFactory>> engineFactories = engineFactoryProviders.stream()

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

@@ -9,24 +9,30 @@
 package org.opensearch.indices.pollingingest;
 
 import org.opensearch.cluster.metadata.IngestionSource;
+import org.opensearch.common.Nullable;
 import org.opensearch.index.IngestionConsumerFactory;
 import org.opensearch.index.engine.Engine;
 import org.opensearch.index.engine.EngineConfig;
 import org.opensearch.index.engine.EngineFactory;
 import org.opensearch.index.engine.IngestionEngine;
 import org.opensearch.index.engine.NRTReplicationEngine;
+import org.opensearch.ingest.IngestService;
 
 import java.util.Objects;
+import java.util.function.Supplier;
 
 /**
  * Engine Factory implementation used with streaming ingestion.
  */
 public class IngestionEngineFactory implements EngineFactory {
 
     private final IngestionConsumerFactory ingestionConsumerFactory;
+    @Nullable
+    private final Supplier<IngestService> ingestServiceSupplier;
 
-    public IngestionEngineFactory(IngestionConsumerFactory ingestionConsumerFactory) {
+    public IngestionEngineFactory(IngestionConsumerFactory ingestionConsumerFactory, Supplier<IngestService> ingestServiceSupplier) {
         this.ingestionConsumerFactory = Objects.requireNonNull(ingestionConsumerFactory);
+        this.ingestServiceSupplier = ingestServiceSupplier;
     }
 
     /**
@@ -39,9 +45,13 @@ public Engine newReadWriteEngine(EngineConfig config) {
         IngestionSource ingestionSource = config.getIndexSettings().getIndexMetadata().getIngestionSource();
         boolean isAllActiveIngestion = ingestionSource != null && ingestionSource.isAllActiveIngestionEnabled();
 
+        IngestService ingestService = ingestServiceSupplier != null ? ingestServiceSupplier.get() : null;
+        assert ingestService != null || ingestServiceSupplier == null
+            : "IngestService supplier returned null. This indicates a initialization ordering issue.";
+
         if (isAllActiveIngestion) {
             // use ingestion engine on both primary and replica in all-active mode
-            IngestionEngine ingestionEngine = new IngestionEngine(config, ingestionConsumerFactory);
+            IngestionEngine ingestionEngine = new IngestionEngine(config, ingestionConsumerFactory, ingestService);
             ingestionEngine.start();
             return ingestionEngine;
         }
@@ -52,7 +62,7 @@ public Engine newReadWriteEngine(EngineConfig config) {
             return new NRTReplicationEngine(config);
         }
 
-        IngestionEngine ingestionEngine = new IngestionEngine(config, ingestionConsumerFactory);
+        IngestionEngine ingestionEngine = new IngestionEngine(config, ingestionConsumerFactory, ingestService);
         ingestionEngine.start();
         return ingestionEngine;
     }

server/src/main/java/org/opensearch/node/Node.java

@@ -626,6 +626,7 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
             }
 
             final SetOnce<RepositoriesService> repositoriesServiceReference = new SetOnce<>();
+            final SetOnce<IngestService> ingestServiceReference = new SetOnce<>();
             final RemoteStoreNodeService remoteStoreNodeService = new RemoteStoreNodeService(repositoriesServiceReference::get, threadPool);
             localNodeFactory = new RemoteStoreVerifyingLocalNodeFactory(settings, nodeEnvironment.nodeId(), remoteStoreNodeService);
             resourcesToClose.add(() -> ThreadPool.terminate(threadPool, 10, TimeUnit.SECONDS));
@@ -1001,6 +1002,7 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
                 searchRequestStats,
                 remoteStoreStatsTrackerFactory,
                 ingestionConsumerFactories,
+                ingestServiceReference::get,
                 recoverySettings,
                 cacheService,
                 remoteStoreSettings,
@@ -1022,6 +1024,7 @@ protected Node(final Environment initialEnvironment, Collection<PluginInfo> clas
                 xContentRegistry,
                 new SystemIngestPipelineCache()
             );
+            ingestServiceReference.set(ingestService);
 
             final FsServiceProvider fsServiceProvider = new FsServiceProvider(
                 settings,

server/src/test/java/org/opensearch/index/engine/IngestionEngineTests.java

@@ -25,6 +25,7 @@
 import org.opensearch.indices.pollingingest.PollingIngestStats;
 import org.opensearch.indices.pollingingest.StreamPoller;
 import org.opensearch.indices.replication.common.ReplicationType;
+import org.opensearch.ingest.IngestService;
 import org.opensearch.test.IndexSettingsModule;
 import org.junit.After;
 import org.junit.Assert;
@@ -265,4 +266,59 @@ private boolean resultsFound(Engine engine, int numDocs) {
             return searcher.getIndexReader().numDocs() == numDocs;
         }
     }
+
+    public void testConstructorWithNullIngestService() throws IOException {
+        final AtomicLong globalCheckpoint = new AtomicLong(SequenceNumbers.NO_OPS_PERFORMED);
+        Store testStore = createStore(indexSettings, newDirectory());
+        FakeIngestionSource.FakeIngestionConsumerFactory consumerFactory = new FakeIngestionSource.FakeIngestionConsumerFactory(messages);
+
+        EngineConfig config = config(indexSettings, testStore, createTempDir(), NoMergePolicy.INSTANCE, null, null, globalCheckpoint::get);
+        String mapping = "{\"properties\":{\"name\":{\"type\": \"text\"},\"age\":{\"type\": \"integer\"}}}}";
+        MapperService mapperService = createMapperService(mapping);
+        config = config(config, () -> new DocumentMapperForType(mapperService.documentMapper(), null), clusterApplierService);
+
+        testStore.createEmpty(config.getIndexSettings().getIndexVersionCreated().luceneVersion);
+        final String translogUuid = Translog.createEmptyTranslog(
+            config.getTranslogConfig().getTranslogPath(),
+            SequenceNumbers.NO_OPS_PERFORMED,
+            shardId,
+            primaryTerm.get()
+        );
+        testStore.associateIndexWithNewTranslog(translogUuid);
+
+        // null IngestService — engine should start without pipeline support
+        IngestionEngine engine = new IngestionEngine(config, consumerFactory, null);
+        engine.start();
+        waitForResults(engine, 2);
+        engine.close();
+        testStore.close();
+    }
+
+    public void testConstructorWithNonNullIngestService() throws IOException {
+        final AtomicLong globalCheckpoint = new AtomicLong(SequenceNumbers.NO_OPS_PERFORMED);
+        Store testStore = createStore(indexSettings, newDirectory());
+        FakeIngestionSource.FakeIngestionConsumerFactory consumerFactory = new FakeIngestionSource.FakeIngestionConsumerFactory(messages);
+
+        EngineConfig config = config(indexSettings, testStore, createTempDir(), NoMergePolicy.INSTANCE, null, null, globalCheckpoint::get);
+        String mapping = "{\"properties\":{\"name\":{\"type\": \"text\"},\"age\":{\"type\": \"integer\"}}}}";
+        MapperService mapperService = createMapperService(mapping);
+        config = config(config, () -> new DocumentMapperForType(mapperService.documentMapper(), null), clusterApplierService);
+
+        testStore.createEmpty(config.getIndexSettings().getIndexVersionCreated().luceneVersion);
+        final String translogUuid = Translog.createEmptyTranslog(
+            config.getTranslogConfig().getTranslogPath(),
+            SequenceNumbers.NO_OPS_PERFORMED,
+            shardId,
+            primaryTerm.get()
+        );
+        testStore.associateIndexWithNewTranslog(translogUuid);
+
+        // non-null IngestService — engine should start with pipeline support available
+        IngestService ingestService = mock(IngestService.class);
+        IngestionEngine engine = new IngestionEngine(config, consumerFactory, ingestService);
+        engine.start();
+        waitForResults(engine, 2);
+        engine.close();
+        testStore.close();
+    }
 }

server/src/test/java/org/opensearch/indices/pollingingest/IngestionEngineFactoryTests.java

@@ -0,0 +1,34 @@
+/*
+ * 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.index.IngestionConsumerFactory;
+import org.opensearch.ingest.IngestService;
+import org.opensearch.test.OpenSearchTestCase;
+
+import static org.mockito.Mockito.mock;
+
+public class IngestionEngineFactoryTests extends OpenSearchTestCase {
+
+    public void testConstructorWithIngestServiceSupplier() {
+        IngestionConsumerFactory consumerFactory = mock(IngestionConsumerFactory.class);
+        IngestService ingestService = mock(IngestService.class);
+
+        IngestionEngineFactory factory = new IngestionEngineFactory(consumerFactory, () -> ingestService);
+        assertNotNull(factory);
+    }
+
+    public void testConstructorWithNullIngestServiceSupplier() {
+        IngestionConsumerFactory consumerFactory = mock(IngestionConsumerFactory.class);
+
+        // Null supplier should be accepted
+        IngestionEngineFactory factory = new IngestionEngineFactory(consumerFactory, null);
+        assertNotNull(factory);
+    }
+}