merging masseyke:sampling-service-shell

Author: masseyke

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

@@ -90,7 +90,7 @@ public abstract class TransportAbstractBulkAction extends HandledTransportAction
  protected final Executor systemCoordinationExecutor;
  private final ActionType<BulkResponse> bulkAction;
  protected final FeatureService featureService;
- private final SamplingService samplingService;
+ protected final SamplingService samplingService;
 
  public TransportAbstractBulkAction(
  ActionType<BulkResponse> action,
@@ -207,7 +207,7 @@ 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, true);
+ applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, releasingListener, false);
  }
  });
  }
@@ -217,7 +217,7 @@ private boolean applyPipelines(
  BulkRequest bulkRequest,
  Executor executor,
  ActionListener<BulkResponse> listener,
- boolean firstTime
+ boolean haveRunIngestService
  ) throws IOException {
  boolean hasIndexRequestsWithPipelines = false;
  ClusterState state = clusterService.state();
@@ -311,8 +311,11 @@ private boolean applyPipelines(
  }
  });
  return true;
- } else if (samplingService != null && firstTime) {
- // else sample, but only if this is the first time through. Otherwise we had pipelines and sampled in IngestService
+ } 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);
@@ -353,7 +356,7 @@ private void processBulkIndexIngestRequest(
  ActionRunnable<BulkResponse> runnable = new ActionRunnable<>(actionListener) {
  @Override
  protected void doRun() throws IOException {
- applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, false);
+ applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, true);
  }
 
  @Override
@@ -432,7 +435,7 @@ private void applyPipelinesAndDoInternalExecute(
  BulkRequest bulkRequest,
  Executor executor,
  ActionListener<BulkResponse> listener,
- boolean firstTime
+ boolean haveRunIngestService
  ) throws IOException {
  final long relativeStartTimeNanos = relativeTimeNanos();
 
@@ -450,7 +453,7 @@ private void applyPipelinesAndDoInternalExecute(
 
  var wrappedListener = bulkRequestModifier.wrapActionListenerIfNeeded(listener);
 
- if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, firstTime) == false) {
+ if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, haveRunIngestService) == false) {
  doInternalExecute(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, relativeStartTimeNanos);
  }
  }

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

@@ -83,7 +83,6 @@
 import org.elasticsearch.threadpool.Scheduler;
 import org.elasticsearch.threadpool.ThreadPool;
 
-import java.io.IOException;
 import java.time.Instant;
 import java.time.InstantSource;
 import java.util.ArrayList;
@@ -981,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;
  }
 
@@ -1199,11 +1199,11 @@ private void executePipelines(
  }
  };
  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");
  }
@@ -1355,37 +1355,22 @@ private void executePipelines(
  if (newPipelines.hasNext()) {
  executePipelines(newPipelines, indexRequest, ingestDocument, resolveFailureStore, listener, originalDocumentMetadata);
  } else {
- try {
- /*
- * At this point, all pipelines have been executed, and we are about to overwrite ingestDocument with the results.
- * 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.
- */
- haveAttemptedSampling.set(true);
- IndexRequest original = copyIndexRequest(indexRequest);
- updateIndexRequestMetadata(original, originalDocumentMetadata);
- samplingService.maybeSample(project, original, ingestDocument);
- } catch (IOException ex) {
- logger.warn("unable to sample data");
- }
-
+ /*
+ * 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) {
- try {
- if (haveAttemptedSampling.get() == false) {
- IndexRequest original = copyIndexRequest(indexRequest);
- updateIndexRequestMetadata(original, originalDocumentMetadata);
- samplingService.maybeSample(state.projectState(projectResolver.getProjectId()).metadata(), original, ingestDocument);
- }
- } catch (IOException ex) {
- logger.warn("unable to sample data");
+ 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);
  }
- // Maybe also sample here? Or put it in the exceptionHandler? We want to make sure the exception didn't come of out the
- // listener though.
  logger.debug(
  () -> format("failed to execute pipeline [%s] for document [%s/%s]", pipelineId, indexRequest.index(), indexRequest.id()),
  e
@@ -1394,20 +1379,49 @@ private void executePipelines(
  }
  }
 
- private IndexRequest copyIndexRequest(IndexRequest original) throws IOException {
+ 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());

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

@@ -184,8 +184,13 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) {
  );
  }
 
+ /**
+ * Potentially samples the given indexRequest, depending on the existing sampling configuration.
+ * @param projectMetadata Used to get the sampling configuration
+ * @param indexRequest The raw request to potentially sample
+ */
  public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) {
- maybeSample(projectMetadata, indexRequest, () -> {
+ maybeSample(projectMetadata, indexRequest.index(), () -> indexRequest, () -> {
  Map<String, Object> sourceAsMap;
  try {
  sourceAsMap = indexRequest.sourceAsMap();
@@ -204,19 +209,20 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque
  });
  }
 
- public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) {
- maybeSample(projectMetadata, indexRequest, () -> ingestDocument);
- }
-
- private void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, Supplier<IngestDocument> ingestDocumentSupplier) {
+ private void maybeSample(
+ ProjectMetadata projectMetadata,
+ String indexName,
+ Supplier<IndexRequest> indexRequestSupplier,
+ Supplier<IngestDocument> ingestDocumentSupplier
+ ) {
  long startTime = relativeNanoTimeSupplier.getAsLong();
  TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom(
  TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME
  );
  ProjectId projectId = projectMetadata.id();
  if (samplingConfig != null) {
  String samplingIndex = samplingConfig.indexName;
- if (samplingIndex.equals(indexRequest.index())) {
+ if (samplingIndex.equals(indexName)) {
  SoftReference<SampleInfo> sampleInfoReference = samples.compute(
  new ProjectIndex(projectId, samplingIndex),
  (k, v) -> v == null || v.get() == null
@@ -266,6 +272,7 @@ private void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequ
  sampleInfo.stats
  )) {
  stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime));
+ IndexRequest indexRequest = indexRequestSupplier.get();
  indexRequest.incRef();
  if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) {
  releaseableSource.incRef();
@@ -371,6 +378,25 @@ private boolean isClusterServiceStoppedOrClosed() {
  return state == Lifecycle.State.STOPPED || state == Lifecycle.State.CLOSED;
  }
 
+ /**
+ *
+ * @param projectMetadata Used to get the sampling configuration
+ * @param indexRequestSupplier A supplier for the raw request to potentially sample
+ * @param ingestDocument The IngestDocument used for evaluating any conditionals that are part of the sample configuration
+ */
+ public void maybeSample(
+ ProjectMetadata projectMetadata,
+ String indexName,
+ Supplier<IndexRequest> indexRequestSupplier,
+ IngestDocument ingestDocument
+ ) {
+ maybeSample(projectMetadata, indexName, indexRequestSupplier, () -> ingestDocument);
+ }
+
+ public boolean atLeastOneSampleConfigured() {
+ return false; // TODO Return true if there is at least one sample in the cluster state
+ }
+
  @Override
  public void clusterChanged(ClusterChangedEvent event) {
  final boolean prevIsMaster = this.isMaster;
@@ -609,4 +635,5 @@ static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask {
  }
 
  record ProjectIndex(ProjectId projectId, String indexName) {};
+
 }

server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java

@@ -174,10 +174,16 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) {
  return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature);
  }
  },
- mock(SamplingService.class)
+ initializeSamplingService()
  );
  }
 
+ private static SamplingService initializeSamplingService() {
+ SamplingService samplingService = mock(SamplingService.class);
+ when(samplingService.atLeastOneSampleConfigured()).thenReturn(true);
+ return samplingService;
+ }
+
  @Override
  void executeBulk(
  Task task,

server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java

@@ -82,6 +82,7 @@
 import java.util.SortedMap;
 import java.util.TreeMap;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -93,6 +94,8 @@
 import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class TransportBulkActionTests extends ESTestCase {
@@ -154,10 +157,16 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) {
  return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature);
  }
  },
- mock(SamplingService.class)
+ initializeSamplingService()
  );
  }
 
+ private static SamplingService initializeSamplingService() {
+ SamplingService samplingService = mock(SamplingService.class);
+ when(samplingService.atLeastOneSampleConfigured()).thenReturn(true);
+ return samplingService;
+ }
+
  @Override
  void createIndex(CreateIndexRequest createIndexRequest, ActionListener<CreateIndexResponse> listener) {
  indexCreated = true;
@@ -735,6 +744,19 @@ public void testFailuresDuringPrerequisiteActions() throws InterruptedException
  assertNull(bulkRequest.requests.get(2));
  }
 
+ public void testSampling() throws ExecutionException, InterruptedException {
+ // This test makes sure that the sampling service is called once per IndexRequest
+ BulkRequest bulkRequest = new BulkRequest().add(new IndexRequest("index").id("id1").source(Collections.emptyMap()))
+ .add(new IndexRequest("index").id("id2").source(Collections.emptyMap()))
+ .add(new DeleteRequest("index2").id("id3"));
+ PlainActionFuture<BulkResponse> future = new PlainActionFuture<>();
+ ActionTestUtils.execute(bulkAction, null, bulkRequest, future);
+ future.get();
+ assertTrue(bulkAction.indexCreated);
+ // We expect 2 sampling calls since there are 2 index requests:
+ verify(bulkAction.samplingService, times(2)).maybeSample(any(), any());
+ }
+
  private BulkRequest buildBulkRequest(List<String> indices) {
  BulkRequest request = new BulkRequest();
  for (String index : indices) {

server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java

@@ -39,6 +39,7 @@
 import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.index.IndexingPressure;
 import org.elasticsearch.indices.EmptySystemIndices;
+import org.elasticsearch.ingest.SamplingService;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.VersionUtils;
@@ -66,6 +67,7 @@
 import static org.elasticsearch.test.StreamsUtils.copyToStringFromClasspath;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.mockito.Mockito.mock;
 
 public class TransportBulkActionTookTests extends ESTestCase {
 
@@ -268,7 +270,7 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) {
  return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature);
  }
  },
- null
+ mock(SamplingService.class)
  );
  }
  }