Adding the shell of the SamplingService (#134766)

Author: masseyke

server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java

@@ -42,6 +42,7 @@
 import org.elasticsearch.index.IndexingPressure;
 import org.elasticsearch.indices.SystemIndices;
 import org.elasticsearch.ingest.IngestService;
+import org.elasticsearch.ingest.SamplingService;
 import org.elasticsearch.node.NodeClosedException;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -90,6 +91,7 @@ public abstract class TransportAbstractBulkAction extends HandledTransportAction
     protected final Executor systemCoordinationExecutor;
     private final ActionType<BulkResponse> bulkAction;
     protected final FeatureService featureService;
+    protected final SamplingService samplingService;
 
     public TransportAbstractBulkAction(
         ActionType<BulkResponse> action,
@@ -103,7 +105,8 @@ public TransportAbstractBulkAction(
         SystemIndices systemIndices,
         ProjectResolver projectResolver,
         LongSupplier relativeTimeNanosProvider,
-        FeatureService featureService
+        FeatureService featureService,
+        SamplingService samplingService
     ) {
         super(action.name(), transportService, actionFilters, requestReader, EsExecutors.DIRECT_EXECUTOR_SERVICE);
         this.threadPool = threadPool;
@@ -119,6 +122,7 @@ public TransportAbstractBulkAction(
         clusterService.addStateApplier(this.ingestForwarder);
         this.relativeTimeNanosProvider = relativeTimeNanosProvider;
         this.bulkAction = action;
+        this.samplingService = samplingService;
     }
 
     @Override
@@ -204,13 +208,18 @@ private void forkAndExecute(Task task, BulkRequest bulkRequest, Executor executo
         executor.execute(new ActionRunnable<>(releasingListener) {
             @Override
             protected void doRun() throws IOException {
-                applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, releasingListener);
+                applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, releasingListener, false);
             }
         });
     }
 
-    private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor executor, ActionListener<BulkResponse> listener)
-        throws IOException {
+    private boolean applyPipelines(
+        Task task,
+        BulkRequest bulkRequest,
+        Executor executor,
+        ActionListener<BulkResponse> listener,
+        boolean haveRunIngestService
+    ) throws IOException {
         boolean hasIndexRequestsWithPipelines = false;
         ClusterState state = clusterService.state();
         ProjectId projectId = projectResolver.getProjectId();
@@ -303,6 +312,16 @@ private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor exec
                 }
             });
             return true;
+        } else if (haveRunIngestService == false && samplingService != null && samplingService.atLeastOneSampleConfigured()) {
+            /*
+             * Else ample only if this request has not passed through IngestService::executeBulkRequest. Otherwise, some request within the
+             * bulk had pipelines and we sampled in IngestService already.
+             */
+            for (DocWriteRequest<?> actionRequest : bulkRequest.requests) {
+                if (actionRequest instanceof IndexRequest ir) {
+                    samplingService.maybeSample(project, ir);
+                }
+            }
         }
         return false;
     }
@@ -338,7 +357,7 @@ private void processBulkIndexIngestRequest(
                     ActionRunnable<BulkResponse> runnable = new ActionRunnable<>(actionListener) {
                         @Override
                         protected void doRun() throws IOException {
-                            applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener);
+                            applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, true);
                         }
 
                         @Override
@@ -416,7 +435,8 @@ private void applyPipelinesAndDoInternalExecute(
         Task task,
         BulkRequest bulkRequest,
         Executor executor,
-        ActionListener<BulkResponse> listener
+        ActionListener<BulkResponse> listener,
+        boolean haveRunIngestService
     ) throws IOException {
         final long relativeStartTimeNanos = relativeTimeNanos();
 
@@ -434,7 +454,7 @@ private void applyPipelinesAndDoInternalExecute(
 
         var wrappedListener = bulkRequestModifier.wrapActionListenerIfNeeded(listener);
 
-        if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener) == false) {
+        if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, haveRunIngestService) == false) {
             doInternalExecute(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, relativeStartTimeNanos);
         }
     }

server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java

@@ -50,6 +50,7 @@
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.indices.SystemIndices;
 import org.elasticsearch.ingest.IngestService;
+import org.elasticsearch.ingest.SamplingService;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -101,7 +102,8 @@ public TransportBulkAction(
         ProjectResolver projectResolver,
         FailureStoreMetrics failureStoreMetrics,
         DataStreamFailureStoreSettings dataStreamFailureStoreSettings,
-        FeatureService featureService
+        FeatureService featureService,
+        SamplingService samplingService
     ) {
         this(
             threadPool,
@@ -117,7 +119,8 @@ public TransportBulkAction(
             threadPool::relativeTimeInNanos,
             failureStoreMetrics,
             dataStreamFailureStoreSettings,
-            featureService
+            featureService,
+            samplingService
         );
     }
 
@@ -135,7 +138,8 @@ public TransportBulkAction(
         LongSupplier relativeTimeProvider,
         FailureStoreMetrics failureStoreMetrics,
         DataStreamFailureStoreSettings dataStreamFailureStoreSettings,
-        FeatureService featureService
+        FeatureService featureService,
+        SamplingService samplingService
     ) {
         this(
             TYPE,
@@ -153,7 +157,8 @@ public TransportBulkAction(
             relativeTimeProvider,
             failureStoreMetrics,
             dataStreamFailureStoreSettings,
-            featureService
+            featureService,
+            samplingService
         );
     }
 
@@ -173,7 +178,8 @@ public TransportBulkAction(
         LongSupplier relativeTimeProvider,
         FailureStoreMetrics failureStoreMetrics,
         DataStreamFailureStoreSettings dataStreamFailureStoreSettings,
-        FeatureService featureService
+        FeatureService featureService,
+        SamplingService samplingService
     ) {
         super(
             bulkAction,
@@ -187,7 +193,8 @@ public TransportBulkAction(
             systemIndices,
             projectResolver,
             relativeTimeProvider,
-            featureService
+            featureService,
+            samplingService
         );
         this.dataStreamFailureStoreSettings = dataStreamFailureStoreSettings;
         Objects.requireNonNull(relativeTimeProvider);

server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java

@@ -117,7 +117,8 @@ public TransportSimulateBulkAction(
             systemIndices,
             projectResolver,
             threadPool::relativeTimeInNanos,
-            featureService
+            featureService,
+            null
         );
         this.indicesService = indicesService;
         this.xContentRegistry = xContentRegistry;

server/src/main/java/org/elasticsearch/ingest/IngestService.java

@@ -52,6 +52,7 @@
 import org.elasticsearch.cluster.service.MasterServiceTaskQueue;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.TriConsumer;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
@@ -77,6 +78,7 @@
 import org.elasticsearch.index.analysis.AnalysisRegistry;
 import org.elasticsearch.node.ReportingService;
 import org.elasticsearch.plugins.IngestPlugin;
+import org.elasticsearch.script.Metadata;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.threadpool.Scheduler;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -98,6 +100,7 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -151,6 +154,7 @@ public static boolean locallySupportedIngestFeature(NodeFeature nodeFeature) {
     private volatile ClusterState state;
     private final ProjectResolver projectResolver;
     private final FeatureService featureService;
+    private final SamplingService samplingService;
     private final Consumer<ActionListener<NodesInfoResponse>> nodeInfoListener;
 
     private static BiFunction<Long, Runnable, Scheduler.ScheduledCancellable> createScheduler(ThreadPool threadPool) {
@@ -252,6 +256,7 @@ public IngestService(
         FailureStoreMetrics failureStoreMetrics,
         ProjectResolver projectResolver,
         FeatureService featureService,
+        SamplingService samplingService,
         Consumer<ActionListener<NodesInfoResponse>> nodeInfoListener
     ) {
         this.clusterService = clusterService;
@@ -276,6 +281,7 @@ public IngestService(
         this.failureStoreMetrics = failureStoreMetrics;
         this.projectResolver = projectResolver;
         this.featureService = featureService;
+        this.samplingService = samplingService;
         this.nodeInfoListener = nodeInfoListener;
     }
 
@@ -290,7 +296,8 @@ public IngestService(
         MatcherWatchdog matcherWatchdog,
         FailureStoreMetrics failureStoreMetrics,
         ProjectResolver projectResolver,
-        FeatureService featureService
+        FeatureService featureService,
+        SamplingService samplingService
     ) {
         this(
             clusterService,
@@ -304,6 +311,7 @@ public IngestService(
             failureStoreMetrics,
             projectResolver,
             featureService,
+            samplingService,
             createNodeInfoListener(client)
         );
     }
@@ -324,6 +332,7 @@ public IngestService(
         this.failureStoreMetrics = ingestService.failureStoreMetrics;
         this.projectResolver = ingestService.projectResolver;
         this.featureService = ingestService.featureService;
+        this.samplingService = ingestService.samplingService;
         this.nodeInfoListener = ingestService.nodeInfoListener;
     }
 
@@ -971,6 +980,7 @@ protected void doRun() {
                         Pipeline firstPipeline = pipelines.peekFirst();
                         if (pipelines.hasNext() == false) {
                             i++;
+                            samplingService.maybeSample(state.metadata().projects().get(pipelines.projectId()), indexRequest);
                             continue;
                         }
 
@@ -983,7 +993,7 @@ protected void doRun() {
                         final int slot = i;
                         final Releasable ref = refs.acquire();
                         final IngestDocument ingestDocument = newIngestDocument(indexRequest);
-                        final org.elasticsearch.script.Metadata originalDocumentMetadata = ingestDocument.getMetadata().clone();
+                        final Metadata originalDocumentMetadata = ingestDocument.getMetadata().clone();
                         // the document listener gives us three-way logic: a document can fail processing (1), or it can
                         // be successfully processed. a successfully processed document can be kept (2) or dropped (3).
                         final ActionListener<IngestPipelinesExecutionResult> documentListener = ActionListener.runAfter(
@@ -1030,7 +1040,14 @@ public void onFailure(Exception e) {
                             }
                         );
 
-                        executePipelines(pipelines, indexRequest, ingestDocument, adaptedResolveFailureStore, documentListener);
+                        executePipelines(
+                            pipelines,
+                            indexRequest,
+                            ingestDocument,
+                            adaptedResolveFailureStore,
+                            documentListener,
+                            originalDocumentMetadata
+                        );
                         assert actionRequest.index() != null;
 
                         i++;
@@ -1149,7 +1166,8 @@ private void executePipelines(
         final IndexRequest indexRequest,
         final IngestDocument ingestDocument,
         final Function<String, Boolean> resolveFailureStore,
-        final ActionListener<IngestPipelinesExecutionResult> listener
+        final ActionListener<IngestPipelinesExecutionResult> listener,
+        final Metadata originalDocumentMetadata
     ) {
         assert pipelines.hasNext();
         PipelineSlot slot = pipelines.next();
@@ -1180,12 +1198,12 @@ private void executePipelines(
                 listener.onFailure(e);
             }
         };
-
+        AtomicBoolean haveAttemptedSampling = new AtomicBoolean(false);
+        final var project = state.metadata().projects().get(pipelines.projectId());
         try {
             if (pipeline == null) {
                 throw new IllegalArgumentException("pipeline with id [" + pipelineId + "] does not exist");
             }
-            final var project = state.metadata().projects().get(pipelines.projectId());
             if (project == null) {
                 throw new IllegalArgumentException("project with id [" + pipelines.projectId() + "] does not exist");
             }
@@ -1335,15 +1353,24 @@ private void executePipelines(
                 }
 
                 if (newPipelines.hasNext()) {
-                    executePipelines(newPipelines, indexRequest, ingestDocument, resolveFailureStore, listener);
+                    executePipelines(newPipelines, indexRequest, ingestDocument, resolveFailureStore, listener, originalDocumentMetadata);
                 } else {
-                    // update the index request's source and (potentially) cache the timestamp for TSDB
+                    /*
+                     * At this point, all pipelines have been executed, and we are about to overwrite ingestDocument with the results.
+                     * This is our chance to sample with both the original document and all changes.
+                     */
+                    haveAttemptedSampling.set(true);
+                    attemptToSampleData(project, indexRequest, ingestDocument, originalDocumentMetadata);
                     updateIndexRequestSource(indexRequest, ingestDocument);
                     cacheRawTimestamp(indexRequest, ingestDocument);
                     listener.onResponse(IngestPipelinesExecutionResult.SUCCESSFUL_RESULT); // document succeeded!
                 }
             });
         } catch (Exception e) {
+            if (haveAttemptedSampling.get() == false) {
+                // It is possible that an exception happened after we sampled. We do not want to sample the same document twice.
+                attemptToSampleData(project, indexRequest, ingestDocument, originalDocumentMetadata);
+            }
             logger.debug(
                 () -> format("failed to execute pipeline [%s] for document [%s/%s]", pipelineId, indexRequest.index(), indexRequest.id()),
                 e
@@ -1352,6 +1379,56 @@ private void executePipelines(
         }
     }
 
+    private void attemptToSampleData(
+        ProjectMetadata projectMetadata,
+        IndexRequest indexRequest,
+        IngestDocument ingestDocument,
+        Metadata originalDocumentMetadata
+    ) {
+        if (samplingService != null && samplingService.atLeastOneSampleConfigured()) {
+            /*
+             * We need both the original document and the fully updated document for sampling, so we make a copy of the original
+             * before overwriting it here. We can discard it after sampling.
+             */
+            samplingService.maybeSample(projectMetadata, indexRequest.index(), () -> {
+                IndexRequest original = copyIndexRequestForSampling(indexRequest);
+                updateIndexRequestMetadata(original, originalDocumentMetadata);
+                return original;
+            }, ingestDocument);
+
+        }
+    }
+
+    /**
+     * Creates a copy of an IndexRequest to be used by random sampling.
+     * @param original The IndexRequest to be copied
+     * @return A copy of the IndexRequest
+     */
+    private IndexRequest copyIndexRequestForSampling(IndexRequest original) {
+        IndexRequest clonedRequest = new IndexRequest(original.index());
+        clonedRequest.id(original.id());
+        clonedRequest.routing(original.routing());
+        clonedRequest.version(original.version());
+        clonedRequest.versionType(original.versionType());
+        clonedRequest.setPipeline(original.getPipeline());
+        clonedRequest.setFinalPipeline(original.getFinalPipeline());
+        clonedRequest.setIfSeqNo(original.ifSeqNo());
+        clonedRequest.setIfPrimaryTerm(original.ifPrimaryTerm());
+        clonedRequest.setRefreshPolicy(original.getRefreshPolicy());
+        clonedRequest.waitForActiveShards(original.waitForActiveShards());
+        clonedRequest.timeout(original.timeout());
+        clonedRequest.opType(original.opType());
+        clonedRequest.setParentTask(original.getParentTask());
+        clonedRequest.setRequireDataStream(original.isRequireDataStream());
+        clonedRequest.setRequireAlias(original.isRequireAlias());
+        clonedRequest.setIncludeSourceOnError(original.getIncludeSourceOnError());
+        BytesReference source = original.source();
+        if (source != null) {
+            clonedRequest.source(source, original.getContentType());
+        }
+        return clonedRequest;
+    }
+
     private static void executePipeline(
         final IngestDocument ingestDocument,
         final Pipeline pipeline,