From 04be99a72c77e7902a32a312c0762402298037ed Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 5 Sep 2025 10:51:24 -0500 Subject: [PATCH 01/37] Prototype of random sampling of original documents --- muted-tests.yml | 3 - .../org/elasticsearch/TransportVersions.java | 1 + .../elasticsearch/action/ActionModule.java | 12 + .../bulk/TransportAbstractBulkAction.java | 58 +++- .../action/bulk/TransportBulkAction.java | 19 +- .../bulk/TransportSimulateBulkAction.java | 3 +- .../elasticsearch/indices/IndicesModule.java | 9 +- .../ingest/ConditionalProcessor.java | 4 +- .../elasticsearch/ingest/IngestService.java | 65 +++- .../elasticsearch/node/NodeConstruction.java | 6 +- .../sample/PutSampleConfigAction.java | 160 +++++++++ .../sample/RestPutSampleConfigAction.java | 79 +++++ .../TransportPutSampleConfigAction.java | 313 ++++++++++++++++++ 13 files changed, 704 insertions(+), 28 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/sample/PutSampleConfigAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java diff --git a/muted-tests.yml b/muted-tests.yml index b5cba9a59e697..c57e584993e5e 100644 --- a/muted-tests.yml +++ b/muted-tests.yml @@ -414,9 +414,6 @@ tests: - class: org.elasticsearch.xpack.ml.integration.RevertModelSnapshotIT method: testRevertModelSnapshot_DeleteInterveningResults issue: https://github.com/elastic/elasticsearch/issues/132349 -#- class: org.elasticsearch.xpack.ml.integration.TextEmbeddingQueryIT -# method: testHybridSearch -# issue: https://github.com/elastic/elasticsearch/issues/132703 - class: org.elasticsearch.xpack.ml.integration.RevertModelSnapshotIT method: testRevertModelSnapshot issue: https://github.com/elastic/elasticsearch/issues/132733 diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index 85f7d75869085..0c8b4c9835d81 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -359,6 +359,7 @@ static TransportVersion def(int id) { public static final TransportVersion SEMANTIC_QUERY_MULTIPLE_INFERENCE_IDS = def(9_150_0_00); public static final TransportVersion ESQL_LOOKUP_JOIN_PRE_JOIN_FILTER = def(9_151_0_00); public static final TransportVersion INFERENCE_API_DISABLE_EIS_RATE_LIMITING = def(9_152_0_00); + public static final TransportVersion RANDOM_SAMPLING = def(9_153_0_00); /* * STOP! READ THIS FIRST! No, really, diff --git a/server/src/main/java/org/elasticsearch/action/ActionModule.java b/server/src/main/java/org/elasticsearch/action/ActionModule.java index 211adffba5ec8..a8bdd7bd4f3d2 100644 --- a/server/src/main/java/org/elasticsearch/action/ActionModule.java +++ b/server/src/main/java/org/elasticsearch/action/ActionModule.java @@ -405,6 +405,12 @@ import org.elasticsearch.rest.action.synonyms.RestGetSynonymsSetsAction; import org.elasticsearch.rest.action.synonyms.RestPutSynonymRuleAction; import org.elasticsearch.rest.action.synonyms.RestPutSynonymsAction; +import org.elasticsearch.sample.GetSampleAction; +import org.elasticsearch.sample.PutSampleConfigAction; +import org.elasticsearch.sample.RestGetSampleAction; +import org.elasticsearch.sample.RestPutSampleConfigAction; +import org.elasticsearch.sample.TransportGetSampleAction; +import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.snapshots.TransportUpdateSnapshotStatusAction; import org.elasticsearch.tasks.Task; import org.elasticsearch.telemetry.TelemetryProvider; @@ -813,6 +819,9 @@ public void reg actions.register(GetSynonymRuleAction.INSTANCE, TransportGetSynonymRuleAction.class); actions.register(DeleteSynonymRuleAction.INSTANCE, TransportDeleteSynonymRuleAction.class); + actions.register(PutSampleConfigAction.INSTANCE, TransportPutSampleConfigAction.class); + actions.register(GetSampleAction.INSTANCE, TransportGetSampleAction.class); + return unmodifiableMap(actions.getRegistry()); } @@ -1040,6 +1049,9 @@ public void initRestHandlers(Supplier nodesInCluster, Predicate< registerHandler.accept(new RestPutSynonymRuleAction()); registerHandler.accept(new RestGetSynonymRuleAction()); registerHandler.accept(new RestDeleteSynonymRuleAction()); + + registerHandler.accept(new RestPutSampleConfigAction()); + registerHandler.accept(new RestGetSampleAction()); } @Override diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java index 03029258cb6c4..73ebe3fca5234 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java +++ b/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; @@ -50,6 +51,7 @@ import java.io.IOException; import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; @@ -63,7 +65,7 @@ */ public abstract class TransportAbstractBulkAction extends HandledTransportAction { private static final Logger logger = LogManager.getLogger(TransportAbstractBulkAction.class); - + private final Map> samples = new HashMap<>(); public static final Set STREAMS_ALLOWED_PARAMS = new HashSet<>(8) { { add("error_trace"); @@ -89,6 +91,7 @@ public abstract class TransportAbstractBulkAction extends HandledTransportAction protected final Executor systemCoordinationExecutor; private final ActionType bulkAction; protected final FeatureService featureService; + private final SamplingService samplingService; public TransportAbstractBulkAction( ActionType action, @@ -102,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; @@ -118,6 +122,7 @@ public TransportAbstractBulkAction( clusterService.addStateApplier(this.ingestForwarder); this.relativeTimeNanosProvider = relativeTimeNanosProvider; this.bulkAction = action; + this.samplingService = samplingService; } @Override @@ -203,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, true); } }); } - private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor executor, ActionListener listener) - throws IOException { + private boolean applyPipelines( + Task task, + BulkRequest bulkRequest, + Executor executor, + ActionListener listener, + boolean firstTime + ) throws IOException { boolean hasIndexRequestsWithPipelines = false; ClusterState state = clusterService.state(); ProjectId projectId = projectResolver.getProjectId(); @@ -302,6 +312,13 @@ private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor exec } }); return true; + } else if (samplingService != null && firstTime) { + // else sample, but only if this is the first time through? + for (DocWriteRequest actionRequest : bulkRequest.requests) { + if (actionRequest instanceof IndexRequest ir) { + samplingService.maybeSample(project, ir); + } + } } return false; } @@ -316,6 +333,10 @@ private void processBulkIndexIngestRequest( final long ingestStartTimeInNanos = relativeTimeNanos(); final BulkRequestModifier bulkRequestModifier = new BulkRequestModifier(original); final Thread originalThread = Thread.currentThread(); + // Function rawDocSaver = indexRequest -> { + // indexRequest.source(); + // return null; + // }; getIngestService(original).executeBulkRequest( metadata.id(), original.numberOfActions(), @@ -337,7 +358,7 @@ private void processBulkIndexIngestRequest( ActionRunnable runnable = new ActionRunnable<>(actionListener) { @Override protected void doRun() throws IOException { - applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener); + applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, false); } @Override @@ -415,7 +436,8 @@ private void applyPipelinesAndDoInternalExecute( Task task, BulkRequest bulkRequest, Executor executor, - ActionListener listener + ActionListener listener, + boolean firstTime ) throws IOException { final long relativeStartTimeNanos = relativeTimeNanos(); @@ -432,8 +454,22 @@ private void applyPipelinesAndDoInternalExecute( } var wrappedListener = bulkRequestModifier.wrapActionListenerIfNeeded(listener); - - if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener) == false) { + boolean noPipelinesRemaining; + try { + noPipelinesRemaining = applyPipelines( + task, + bulkRequestModifier.getBulkRequest(), + executor, + wrappedListener, + firstTime + ) == false; + } catch (Exception e) { + maybeSample(); + throw e; + } + if (noPipelinesRemaining) { + // we have run all pipelines, so maybe sample + maybeSample(); doInternalExecute(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, relativeStartTimeNanos); } } @@ -490,6 +526,10 @@ private boolean streamsRestrictedParamsUsed(BulkRequest bulkRequest) { return Sets.difference(bulkRequest.requestParamsUsed(), STREAMS_ALLOWED_PARAMS).isEmpty() == false; } + private void maybeSample() { + // Is sampling enabled for this index? + } + /** * This method creates any missing resources and actually applies the BulkRequest to the relevant indices * @param task The task in which this work is being done diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java index 7e443e055cc90..81d60886b7bab 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java +++ b/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); diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java index b52f5447b9311..6712920b3bf85 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java +++ b/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; diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java index 09be98630d5c4..1790a2397d8e3 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java @@ -22,6 +22,7 @@ import org.elasticsearch.action.admin.indices.rollover.MinSizeCondition; import org.elasticsearch.action.admin.indices.rollover.OptimalShardCountCondition; import org.elasticsearch.action.resync.TransportResyncReplicationAction; +import org.elasticsearch.cluster.metadata.Metadata; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.index.mapper.BinaryFieldMapper; import org.elasticsearch.index.mapper.BooleanFieldMapper; @@ -77,6 +78,7 @@ import org.elasticsearch.injection.guice.AbstractModule; import org.elasticsearch.plugins.FieldPredicate; import org.elasticsearch.plugins.MapperPlugin; +import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.xcontent.NamedXContentRegistry; import org.elasticsearch.xcontent.ParseField; @@ -115,7 +117,12 @@ public static List getNamedWriteables() { new NamedWriteableRegistry.Entry(Condition.class, MaxSizeCondition.NAME, MaxSizeCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardSizeCondition.NAME, MaxPrimaryShardSizeCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardDocsCondition.NAME, MaxPrimaryShardDocsCondition::new), - new NamedWriteableRegistry.Entry(Condition.class, OptimalShardCountCondition.NAME, OptimalShardCountCondition::new) + new NamedWriteableRegistry.Entry(Condition.class, OptimalShardCountCondition.NAME, OptimalShardCountCondition::new), + new NamedWriteableRegistry.Entry( + Metadata.ProjectCustom.class, + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME, + TransportPutSampleConfigAction.SamplingConfigCustomMetadata::new + ) ); } diff --git a/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java b/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java index 39c3882c4ee49..20db72db39f91 100644 --- a/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java +++ b/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java @@ -40,7 +40,7 @@ public class ConditionalProcessor extends AbstractProcessor implements WrappingProcessor { private static final DeprecationLogger deprecationLogger = DeprecationLogger.getLogger(DynamicMap.class); - private static final Map> FUNCTIONS = Map.of("_type", value -> { + public static final Map> FUNCTIONS = Map.of("_type", value -> { deprecationLogger.warn( DeprecationCategory.INDICES, "conditional-processor__type", @@ -185,7 +185,7 @@ private static UnsupportedOperationException unmodifiableException() { return new UnsupportedOperationException("Mutating ingest documents in conditionals is not supported"); } - private static final class UnmodifiableIngestData implements Map { + public static final class UnmodifiableIngestData implements Map { private final Map data; diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 250ab60f7442c..07458a97d60bc 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -16,6 +16,7 @@ import org.elasticsearch.ElasticsearchParseException; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.ResourceNotFoundException; +import org.elasticsearch.TransportVersion; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.DocWriteRequest; import org.elasticsearch.action.admin.cluster.node.info.NodeInfo; @@ -53,6 +54,8 @@ import org.elasticsearch.common.Priority; import org.elasticsearch.common.TriConsumer; import org.elasticsearch.common.collect.ImmutableOpenMap; +import org.elasticsearch.common.io.stream.BytesStreamOutput; +import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.common.regex.Regex; @@ -78,10 +81,12 @@ import org.elasticsearch.node.ReportingService; import org.elasticsearch.plugins.IngestPlugin; import org.elasticsearch.plugins.internal.XContentParserDecorator; +import org.elasticsearch.script.Metadata; import org.elasticsearch.script.ScriptService; 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; @@ -157,6 +162,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> nodeInfoListener; private static BiFunction createScheduler(ThreadPool threadPool) { @@ -258,6 +264,7 @@ public IngestService( FailureStoreMetrics failureStoreMetrics, ProjectResolver projectResolver, FeatureService featureService, + SamplingService samplingService, Consumer> nodeInfoListener ) { this.clusterService = clusterService; @@ -282,6 +289,7 @@ public IngestService( this.failureStoreMetrics = failureStoreMetrics; this.projectResolver = projectResolver; this.featureService = featureService; + this.samplingService = samplingService; this.nodeInfoListener = nodeInfoListener; } @@ -296,7 +304,8 @@ public IngestService( MatcherWatchdog matcherWatchdog, FailureStoreMetrics failureStoreMetrics, ProjectResolver projectResolver, - FeatureService featureService + FeatureService featureService, + SamplingService samplingService ) { this( clusterService, @@ -310,6 +319,7 @@ public IngestService( failureStoreMetrics, projectResolver, featureService, + samplingService, createNodeInfoListener(client) ); } @@ -330,6 +340,7 @@ public IngestService( this.failureStoreMetrics = ingestService.failureStoreMetrics; this.projectResolver = ingestService.projectResolver; this.featureService = ingestService.featureService; + this.samplingService = ingestService.samplingService; this.nodeInfoListener = ingestService.nodeInfoListener; } @@ -989,7 +1000,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 documentListener = ActionListener.runAfter( @@ -1036,7 +1047,14 @@ public void onFailure(Exception e) { } ); - executePipelines(pipelines, indexRequest, ingestDocument, adaptedResolveFailureStore, documentListener); + executePipelines( + pipelines, + indexRequest, + ingestDocument, + adaptedResolveFailureStore, + documentListener, + originalDocumentMetadata + ); assert actionRequest.index() != null; i++; @@ -1155,7 +1173,8 @@ private void executePipelines( final IndexRequest indexRequest, final IngestDocument ingestDocument, final Function resolveFailureStore, - final ActionListener listener + final ActionListener listener, + final Metadata originalDocumentMetadata ) { assert pipelines.hasNext(); PipelineSlot slot = pipelines.next(); @@ -1341,15 +1360,38 @@ 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 + // Here is where we finally overwrite source. Somehow sample before this? + // But also somewhere else if there are no pipelines + // Can't do this in the listener though b/c the source is already gone + // byte[] originalBytes; + // BytesReference sourceBytesRef = indexRequest.source(); + // if (sourceBytesRef.hasArray()) { + // originalBytes = Arrays.copyOfRange(sourceBytesRef.array(), sourceBytesRef.arrayOffset(), sourceBytesRef.arrayOffset() + // + sourceBytesRef.length()); + // } else { + // originalBytes = sourceBytesRef.toBytesRef().bytes; + // } + // + // if sampling enabled for this index AND condition is met AND sample THEN + try { + IndexRequest sample = copyIndexRequest(indexRequest); + updateIndexRequestMetadata(sample, originalDocumentMetadata); + samplingService.maybeSample(project, sample, ingestDocument); + } catch (IOException ex) { + logger.warn("unable to sample data"); + } + updateIndexRequestSource(indexRequest, ingestDocument); cacheRawTimestamp(indexRequest, ingestDocument); listener.onResponse(IngestPipelinesExecutionResult.SUCCESSFUL_RESULT); // document succeeded! } }); } catch (Exception e) { + // 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 @@ -1358,6 +1400,19 @@ private void executePipelines( } } + private IndexRequest copyIndexRequest(IndexRequest original) throws IOException { + // This makes a whole new copy of the source, and removes it from the BytesReference. I think this is what we want here? That way + // if we have a huge bulk request and only one thing is sampled, we don't keep it all. + try (BytesStreamOutput output = new BytesStreamOutput()) { + output.setTransportVersion(TransportVersion.current()); + original.writeTo(output); + try (StreamInput in = output.bytes().streamInput()) { + in.setTransportVersion(TransportVersion.current()); + return new IndexRequest(in); + } + } + } + private static void executePipeline( final IngestDocument ingestDocument, final Pipeline pipeline, diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index c909c9a25e5b8..2f8abf53ee838 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -142,6 +142,7 @@ import org.elasticsearch.indices.recovery.plan.RecoveryPlannerService; import org.elasticsearch.indices.recovery.plan.ShardSnapshotsService; import org.elasticsearch.ingest.IngestService; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.injection.guice.Injector; import org.elasticsearch.injection.guice.Key; import org.elasticsearch.injection.guice.Module; @@ -713,6 +714,8 @@ private void construct( modules.bindToInstance(DocumentParsingProvider.class, documentParsingProvider); FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); + SamplingService samplingService = new SamplingService(scriptService); + modules.bindToInstance(SamplingService.class, samplingService); FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); final IngestService ingestService = new IngestService( @@ -726,7 +729,8 @@ private void construct( IngestService.createGrokThreadWatchdog(environment, threadPool), failureStoreMetrics, projectResolver, - featureService + featureService, + samplingService ); SystemIndices systemIndices = createSystemIndices(settings); diff --git a/server/src/main/java/org/elasticsearch/sample/PutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/PutSampleConfigAction.java new file mode 100644 index 0000000000000..5467e615fe59c --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/PutSampleConfigAction.java @@ -0,0 +1,160 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.ActionType; +import org.elasticsearch.action.IndicesRequest; +import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.action.support.master.AcknowledgedRequest; +import org.elasticsearch.action.support.master.AcknowledgedResponse; +import org.elasticsearch.common.Strings; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.unit.ByteSizeValue; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.tasks.CancellableTask; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.tasks.TaskId; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Map; +import java.util.Objects; + +public class PutSampleConfigAction extends ActionType { + public static final String NAME = "indices:admin/sample/config/update"; + public static final PutSampleConfigAction INSTANCE = new PutSampleConfigAction(); + + public PutSampleConfigAction() { + super(NAME); + } + + public static class Request extends AcknowledgedRequest implements IndicesRequest.Replaceable { + private final double rate; + private final Integer maxSamples; + private final ByteSizeValue maxSize; + private final TimeValue timeToLive; + private final String condition; + private String[] indices = Strings.EMPTY_ARRAY; + + public Request( + double rate, + @Nullable Integer maxSamples, + @Nullable ByteSizeValue maxSize, + @Nullable TimeValue timeToLive, + @Nullable String condition, + @Nullable TimeValue masterNodeTimeout, + @Nullable TimeValue ackTimeout + ) { + super(masterNodeTimeout, ackTimeout); + this.rate = rate; + this.maxSamples = maxSamples; + this.maxSize = maxSize; + this.timeToLive = timeToLive; + this.condition = condition; + } + + public double getRate() { + return rate; + } + + public Integer getMaxSamples() { + return maxSamples; + } + + public ByteSizeValue getMaxSize() { + return maxSize; + } + + public TimeValue getTimeToLive() { + return timeToLive; + } + + public String getCondition() { + return condition; + } + + @Override + public PutSampleConfigAction.Request indices(String... dataStreamNames) { + this.indices = dataStreamNames; + return this; + } + + @Override + public boolean includeDataStreams() { + return true; + } + + public Request(StreamInput in) throws IOException { + super(in); + this.indices = in.readStringArray(); + this.rate = in.readDouble(); + this.maxSamples = in.readOptionalInt(); + this.maxSize = in.readOptionalWriteable(ByteSizeValue::readFrom); + this.timeToLive = in.readOptionalTimeValue(); + this.condition = in.readOptionalString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeStringArray(indices); + out.writeDouble(rate); + out.writeOptionalInt(maxSamples); + out.writeOptionalWriteable(maxSize); + out.writeOptionalTimeValue(timeToLive); + out.writeOptionalString(condition); + } + + @Override + public String[] indices() { + return indices; + } + + @Override + public IndicesOptions indicesOptions() { + return IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED; + } + + @Override + public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { + return new CancellableTask(id, type, action, "", parentTaskId, headers); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + PutSampleConfigAction.Request request = (PutSampleConfigAction.Request) o; + return Arrays.equals(indices, request.indices()) + && rate == request.rate + && Objects.equals(maxSamples, request.maxSamples) + && Objects.equals(maxSize, request.maxSize) + && Objects.equals(timeToLive, request.timeToLive) + && Objects.equals(condition, request.condition); + } + + @Override + public int hashCode() { + return Objects.hash( + Arrays.hashCode(indices), + rate, + maxSamples, + maxSize, + timeToLive, + condition, + masterNodeTimeout(), + ackTimeout() + ); + } + + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java new file mode 100644 index 0000000000000..6592ec3643930 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java @@ -0,0 +1,79 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.support.master.AcknowledgedResponse; +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.common.Strings; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestChannel; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.rest.RestResponse; +import org.elasticsearch.rest.RestStatus; +import org.elasticsearch.rest.RestUtils; +import org.elasticsearch.rest.Scope; +import org.elasticsearch.rest.ServerlessScope; +import org.elasticsearch.rest.action.RestBuilderListener; +import org.elasticsearch.rest.action.RestCancellableNodeClient; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; +import java.util.List; + +import static org.elasticsearch.rest.RestRequest.Method.PUT; + +@ServerlessScope(Scope.PUBLIC) +public class RestPutSampleConfigAction extends BaseRestHandler { + @Override + public String getName() { + return "put_sample_config"; + } + + @Override + public List routes() { + return List.of(new Route(PUT, "/_sample/{name}/_config")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig; + try (XContentParser parser = request.contentParser()) { + samplingConfig = TransportPutSampleConfigAction.SamplingConfigCustomMetadata.fromXContent(parser); + } + PutSampleConfigAction.Request putSampleConfigRequest = new PutSampleConfigAction.Request( + samplingConfig.rate, + samplingConfig.maxSamples, + samplingConfig.maxSize, + samplingConfig.timeToLive, + samplingConfig.condition, + RestUtils.getMasterNodeTimeout(request), + RestUtils.getAckTimeout(request) + ).indices(Strings.splitStringByCommaToArray(request.param("name"))); + return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( + PutSampleConfigAction.INSTANCE, + putSampleConfigRequest, + new ReindexDataStreamRestToXContentListener(channel) + ); + } + + static class ReindexDataStreamRestToXContentListener extends RestBuilderListener { + + ReindexDataStreamRestToXContentListener(RestChannel channel) { + super(channel); + } + + @Override + public RestResponse buildResponse(AcknowledgedResponse response, XContentBuilder builder) throws Exception { + response.toXContent(builder, channel.request()); + return new RestResponse(RestStatus.OK, builder); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java new file mode 100644 index 0000000000000..433499750368e --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java @@ -0,0 +1,313 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.TransportVersion; +import org.elasticsearch.TransportVersions; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.support.ActionFilters; +import org.elasticsearch.action.support.master.AcknowledgedResponse; +import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction; +import org.elasticsearch.cluster.AbstractNamedDiffable; +import org.elasticsearch.cluster.AckedBatchedClusterStateUpdateTask; +import org.elasticsearch.cluster.ClusterState; +import org.elasticsearch.cluster.ClusterStateAckListener; +import org.elasticsearch.cluster.ClusterStateTaskExecutor; +import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor; +import org.elasticsearch.cluster.block.ClusterBlockException; +import org.elasticsearch.cluster.metadata.Metadata; +import org.elasticsearch.cluster.metadata.ProjectId; +import org.elasticsearch.cluster.metadata.ProjectMetadata; +import org.elasticsearch.cluster.project.ProjectResolver; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.cluster.service.MasterServiceTaskQueue; +import org.elasticsearch.common.Priority; +import org.elasticsearch.common.collect.Iterators; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.unit.ByteSizeUnit; +import org.elasticsearch.common.unit.ByteSizeValue; +import org.elasticsearch.common.util.concurrent.EsExecutors; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.core.Tuple; +import org.elasticsearch.injection.guice.Inject; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.TransportService; +import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.ToXContent; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; +import java.util.EnumSet; +import java.util.Iterator; + +import static org.elasticsearch.cluster.metadata.Metadata.ALL_CONTEXTS; +import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; + +public class TransportPutSampleConfigAction extends AcknowledgedTransportMasterNodeAction { + private static final Logger logger = LogManager.getLogger(TransportPutSampleConfigAction.class); + private final ProjectResolver projectResolver; + private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; + + @Inject + public TransportPutSampleConfigAction( + TransportService transportService, + ClusterService clusterService, + ThreadPool threadPool, + ActionFilters actionFilters, + ProjectResolver projectResolver + ) { + super( + PutSampleConfigAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + PutSampleConfigAction.Request::new, + EsExecutors.DIRECT_EXECUTOR_SERVICE + ); + this.projectResolver = projectResolver; + ClusterStateTaskExecutor updateMappingsExecutor = new SimpleBatchedAckListenerTaskExecutor<>() { + + @Override + public Tuple executeTask( + UpdateSampleConfigTask updateSamplingConfigTask, + ClusterState clusterState + ) throws Exception { + ProjectMetadata projectMetadata = clusterState.metadata().getProject(updateSamplingConfigTask.projectId); + SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom(SamplingConfigCustomMetadata.NAME); + ProjectMetadata.Builder projectMetadataBuilder = ProjectMetadata.builder(projectMetadata); + projectMetadataBuilder.putCustom( + SamplingConfigCustomMetadata.NAME, + new SamplingConfigCustomMetadata( + updateSamplingConfigTask.indexName, + updateSamplingConfigTask.rate, + updateSamplingConfigTask.maxSamples, + updateSamplingConfigTask.maxSize, + updateSamplingConfigTask.timeToLive, + updateSamplingConfigTask.condition + ) + ); + ClusterState updatedClusterState = ClusterState.builder(clusterState).putProjectMetadata(projectMetadataBuilder).build(); + return new Tuple<>(updatedClusterState, updateSamplingConfigTask); + } + }; + this.updateSamplingConfigTaskQueue = clusterService.createTaskQueue( + "update-data-stream-mappings", + Priority.NORMAL, + updateMappingsExecutor + ); + } + + @Override + protected void masterOperation( + Task task, + PutSampleConfigAction.Request request, + ClusterState state, + ActionListener listener + ) throws Exception { + ProjectId projectId = projectResolver.getProjectId(); + updateSamplingConfigTaskQueue.submitTask( + "updating mappings on data stream", + new UpdateSampleConfigTask( + projectId, + request.indices()[0], + request.getRate(), + request.getMaxSamples(), + request.getMaxSize(), + request.getTimeToLive(), + request.getCondition(), + request.ackTimeout(), + listener + ), + request.masterNodeTimeout() + ); + state.projectState(projectResolver.getProjectId()).metadata().custom("sample_config"); + } + + @Override + protected ClusterBlockException checkBlock(PutSampleConfigAction.Request request, ClusterState state) { + return null; + } + + static class UpdateSampleConfigTask extends AckedBatchedClusterStateUpdateTask { + final ProjectId projectId; + private final String indexName; + private final double rate; + private final Integer maxSamples; + private final ByteSizeValue maxSize; + private final TimeValue timeToLive; + private final String condition; + + UpdateSampleConfigTask( + ProjectId projectId, + String indexName, + double rate, + Integer maxSamples, + ByteSizeValue maxSize, + TimeValue timeToLive, + String condition, + TimeValue ackTimeout, + ActionListener listener + ) { + super(ackTimeout, listener); + this.projectId = projectId; + this.indexName = indexName; + this.rate = rate; + this.maxSamples = maxSamples; + this.maxSize = maxSize; + this.timeToLive = timeToLive; + this.condition = condition; + } + } + + public static final class SamplingConfigCustomMetadata extends AbstractNamedDiffable + implements + Metadata.ProjectCustom { + public final String indexName; + public final double rate; + public final Integer maxSamples; + public final ByteSizeValue maxSize; + public final TimeValue timeToLive; + public final String condition; + + public static final String NAME = "sampling_config"; + public static final ParseField INDEX_NAME_FIELD = new ParseField("index_name"); + public static final ParseField RATE_FIELD = new ParseField("rate"); + public static final ParseField MAX_SAMPLES_FIELD = new ParseField("max_samples"); + public static final ParseField MAX_SIZE_FIELD = new ParseField("max_size"); + public static final ParseField TIME_TO_LIVE_FIELD = new ParseField("time_to_live"); + public static final ParseField CONDITION_FIELD = new ParseField("condition"); + + @SuppressWarnings("unchecked") + private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + NAME, + args -> new SamplingConfigCustomMetadata( + (String) args[0], + (double) args[1], + (Integer) args[2], + (Long) args[3], + args[4] == null ? null : TimeValue.timeValueMillis((long) args[4]), + (String) args[5] + ) + ); + + static { + PARSER.declareString(constructorArg(), INDEX_NAME_FIELD); + PARSER.declareDouble(constructorArg(), RATE_FIELD); + PARSER.declareInt(optionalConstructorArg(), MAX_SAMPLES_FIELD); + PARSER.declareLong(optionalConstructorArg(), MAX_SIZE_FIELD); + PARSER.declareLong(optionalConstructorArg(), TIME_TO_LIVE_FIELD); + PARSER.declareString(optionalConstructorArg(), CONDITION_FIELD); + } + + public SamplingConfigCustomMetadata( + String indexName, + double rate, + Integer maxSamples, + ByteSizeValue maxSize, + TimeValue timeToLive, + String condition + ) { + this.indexName = indexName; + this.rate = rate; + this.maxSamples = maxSamples; + this.maxSize = maxSize; + this.timeToLive = timeToLive; + this.condition = condition; + } + + public SamplingConfigCustomMetadata( + String indexName, + double rate, + Integer maxSamples, + Long maxSizeInBytes, + TimeValue timeToLive, + String condition + ) { + this( + indexName, + rate, + maxSamples, + maxSizeInBytes == null ? null : ByteSizeValue.of(maxSizeInBytes, ByteSizeUnit.BYTES), + timeToLive, + condition + ); + } + + public SamplingConfigCustomMetadata(StreamInput in) throws IOException { + this( + in.readString(), + in.readDouble(), + in.readOptionalInt(), + in.readOptionalLong(), + in.readOptionalTimeValue(), + in.readOptionalString() + ); + } + + @Override + public EnumSet context() { + return ALL_CONTEXTS; + } + + @Override + public TransportVersion getMinimalSupportedVersion() { + return TransportVersions.RANDOM_SAMPLING; + } + + @Override + public String getWriteableName() { + return "sample_config"; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(indexName); + out.writeDouble(rate); + out.writeOptionalInt(maxSamples); + out.writeOptionalLong(maxSize == null ? null : maxSize.getBytes()); + out.writeOptionalTimeValue(timeToLive); + out.writeOptionalString(condition); + } + + @Override + public Iterator toXContentChunked(ToXContent.Params params) { + return Iterators.single((b, p) -> { + // b.startObject(); + b.field(INDEX_NAME_FIELD.getPreferredName(), indexName); + b.field(RATE_FIELD.getPreferredName(), rate); + if (maxSamples != null) { + b.field(MAX_SAMPLES_FIELD.getPreferredName(), maxSamples); + } + if (maxSize != null) { + b.field(MAX_SIZE_FIELD.getPreferredName(), maxSize.getBytes()); + } + if (timeToLive != null) { + b.field(TIME_TO_LIVE_FIELD.getPreferredName(), timeToLive.millis()); + } + if (condition != null) { + b.field(CONDITION_FIELD.getPreferredName(), condition); + } + // b.endObject(); + return b; + }); + } + + public static SamplingConfigCustomMetadata fromXContent(XContentParser parser) throws IOException { + return PARSER.apply(parser, null); + } + } +} From 7e4f30d9404a070c999c3b6167fbeed3de8b827e Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 5 Sep 2025 10:54:29 -0500 Subject: [PATCH 02/37] reverting accidental change --- muted-tests.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/muted-tests.yml b/muted-tests.yml index e5f6a24f8444d..dad57b96e9846 100644 --- a/muted-tests.yml +++ b/muted-tests.yml @@ -414,6 +414,9 @@ tests: - class: org.elasticsearch.xpack.ml.integration.RevertModelSnapshotIT method: testRevertModelSnapshot_DeleteInterveningResults issue: https://github.com/elastic/elasticsearch/issues/132349 +#- class: org.elasticsearch.xpack.ml.integration.TextEmbeddingQueryIT +# method: testHybridSearch +# issue: https://github.com/elastic/elasticsearch/issues/132703 - class: org.elasticsearch.xpack.ml.integration.RevertModelSnapshotIT method: testRevertModelSnapshot issue: https://github.com/elastic/elasticsearch/issues/132733 From d6e40e36ecc905eea065198f396c3314dd601555 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 5 Sep 2025 10:55:54 -0500 Subject: [PATCH 03/37] adding files --- .../elasticsearch/ingest/SamplingService.java | 114 ++++++++++++++ .../elasticsearch/sample/GetSampleAction.java | 141 ++++++++++++++++++ .../sample/RestGetSampleAction.java | 43 ++++++ .../sample/TransportGetSampleAction.java | 48 ++++++ 4 files changed, 346 insertions(+) create mode 100644 server/src/main/java/org/elasticsearch/ingest/SamplingService.java create mode 100644 server/src/main/java/org/elasticsearch/sample/GetSampleAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java new file mode 100644 index 0000000000000..078d92a7e41d7 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -0,0 +1,114 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.ingest; + +import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.cluster.metadata.ProjectMetadata; +import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; +import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.plugins.internal.XContentParserDecorator; +import org.elasticsearch.sample.TransportPutSampleConfigAction; +import org.elasticsearch.script.DynamicMap; +import org.elasticsearch.script.IngestConditionalScript; +import org.elasticsearch.script.Script; +import org.elasticsearch.script.ScriptService; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xcontent.XContentType; +import org.elasticsearch.xcontent.json.JsonXContent; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.ingest.ConditionalProcessor.FUNCTIONS; + +public class SamplingService { + private final ScriptService scriptService; + private final Map> samples = new HashMap<>(); + + public SamplingService(ScriptService scriptService) { + this.scriptService = scriptService; + } + + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { + maybeSample( + projectMetadata, + indexRequest, + new IngestDocument( + indexRequest.index(), + indexRequest.id(), + indexRequest.version(), + indexRequest.routing(), + indexRequest.versionType(), + indexRequest.sourceAsMap(XContentParserDecorator.NOOP) + ) + ); + } + + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) throws IOException { + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME + ); + if (samplingConfig != null) { + String samplingIndex = samplingConfig.indexName; + if (samplingIndex.equals(indexRequest.index())) { + List samplesForIndex = samples.computeIfAbsent(samplingIndex, k -> new ArrayList<>()); + if (samplesForIndex.size() < samplingConfig.maxSamples) { + String condition = samplingConfig.condition; + if (evaluateCondition(ingestDocument, condition)) { + if (Math.random() < samplingConfig.rate) { + samplesForIndex.add(indexRequest); + System.out.println("Sampling " + indexRequest); + } + } + } + } + } + } + + public List getSamples(String index) { + return samples.get(index); + } + + private boolean evaluateCondition(IngestDocument ingestDocument, String condition) throws IOException { + if (condition == null) { + return true; + } + IngestConditionalScript.Factory factory = null;// precompiledConditionalScriptFactory; + Script script = getScript(condition); + if (factory == null) { + factory = scriptService.compile(script, IngestConditionalScript.CONTEXT); + } + return factory.newInstance( + script.getParams(), + new ConditionalProcessor.UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)) + ).execute(); + } + + private static Script getScript(String conditional) throws IOException { + if (conditional != null) { + try ( + XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent).map(Map.of("source", conditional)); + XContentParser parser = XContentHelper.createParserNotCompressed( + LoggingDeprecationHandler.XCONTENT_PARSER_CONFIG, + BytesReference.bytes(builder), + XContentType.JSON + ) + ) { + return Script.parse(parser); + } + } + return null; + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java new file mode 100644 index 0000000000000..eed6ce2763ed4 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java @@ -0,0 +1,141 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.ActionRequest; +import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.action.ActionResponse; +import org.elasticsearch.action.ActionType; +import org.elasticsearch.action.IndicesRequest; +import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.common.collect.Iterators; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.xcontent.ChunkedToXContent; +import org.elasticsearch.tasks.CancellableTask; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.tasks.TaskId; +import org.elasticsearch.xcontent.ToXContent; + +import java.io.IOException; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +import static org.elasticsearch.common.collect.Iterators.single; +import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.chunk; + +public class GetSampleAction extends ActionType { + + public static final GetSampleAction INSTANCE = new GetSampleAction(); + public static final String NAME = "indices:admin/sample"; + + private GetSampleAction() { + super(NAME); + } + + public static class Response extends ActionResponse implements ChunkedToXContent { + + private final List samples; + + public Response(final List samples) { + this.samples = samples; + } + + public List getSamples() { + return samples; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeCollection(samples); + } + + @Override + public Iterator toXContentChunked(ToXContent.Params params) { + return Iterators.concat( + chunk((builder, p) -> builder.startObject().startArray("samples")), + Iterators.flatMap(samples.iterator(), sample -> single((builder, params1) -> { + Map source = sample.sourceAsMap(); + builder.value(source); + return builder; + })), + chunk((builder, p) -> builder.endArray().endObject()) + ); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + GetSampleAction.Response response = (GetSampleAction.Response) o; + return samples.equals(response.samples); + } + + @Override + public int hashCode() { + return Objects.hash(samples); + } + + @Override + public String toString() { + return "Response{samples=" + samples + '}'; + } + } + + public static class Request extends ActionRequest implements IndicesRequest.Replaceable { + private String[] names; + + public Request(String[] names) { + super(); + this.names = names; + } + + public Request(StreamInput in) throws IOException { + super(in); + this.names = in.readStringArray(); + } + + @Override + public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { + return new CancellableTask(id, type, action, "", parentTaskId, headers); + } + + @Override + public ActionRequestValidationException validate() { + if (this.indices().length != 1) { + return new ActionRequestValidationException(); + } + return null; + } + + @Override + public IndicesRequest indices(String... indices) { + this.names = indices; + return this; + } + + @Override + public String[] indices() { + return names; + } + + @Override + public IndicesOptions indicesOptions() { + return IndicesOptions.DEFAULT; + } + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java new file mode 100644 index 0000000000000..abe16175c8765 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java @@ -0,0 +1,43 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.rest.action.RestCancellableNodeClient; +import org.elasticsearch.rest.action.RestRefCountedChunkedToXContentListener; + +import java.io.IOException; +import java.util.List; + +import static org.elasticsearch.rest.RestRequest.Method.GET; + +public class RestGetSampleAction extends BaseRestHandler { + @Override + public String getName() { + return "get_sample"; + } + + @Override + public List routes() { + return List.of(new Route(GET, "/_sample/{name}")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + GetSampleAction.Request getSampleRequest = new GetSampleAction.Request(new String[] { request.param("name") }); + return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( + GetSampleAction.INSTANCE, + getSampleRequest, + new RestRefCountedChunkedToXContentListener<>(channel) + ); + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java new file mode 100644 index 0000000000000..cd0b591db4eb4 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java @@ -0,0 +1,48 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.action.support.ActionFilters; +import org.elasticsearch.action.support.HandledTransportAction; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.util.concurrent.EsExecutors; +import org.elasticsearch.ingest.SamplingService; +import org.elasticsearch.injection.guice.Inject; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.TransportService; + +import java.util.List; + +public class TransportGetSampleAction extends HandledTransportAction { + private final SamplingService samplingService; + + @Inject + public TransportGetSampleAction( + TransportService transportService, + ClusterService clusterService, + ThreadPool threadPool, + ActionFilters actionFilters, + SamplingService samplingService + ) { + super(GetSampleAction.NAME, transportService, actionFilters, GetSampleAction.Request::new, EsExecutors.DIRECT_EXECUTOR_SERVICE); + this.samplingService = samplingService; + } + + @Override + protected void doExecute(Task task, GetSampleAction.Request request, ActionListener listener) { + String index = request.indices()[0]; + List samples = samplingService.getSamples(index); + GetSampleAction.Response response = new GetSampleAction.Response(samples); + listener.onResponse(response); + } +} From 713384654cfe4287a1d3e367a42f6797b19dbb48 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 5 Sep 2025 15:25:48 -0500 Subject: [PATCH 04/37] making the transport action a TransportNodesAction --- .../elasticsearch/cluster/ClusterModule.java | 14 +++ .../elasticsearch/indices/IndicesModule.java | 10 +- .../elasticsearch/ingest/SamplingService.java | 3 + .../elasticsearch/sample/GetSampleAction.java | 117 +++++++++++++----- .../sample/TransportGetSampleAction.java | 47 +++++-- .../TransportPutSampleConfigAction.java | 7 +- 6 files changed, 151 insertions(+), 47 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/cluster/ClusterModule.java b/server/src/main/java/org/elasticsearch/cluster/ClusterModule.java index 2522001b94168..e56896917ef45 100644 --- a/server/src/main/java/org/elasticsearch/cluster/ClusterModule.java +++ b/server/src/main/java/org/elasticsearch/cluster/ClusterModule.java @@ -91,6 +91,7 @@ import org.elasticsearch.persistent.PersistentTasksCustomMetadata; import org.elasticsearch.persistent.PersistentTasksNodeService; import org.elasticsearch.plugins.ClusterPlugin; +import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.script.ScriptMetadata; import org.elasticsearch.snapshots.RegisteredPolicySnapshots; import org.elasticsearch.snapshots.SnapshotsInfoService; @@ -278,6 +279,12 @@ public static List getNamedWriteables() { PersistentTasksCustomMetadata::new, PersistentTasksCustomMetadata::readDiffFrom ); + registerProjectCustom( + entries, + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME, + TransportPutSampleConfigAction.SamplingConfigCustomMetadata::new, + TransportPutSampleConfigAction.SamplingConfigCustomMetadata::readDiffFrom + ); // Cluster scoped persistent tasks registerMetadataCustom( entries, @@ -358,6 +365,13 @@ public static List getNamedXWriteables() { ClusterPersistentTasksCustomMetadata::fromXContent ) ); + entries.add( + new NamedXContentRegistry.Entry( + Metadata.ProjectCustom.class, + new ParseField(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME), + TransportPutSampleConfigAction.SamplingConfigCustomMetadata::fromXContent + ) + ); entries.add( new NamedXContentRegistry.Entry( Metadata.ProjectCustom.class, diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java index 1790a2397d8e3..9e64dbffc7be6 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java @@ -22,7 +22,6 @@ import org.elasticsearch.action.admin.indices.rollover.MinSizeCondition; import org.elasticsearch.action.admin.indices.rollover.OptimalShardCountCondition; import org.elasticsearch.action.resync.TransportResyncReplicationAction; -import org.elasticsearch.cluster.metadata.Metadata; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.index.mapper.BinaryFieldMapper; import org.elasticsearch.index.mapper.BooleanFieldMapper; @@ -78,7 +77,6 @@ import org.elasticsearch.injection.guice.AbstractModule; import org.elasticsearch.plugins.FieldPredicate; import org.elasticsearch.plugins.MapperPlugin; -import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.xcontent.NamedXContentRegistry; import org.elasticsearch.xcontent.ParseField; @@ -116,13 +114,7 @@ public static List getNamedWriteables() { new NamedWriteableRegistry.Entry(Condition.class, MaxDocsCondition.NAME, MaxDocsCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxSizeCondition.NAME, MaxSizeCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardSizeCondition.NAME, MaxPrimaryShardSizeCondition::new), - new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardDocsCondition.NAME, MaxPrimaryShardDocsCondition::new), - new NamedWriteableRegistry.Entry(Condition.class, OptimalShardCountCondition.NAME, OptimalShardCountCondition::new), - new NamedWriteableRegistry.Entry( - Metadata.ProjectCustom.class, - TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME, - TransportPutSampleConfigAction.SamplingConfigCustomMetadata::new - ) + new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardDocsCondition.NAME, MaxPrimaryShardDocsCondition::new) ); } diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 078d92a7e41d7..a2b594e442095 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -12,6 +12,7 @@ import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.cluster.metadata.ProjectMetadata; import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.plugins.internal.XContentParserDecorator; @@ -68,6 +69,8 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque String condition = samplingConfig.condition; if (evaluateCondition(ingestDocument, condition)) { if (Math.random() < samplingConfig.rate) { + indexRequest.incRef(); + ((ReleasableBytesReference) indexRequest.source()).incRef(); samplesForIndex.add(indexRequest); System.out.println("Sampling " + indexRequest); } diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java index eed6ce2763ed4..8aebf8ee7b206 100644 --- a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java @@ -9,23 +9,30 @@ package org.elasticsearch.sample; -import org.elasticsearch.action.ActionRequest; import org.elasticsearch.action.ActionRequestValidationException; -import org.elasticsearch.action.ActionResponse; import org.elasticsearch.action.ActionType; +import org.elasticsearch.action.FailedNodeException; import org.elasticsearch.action.IndicesRequest; import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.action.support.nodes.BaseNodeResponse; +import org.elasticsearch.action.support.nodes.BaseNodesRequest; +import org.elasticsearch.action.support.nodes.BaseNodesResponse; +import org.elasticsearch.cluster.ClusterName; +import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.collect.Iterators; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.xcontent.ChunkedToXContent; import org.elasticsearch.tasks.CancellableTask; import org.elasticsearch.tasks.Task; import org.elasticsearch.tasks.TaskId; +import org.elasticsearch.transport.AbstractTransportRequest; import org.elasticsearch.xcontent.ToXContent; import java.io.IOException; +import java.util.Collection; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -43,28 +50,35 @@ private GetSampleAction() { super(NAME); } - public static class Response extends ActionResponse implements ChunkedToXContent { + public static class Response extends BaseNodesResponse implements Writeable, ChunkedToXContent { - private final List samples; + public Response(StreamInput in) throws IOException { + super(in); + } - public Response(final List samples) { - this.samples = samples; + public Response(ClusterName clusterName, List nodes, List failures) { + super(clusterName, nodes, failures); } public List getSamples() { - return samples; + return getNodes().stream().map(n -> n.samples).filter(Objects::nonNull).flatMap(Collection::stream).toList(); } @Override - public void writeTo(StreamOutput out) throws IOException { - out.writeCollection(samples); + protected List readNodesFrom(StreamInput in) throws IOException { + return in.readCollectionAsList(NodeResponse::new); + } + + @Override + protected void writeNodesTo(StreamOutput out, List nodes) throws IOException { + out.writeCollection(nodes); } @Override public Iterator toXContentChunked(ToXContent.Params params) { return Iterators.concat( chunk((builder, p) -> builder.startObject().startArray("samples")), - Iterators.flatMap(samples.iterator(), sample -> single((builder, params1) -> { + Iterators.flatMap(getSamples().iterator(), sample -> single((builder, params1) -> { Map source = sample.sourceAsMap(); builder.value(source); return builder; @@ -75,40 +89,64 @@ public Iterator toXContentChunked(ToXContent.Params params @Override public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - GetSampleAction.Response response = (GetSampleAction.Response) o; - return samples.equals(response.samples); + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + Response that = (Response) o; + return Objects.equals(getNodes(), that.getNodes()) && Objects.equals(failures(), that.failures()); } @Override public int hashCode() { - return Objects.hash(samples); + return Objects.hash(getNodes(), failures()); + } + + } + + public static class NodeResponse extends BaseNodeResponse { + private final List samples; + + protected NodeResponse(StreamInput in) throws IOException { + super(in); + samples = in.readCollectionAsList(IndexRequest::new); + } + + protected NodeResponse(DiscoveryNode node, List samples) { + super(node); + this.samples = samples; + } + + public List getSamples() { + return samples; } @Override - public String toString() { - return "Response{samples=" + samples + '}'; + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeCollection(samples); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + NodeResponse that = (NodeResponse) o; + return samples.equals(that.samples); + } + + @Override + public int hashCode() { + return Objects.hash(samples); } } - public static class Request extends ActionRequest implements IndicesRequest.Replaceable { + public static class Request extends BaseNodesRequest implements IndicesRequest.Replaceable { private String[] names; public Request(String[] names) { - super(); + super((String[]) null); this.names = names; } - public Request(StreamInput in) throws IOException { - super(in); - this.names = in.readStringArray(); - } - @Override public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { return new CancellableTask(id, type, action, "", parentTaskId, headers); @@ -138,4 +176,27 @@ public IndicesOptions indicesOptions() { return IndicesOptions.DEFAULT; } } + + public static class NodeRequest extends AbstractTransportRequest { + private final String index; + + public NodeRequest(String index) { + this.index = index; + } + + public NodeRequest(StreamInput in) throws IOException { + super(in); + this.index = in.readString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(index); + } + + public String getIndex() { + return index; + } + } } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java index cd0b591db4eb4..cb027daa3b697 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java @@ -9,21 +9,28 @@ package org.elasticsearch.sample; -import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.FailedNodeException; import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.ActionFilters; -import org.elasticsearch.action.support.HandledTransportAction; +import org.elasticsearch.action.support.nodes.TransportNodesAction; +import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.cluster.service.ClusterService; -import org.elasticsearch.common.util.concurrent.EsExecutors; +import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.injection.guice.Inject; import org.elasticsearch.tasks.Task; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.transport.TransportService; +import java.io.IOException; import java.util.List; -public class TransportGetSampleAction extends HandledTransportAction { +import static org.elasticsearch.sample.GetSampleAction.NodeRequest; +import static org.elasticsearch.sample.GetSampleAction.NodeResponse; +import static org.elasticsearch.sample.GetSampleAction.Request; +import static org.elasticsearch.sample.GetSampleAction.Response; + +public class TransportGetSampleAction extends TransportNodesAction { private final SamplingService samplingService; @Inject @@ -34,15 +41,37 @@ public TransportGetSampleAction( ActionFilters actionFilters, SamplingService samplingService ) { - super(GetSampleAction.NAME, transportService, actionFilters, GetSampleAction.Request::new, EsExecutors.DIRECT_EXECUTOR_SERVICE); + super( + GetSampleAction.NAME, + clusterService, + transportService, + actionFilters, + NodeRequest::new, + threadPool.executor(ThreadPool.Names.MANAGEMENT) + ); this.samplingService = samplingService; } + @SuppressWarnings("checkstyle:LineLength") + @Override + protected Response newResponse(Request request, List nodeResponses, List failures) { + return new Response(clusterService.getClusterName(), nodeResponses, failures); + } + + @Override + protected NodeRequest newNodeRequest(Request request) { + return new NodeRequest(request.indices()[0]); + } + + @Override + protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throws IOException { + return new NodeResponse(in); + } + @Override - protected void doExecute(Task task, GetSampleAction.Request request, ActionListener listener) { - String index = request.indices()[0]; + protected NodeResponse nodeOperation(NodeRequest request, Task task) { + String index = request.getIndex(); List samples = samplingService.getSamples(index); - GetSampleAction.Response response = new GetSampleAction.Response(samples); - listener.onResponse(response); + return new NodeResponse(transportService.getLocalNode(), samples == null ? List.of() : samples); } } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java index 433499750368e..08fbc451def08 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java @@ -22,6 +22,7 @@ import org.elasticsearch.cluster.ClusterState; import org.elasticsearch.cluster.ClusterStateAckListener; import org.elasticsearch.cluster.ClusterStateTaskExecutor; +import org.elasticsearch.cluster.NamedDiff; import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor; import org.elasticsearch.cluster.block.ClusterBlockException; import org.elasticsearch.cluster.metadata.Metadata; @@ -258,6 +259,10 @@ public SamplingConfigCustomMetadata(StreamInput in) throws IOException { ); } + public static NamedDiff readDiffFrom(StreamInput in) throws IOException { + return readDiffFrom(Metadata.ProjectCustom.class, NAME, in); + } + @Override public EnumSet context() { return ALL_CONTEXTS; @@ -270,7 +275,7 @@ public TransportVersion getMinimalSupportedVersion() { @Override public String getWriteableName() { - return "sample_config"; + return NAME; } @Override From 2f1122f0c96a7c021f4aad0095b6a2c157e8ab28 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 5 Sep 2025 15:27:16 -0500 Subject: [PATCH 05/37] removing accidental change --- .../src/main/java/org/elasticsearch/indices/IndicesModule.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java index 9e64dbffc7be6..09be98630d5c4 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesModule.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesModule.java @@ -114,7 +114,8 @@ public static List getNamedWriteables() { new NamedWriteableRegistry.Entry(Condition.class, MaxDocsCondition.NAME, MaxDocsCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxSizeCondition.NAME, MaxSizeCondition::new), new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardSizeCondition.NAME, MaxPrimaryShardSizeCondition::new), - new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardDocsCondition.NAME, MaxPrimaryShardDocsCondition::new) + new NamedWriteableRegistry.Entry(Condition.class, MaxPrimaryShardDocsCondition.NAME, MaxPrimaryShardDocsCondition::new), + new NamedWriteableRegistry.Entry(Condition.class, OptimalShardCountCondition.NAME, OptimalShardCountCondition::new) ); } From 9b3fa2ac7be2f09db21d16197c4981f7083021ae Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Mon, 8 Sep 2025 13:53:06 -0500 Subject: [PATCH 06/37] Avoiding compiling scripts multiple times, adding stats --- .../elasticsearch/ingest/SamplingService.java | 158 ++++++++++++++---- .../elasticsearch/node/NodeConstruction.java | 2 +- 2 files changed, 126 insertions(+), 34 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index a2b594e442095..3c6a8a73d6232 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -15,6 +15,8 @@ import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.logging.LogManager; +import org.elasticsearch.logging.Logger; import org.elasticsearch.plugins.internal.XContentParserDecorator; import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.script.DynamicMap; @@ -31,15 +33,20 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.LongAdder; +import java.util.function.LongSupplier; import static org.elasticsearch.ingest.ConditionalProcessor.FUNCTIONS; public class SamplingService { + private static final Logger logger = LogManager.getLogger(SamplingService.class); private final ScriptService scriptService; - private final Map> samples = new HashMap<>(); + private final LongSupplier relativeNanoTimeSupplier; + private final Map samples = new HashMap<>(); - public SamplingService(ScriptService scriptService) { + public SamplingService(ScriptService scriptService, LongSupplier relativeNanoTimeSupplier) { this.scriptService = scriptService; + this.relativeNanoTimeSupplier = relativeNanoTimeSupplier; } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { @@ -58,41 +65,72 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) throws IOException { + long startTime = relativeNanoTimeSupplier.getAsLong(); TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME ); if (samplingConfig != null) { String samplingIndex = samplingConfig.indexName; if (samplingIndex.equals(indexRequest.index())) { - List samplesForIndex = samples.computeIfAbsent(samplingIndex, k -> new ArrayList<>()); - if (samplesForIndex.size() < samplingConfig.maxSamples) { - String condition = samplingConfig.condition; - if (evaluateCondition(ingestDocument, condition)) { - if (Math.random() < samplingConfig.rate) { - indexRequest.incRef(); - ((ReleasableBytesReference) indexRequest.source()).incRef(); - samplesForIndex.add(indexRequest); - System.out.println("Sampling " + indexRequest); + SampleInfo sampleInfo = samples.computeIfAbsent(samplingIndex, k -> new SampleInfo()); + SampleStats stats = sampleInfo.stats; + stats.potentialSamples.increment(); + try { + if (sampleInfo.getSamples().size() < samplingConfig.maxSamples) { + String condition = samplingConfig.condition; + if (condition != null) { + if (sampleInfo.script == null || sampleInfo.factory == null) { + // We don't want to pay for synchronization because worst case, we compile the script twice + long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); + Script script = getScript(condition); + sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); + stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); + } } + long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); + if (condition == null + || evaluateCondition(ingestDocument, sampleInfo.script, sampleInfo.factory, sampleInfo.stats)) { + stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); + if (Math.random() < samplingConfig.rate) { + indexRequest.incRef(); + if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) { + releaseableSource.incRef(); + } + sampleInfo.getSamples().add(indexRequest); + stats.samples.increment(); + logger.info("Sampling " + indexRequest); + } else { + stats.samplesRejectedForRate.increment(); + } + } else { + stats.samplesRejectedForCondition.increment(); + } + } else { + stats.samplesRejectedForSize.increment(); } + } catch (Exception e) { + stats.samplesRejectedForException.increment(); + stats.lastException = e; + e.printStackTrace(System.out); + throw e; + } finally { + stats.timeSampling.add((relativeNanoTimeSupplier.getAsLong() - startTime)); + logger.info("********* Stats: " + stats); } } } } public List getSamples(String index) { - return samples.get(index); + return samples.get(index).getSamples(); } - private boolean evaluateCondition(IngestDocument ingestDocument, String condition) throws IOException { - if (condition == null) { - return true; - } - IngestConditionalScript.Factory factory = null;// precompiledConditionalScriptFactory; - Script script = getScript(condition); - if (factory == null) { - factory = scriptService.compile(script, IngestConditionalScript.CONTEXT); - } + private boolean evaluateCondition( + IngestDocument ingestDocument, + Script script, + IngestConditionalScript.Factory factory, + SampleStats stats + ) { return factory.newInstance( script.getParams(), new ConditionalProcessor.UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)) @@ -100,18 +138,72 @@ private boolean evaluateCondition(IngestDocument ingestDocument, String conditio } private static Script getScript(String conditional) throws IOException { - if (conditional != null) { - try ( - XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent).map(Map.of("source", conditional)); - XContentParser parser = XContentHelper.createParserNotCompressed( - LoggingDeprecationHandler.XCONTENT_PARSER_CONFIG, - BytesReference.bytes(builder), - XContentType.JSON - ) - ) { - return Script.parse(parser); - } + logger.info("Parsing script for conditional " + conditional); + try ( + XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent).map(Map.of("source", conditional)); + XContentParser parser = XContentHelper.createParserNotCompressed( + LoggingDeprecationHandler.XCONTENT_PARSER_CONFIG, + BytesReference.bytes(builder), + XContentType.JSON + ) + ) { + return Script.parse(parser); + } + } + + private static final class SampleInfo { + private final List samples; + private final SampleStats stats; + private volatile Script script; + private volatile IngestConditionalScript.Factory factory; + + SampleInfo() { + this.samples = new ArrayList<>(); + this.stats = new SampleStats(); + } + + public List getSamples() { + return samples; + } + + void setScript(Script script, IngestConditionalScript.Factory factory) { + this.script = script; + this.factory = factory; + } + } + + private static final class SampleStats { + LongAdder potentialSamples = new LongAdder(); + LongAdder samplesRejectedForSize = new LongAdder(); + LongAdder samplesRejectedForCondition = new LongAdder(); + LongAdder samplesRejectedForRate = new LongAdder(); + LongAdder samplesRejectedForException = new LongAdder(); + LongAdder samples = new LongAdder(); + LongAdder timeSampling = new LongAdder(); + LongAdder timeEvaluatingCondition = new LongAdder(); + LongAdder timeCompilingCondition = new LongAdder(); + Exception lastException = null; + + @Override + public String toString() { + return "potentialSamples: " + + potentialSamples + + ", samplesRejectedForSize: " + + samplesRejectedForSize + + ", samplesRejectedForCondition: " + + samplesRejectedForCondition + + ", samplesRejectedForRate: " + + samplesRejectedForRate + + ", samplesRejectedForException: " + + samplesRejectedForException + + ", samples: " + + samples + + ", timeSampling: " + + (timeSampling.longValue() / 1000000) + + ", timeEvaluatingCondition: " + + (timeEvaluatingCondition.longValue() / 1000000) + + ", timeCompilingCondition: " + + (timeCompilingCondition.longValue() / 1000000); } - return null; } } diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index 2f8abf53ee838..65ee7e3bc4895 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -714,7 +714,7 @@ private void construct( modules.bindToInstance(DocumentParsingProvider.class, documentParsingProvider); FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); - SamplingService samplingService = new SamplingService(scriptService); + SamplingService samplingService = new SamplingService(scriptService, System::nanoTime); modules.bindToInstance(SamplingService.class, samplingService); FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); From a8bf002ec37a1f94ddd2d665ed66cfc3989a8cec Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Mon, 8 Sep 2025 15:04:48 -0500 Subject: [PATCH 07/37] moved cluster state update logic into SamplingService --- .../elasticsearch/ingest/SamplingService.java | 133 +++++++++++++++++- .../elasticsearch/node/NodeConstruction.java | 2 +- .../TransportPutSampleConfigAction.java | 104 +++----------- 3 files changed, 145 insertions(+), 94 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 3c6a8a73d6232..c7b86622ec0dc 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -9,12 +9,27 @@ package org.elasticsearch.ingest; +import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.action.support.master.AcknowledgedResponse; +import org.elasticsearch.cluster.AckedBatchedClusterStateUpdateTask; +import org.elasticsearch.cluster.ClusterState; +import org.elasticsearch.cluster.ClusterStateAckListener; +import org.elasticsearch.cluster.ClusterStateTaskExecutor; +import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor; +import org.elasticsearch.cluster.metadata.ProjectId; import org.elasticsearch.cluster.metadata.ProjectMetadata; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.cluster.service.MasterServiceTaskQueue; +import org.elasticsearch.common.Priority; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.bytes.ReleasableBytesReference; +import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.core.Tuple; import org.elasticsearch.logging.LogManager; import org.elasticsearch.logging.Logger; import org.elasticsearch.plugins.internal.XContentParserDecorator; @@ -42,11 +57,73 @@ public class SamplingService { private static final Logger logger = LogManager.getLogger(SamplingService.class); private final ScriptService scriptService; private final LongSupplier relativeNanoTimeSupplier; + private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; private final Map samples = new HashMap<>(); - public SamplingService(ScriptService scriptService, LongSupplier relativeNanoTimeSupplier) { + public SamplingService(ScriptService scriptService, ClusterService clusterService, LongSupplier relativeNanoTimeSupplier) { this.scriptService = scriptService; this.relativeNanoTimeSupplier = relativeNanoTimeSupplier; + ClusterStateTaskExecutor updateSampleConfigExecutor = new SimpleBatchedAckListenerTaskExecutor<>() { + + @Override + public Tuple executeTask( + UpdateSampleConfigTask updateSamplingConfigTask, + ClusterState clusterState + ) { + ProjectMetadata projectMetadata = clusterState.metadata().getProject(updateSamplingConfigTask.projectId); + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME + ); + ProjectMetadata.Builder projectMetadataBuilder = ProjectMetadata.builder(projectMetadata); + projectMetadataBuilder.putCustom( + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME, + new TransportPutSampleConfigAction.SamplingConfigCustomMetadata( + updateSamplingConfigTask.indexName, + updateSamplingConfigTask.rate, + updateSamplingConfigTask.maxSamples, + updateSamplingConfigTask.maxSize, + updateSamplingConfigTask.timeToLive, + updateSamplingConfigTask.condition + ) + ); + ClusterState updatedClusterState = ClusterState.builder(clusterState).putProjectMetadata(projectMetadataBuilder).build(); + return new Tuple<>(updatedClusterState, updateSamplingConfigTask); + } + }; + this.updateSamplingConfigTaskQueue = clusterService.createTaskQueue( + "update-data-stream-mappings", + Priority.NORMAL, + updateSampleConfigExecutor + ); + } + + public void updateSampleConfiguration( + ProjectId projectId, + String index, + double rate, + @Nullable Integer maxSamples, + @Nullable ByteSizeValue maxSize, + @Nullable TimeValue timeToLive, + @Nullable String condition, + @Nullable TimeValue masterNodeTimeout, + @Nullable TimeValue ackTimeout, + ActionListener listener + ) { + updateSamplingConfigTaskQueue.submitTask( + "updating mappings on data stream", + new UpdateSampleConfigTask( + projectId, + index, + rate, + maxSamples, + maxSize, + timeToLive, + condition, + ackTimeout, + ActionListener.runBefore(listener, () -> samples.computeIfPresent(index, (s, sampleInfo) -> new SampleInfo())) + ), + masterNodeTimeout + ); } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { @@ -82,9 +159,21 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque if (sampleInfo.script == null || sampleInfo.factory == null) { // We don't want to pay for synchronization because worst case, we compile the script twice long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); - Script script = getScript(condition); - sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); - stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); + try { + if (sampleInfo.compilationFailed) { + // we don't want to waste time + stats.samplesRejectedForException.increment(); + return; + } else { + Script script = getScript(condition); + sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); + } + } catch (Exception e) { + sampleInfo.compilationFailed = true; + throw e; + } finally { + stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); + } } } long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); @@ -111,8 +200,8 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque } catch (Exception e) { stats.samplesRejectedForException.increment(); stats.lastException = e; - e.printStackTrace(System.out); - throw e; + logger.info("Error performing sampling for " + samplingIndex, e); + // e.printStackTrace(System.out); } finally { stats.timeSampling.add((relativeNanoTimeSupplier.getAsLong() - startTime)); logger.info("********* Stats: " + stats); @@ -156,6 +245,7 @@ private static final class SampleInfo { private final SampleStats stats; private volatile Script script; private volatile IngestConditionalScript.Factory factory; + private volatile boolean compilationFailed = false; SampleInfo() { this.samples = new ArrayList<>(); @@ -206,4 +296,35 @@ public String toString() { + (timeCompilingCondition.longValue() / 1000000); } } + + static class UpdateSampleConfigTask extends AckedBatchedClusterStateUpdateTask { + final ProjectId projectId; + private final String indexName; + private final double rate; + private final Integer maxSamples; + private final ByteSizeValue maxSize; + private final TimeValue timeToLive; + private final String condition; + + UpdateSampleConfigTask( + ProjectId projectId, + String indexName, + double rate, + Integer maxSamples, + ByteSizeValue maxSize, + TimeValue timeToLive, + String condition, + TimeValue ackTimeout, + ActionListener listener + ) { + super(ackTimeout, listener); + this.projectId = projectId; + this.indexName = indexName; + this.rate = rate; + this.maxSamples = maxSamples; + this.maxSize = maxSize; + this.timeToLive = timeToLive; + this.condition = condition; + } + } } diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index 65ee7e3bc4895..418762ba6ee47 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -714,7 +714,7 @@ private void construct( modules.bindToInstance(DocumentParsingProvider.class, documentParsingProvider); FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); - SamplingService samplingService = new SamplingService(scriptService, System::nanoTime); + SamplingService samplingService = new SamplingService(scriptService, clusterService, System::nanoTime); modules.bindToInstance(SamplingService.class, samplingService); FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); diff --git a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java index 08fbc451def08..f8c97955bd6d4 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java @@ -18,20 +18,13 @@ import org.elasticsearch.action.support.master.AcknowledgedResponse; import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction; import org.elasticsearch.cluster.AbstractNamedDiffable; -import org.elasticsearch.cluster.AckedBatchedClusterStateUpdateTask; import org.elasticsearch.cluster.ClusterState; -import org.elasticsearch.cluster.ClusterStateAckListener; -import org.elasticsearch.cluster.ClusterStateTaskExecutor; import org.elasticsearch.cluster.NamedDiff; -import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor; import org.elasticsearch.cluster.block.ClusterBlockException; import org.elasticsearch.cluster.metadata.Metadata; import org.elasticsearch.cluster.metadata.ProjectId; -import org.elasticsearch.cluster.metadata.ProjectMetadata; import org.elasticsearch.cluster.project.ProjectResolver; import org.elasticsearch.cluster.service.ClusterService; -import org.elasticsearch.cluster.service.MasterServiceTaskQueue; -import org.elasticsearch.common.Priority; import org.elasticsearch.common.collect.Iterators; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; @@ -39,7 +32,7 @@ import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.util.concurrent.EsExecutors; import org.elasticsearch.core.TimeValue; -import org.elasticsearch.core.Tuple; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.injection.guice.Inject; import org.elasticsearch.tasks.Task; import org.elasticsearch.threadpool.ThreadPool; @@ -60,7 +53,7 @@ public class TransportPutSampleConfigAction extends AcknowledgedTransportMasterNodeAction { private static final Logger logger = LogManager.getLogger(TransportPutSampleConfigAction.class); private final ProjectResolver projectResolver; - private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; + private final SamplingService samplingService; @Inject public TransportPutSampleConfigAction( @@ -68,7 +61,8 @@ public TransportPutSampleConfigAction( ClusterService clusterService, ThreadPool threadPool, ActionFilters actionFilters, - ProjectResolver projectResolver + ProjectResolver projectResolver, + SamplingService samplingService ) { super( PutSampleConfigAction.NAME, @@ -80,36 +74,7 @@ public TransportPutSampleConfigAction( EsExecutors.DIRECT_EXECUTOR_SERVICE ); this.projectResolver = projectResolver; - ClusterStateTaskExecutor updateMappingsExecutor = new SimpleBatchedAckListenerTaskExecutor<>() { - - @Override - public Tuple executeTask( - UpdateSampleConfigTask updateSamplingConfigTask, - ClusterState clusterState - ) throws Exception { - ProjectMetadata projectMetadata = clusterState.metadata().getProject(updateSamplingConfigTask.projectId); - SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom(SamplingConfigCustomMetadata.NAME); - ProjectMetadata.Builder projectMetadataBuilder = ProjectMetadata.builder(projectMetadata); - projectMetadataBuilder.putCustom( - SamplingConfigCustomMetadata.NAME, - new SamplingConfigCustomMetadata( - updateSamplingConfigTask.indexName, - updateSamplingConfigTask.rate, - updateSamplingConfigTask.maxSamples, - updateSamplingConfigTask.maxSize, - updateSamplingConfigTask.timeToLive, - updateSamplingConfigTask.condition - ) - ); - ClusterState updatedClusterState = ClusterState.builder(clusterState).putProjectMetadata(projectMetadataBuilder).build(); - return new Tuple<>(updatedClusterState, updateSamplingConfigTask); - } - }; - this.updateSamplingConfigTaskQueue = clusterService.createTaskQueue( - "update-data-stream-mappings", - Priority.NORMAL, - updateMappingsExecutor - ); + this.samplingService = samplingService; } @Override @@ -120,22 +85,19 @@ protected void masterOperation( ActionListener listener ) throws Exception { ProjectId projectId = projectResolver.getProjectId(); - updateSamplingConfigTaskQueue.submitTask( - "updating mappings on data stream", - new UpdateSampleConfigTask( - projectId, - request.indices()[0], - request.getRate(), - request.getMaxSamples(), - request.getMaxSize(), - request.getTimeToLive(), - request.getCondition(), - request.ackTimeout(), - listener - ), - request.masterNodeTimeout() + samplingService.updateSampleConfiguration( + projectId, + request.indices()[0], + request.getRate(), + request.getMaxSamples(), + request.getMaxSize(), + request.getTimeToLive(), + request.getCondition(), + request.masterNodeTimeout(), + request.ackTimeout(), + listener ); - state.projectState(projectResolver.getProjectId()).metadata().custom("sample_config"); + state.projectState(projectId).metadata().custom("sample_config"); } @Override @@ -143,37 +105,6 @@ protected ClusterBlockException checkBlock(PutSampleConfigAction.Request request return null; } - static class UpdateSampleConfigTask extends AckedBatchedClusterStateUpdateTask { - final ProjectId projectId; - private final String indexName; - private final double rate; - private final Integer maxSamples; - private final ByteSizeValue maxSize; - private final TimeValue timeToLive; - private final String condition; - - UpdateSampleConfigTask( - ProjectId projectId, - String indexName, - double rate, - Integer maxSamples, - ByteSizeValue maxSize, - TimeValue timeToLive, - String condition, - TimeValue ackTimeout, - ActionListener listener - ) { - super(ackTimeout, listener); - this.projectId = projectId; - this.indexName = indexName; - this.rate = rate; - this.maxSamples = maxSamples; - this.maxSize = maxSize; - this.timeToLive = timeToLive; - this.condition = condition; - } - } - public static final class SamplingConfigCustomMetadata extends AbstractNamedDiffable implements Metadata.ProjectCustom { @@ -306,7 +237,6 @@ public Iterator toXContentChunked(ToXContent.Params params if (condition != null) { b.field(CONDITION_FIELD.getPreferredName(), condition); } - // b.endObject(); return b; }); } From ca393b9daf03d6810791fdd20b01a6c45d016b7a Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Mon, 8 Sep 2025 15:19:06 -0500 Subject: [PATCH 08/37] fixing tests --- .../bulk/TransportBulkActionIngestTests.java | 4 +++- .../action/bulk/TransportBulkActionTests.java | 4 +++- .../bulk/TransportBulkActionTookTests.java | 3 ++- .../elasticsearch/ingest/IngestServiceTests.java | 16 +++++++++++----- .../ingest/SimulateIngestServiceTests.java | 3 ++- .../snapshots/SnapshotResiliencyTests.java | 6 ++++-- 6 files changed, 25 insertions(+), 11 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java index a54cd08c3738a..23b9f50a1d549 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java +++ b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java @@ -54,6 +54,7 @@ import org.elasticsearch.indices.EmptySystemIndices; import org.elasticsearch.indices.TestIndexNameExpressionResolver; import org.elasticsearch.ingest.IngestService; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.tasks.Task; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; @@ -172,7 +173,8 @@ class TestTransportBulkAction extends TransportBulkAction { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java index 481fdf5ea3530..20c3ee88b3de0 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java +++ b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java @@ -63,6 +63,7 @@ import org.elasticsearch.index.VersionType; import org.elasticsearch.indices.SystemIndexDescriptorUtils; import org.elasticsearch.indices.SystemIndices; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.VersionUtils; import org.elasticsearch.test.index.IndexVersionUtils; @@ -152,7 +153,8 @@ public ProjectId getProjectId() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java index 0077f739bf7b6..3b615b960d222 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java +++ b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java @@ -267,7 +267,8 @@ static class TestTransportBulkAction extends TransportBulkAction { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + null ); } } diff --git a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java index df3211331c7f5..21ba2e50e8f3a 100644 --- a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java @@ -171,7 +171,8 @@ public void testIngestPlugin() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); Map factories = ingestService.getProcessorFactories(); assertTrue(factories.containsKey("foo")); @@ -198,7 +199,8 @@ public void testIngestPluginDuplicate() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ) ); assertTrue(e.getMessage(), e.getMessage().contains("already registered")); @@ -222,7 +224,8 @@ public void testExecuteIndexPipelineDoesNotExist() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); final IndexRequest indexRequest = new IndexRequest("_index").id("_id") .source(Map.of()) @@ -2567,7 +2570,8 @@ public Map getProcessors(Processor.Parameters paramet public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); ingestService.addIngestClusterStateListener(ingestClusterStateListener); @@ -3075,6 +3079,7 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } }, + mock(SamplingService.class), consumer ); ingestService.applyClusterState(new ClusterChangedEvent("", clusterState, clusterState)); @@ -3406,7 +3411,8 @@ public Map getProcessors(final Processor.Parameters p public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return featureTest.test(feature); } - } + }, + mock(SamplingService.class) ); if (randomBoolean()) { /* diff --git a/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java index 6733ff4c44e6e..79b4e76932fb3 100644 --- a/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java @@ -173,7 +173,8 @@ public Map getProcessors(final Processor.Parameters p public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } } diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java index 6a459cab07328..91446eb5cfc4b 100644 --- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java @@ -2666,7 +2666,8 @@ public RecyclerBytesStreamOutput newNetworkBytesStream() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + null ), client, actionFilters, @@ -2681,7 +2682,8 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + null ) ); final TransportShardBulkAction transportShardBulkAction = new TransportShardBulkAction( From ca24196ca34e33071d937657b6a02ff56c1ccb01 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 09:29:22 -0500 Subject: [PATCH 09/37] enforcing TTL --- .../elasticsearch/ingest/SamplingService.java | 121 +++++++++++++++--- .../elasticsearch/node/NodeConstruction.java | 1 + .../TransportPutSampleConfigAction.java | 21 ++- 3 files changed, 126 insertions(+), 17 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index c7b86622ec0dc..e0b5461809aa9 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -13,8 +13,10 @@ import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.master.AcknowledgedResponse; import org.elasticsearch.cluster.AckedBatchedClusterStateUpdateTask; +import org.elasticsearch.cluster.ClusterChangedEvent; import org.elasticsearch.cluster.ClusterState; import org.elasticsearch.cluster.ClusterStateAckListener; +import org.elasticsearch.cluster.ClusterStateListener; import org.elasticsearch.cluster.ClusterStateTaskExecutor; import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor; import org.elasticsearch.cluster.metadata.ProjectId; @@ -48,16 +50,18 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.atomic.LongAdder; import java.util.function.LongSupplier; import static org.elasticsearch.ingest.ConditionalProcessor.FUNCTIONS; -public class SamplingService { +public class SamplingService implements ClusterStateListener { private static final Logger logger = LogManager.getLogger(SamplingService.class); private final ScriptService scriptService; private final LongSupplier relativeNanoTimeSupplier; private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; + private final MasterServiceTaskQueue deleteSamplingConfigTaskQueue; private final Map samples = new HashMap<>(); public SamplingService(ScriptService scriptService, ClusterService clusterService, LongSupplier relativeNanoTimeSupplier) { @@ -95,6 +99,34 @@ public Tuple executeTask( Priority.NORMAL, updateSampleConfigExecutor ); + ClusterStateTaskExecutor deleteSampleConfigExecutor = new SimpleBatchedAckListenerTaskExecutor<>() { + + @Override + public Tuple executeTask( + DeleteSampleConfigTask deleteSamplingConfigTask, + ClusterState clusterState + ) { + ProjectMetadata projectMetadata = clusterState.metadata().getProject(deleteSamplingConfigTask.projectId); + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME + ); + if (samplingConfig != null) { + ProjectMetadata.Builder projectMetadataBuilder = ProjectMetadata.builder(projectMetadata); + projectMetadataBuilder.removeCustom(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME); + ClusterState updatedClusterState = ClusterState.builder(clusterState) + .putProjectMetadata(projectMetadataBuilder) + .build(); + return new Tuple<>(updatedClusterState, deleteSamplingConfigTask); + } else { + return null; // someone beat us to it. This seems like a bad plan TODO + } + } + }; + this.deleteSamplingConfigTaskQueue = clusterService.createTaskQueue( + "delete-data-stream-mappings", + Priority.NORMAL, + deleteSampleConfigExecutor + ); } public void updateSampleConfiguration( @@ -110,22 +142,20 @@ public void updateSampleConfiguration( ActionListener listener ) { updateSamplingConfigTaskQueue.submitTask( - "updating mappings on data stream", - new UpdateSampleConfigTask( - projectId, - index, - rate, - maxSamples, - maxSize, - timeToLive, - condition, - ackTimeout, - ActionListener.runBefore(listener, () -> samples.computeIfPresent(index, (s, sampleInfo) -> new SampleInfo())) - ), + "updating sampling config", + new UpdateSampleConfigTask(projectId, index, rate, maxSamples, maxSize, timeToLive, condition, ackTimeout, listener), masterNodeTimeout ); } + public void deleteSampleConfiguration(ProjectId projectId, String index) { + deleteSamplingConfigTaskQueue.submitTask( + "deleting sampling config", + new DeleteSampleConfigTask(projectId, index, TimeValue.THIRTY_SECONDS, ActionListener.noop()), + TimeValue.THIRTY_SECONDS + ); + } + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { maybeSample( projectMetadata, @@ -149,7 +179,10 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque if (samplingConfig != null) { String samplingIndex = samplingConfig.indexName; if (samplingIndex.equals(indexRequest.index())) { - SampleInfo sampleInfo = samples.computeIfAbsent(samplingIndex, k -> new SampleInfo()); + SampleInfo sampleInfo = samples.computeIfAbsent( + samplingIndex, + k -> new SampleInfo(samplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) + ); SampleStats stats = sampleInfo.stats; stats.potentialSamples.increment(); try { @@ -201,13 +234,13 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque stats.samplesRejectedForException.increment(); stats.lastException = e; logger.info("Error performing sampling for " + samplingIndex, e); - // e.printStackTrace(System.out); } finally { stats.timeSampling.add((relativeNanoTimeSupplier.getAsLong() - startTime)); logger.info("********* Stats: " + stats); } } } + checkTTLs(projectMetadata.id()); // TODO make this happen less often? } public List getSamples(String index) { @@ -240,16 +273,56 @@ private static Script getScript(String conditional) throws IOException { } } + @Override + public void clusterChanged(ClusterChangedEvent event) { + if (event.metadataChanged()) { + for (ProjectMetadata projectMetadata : event.state().metadata().projects().values()) { + ProjectId projectId = projectMetadata.id(); + if (event.customMetadataChanged(projectId, TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME)) { + TransportPutSampleConfigAction.SamplingConfigCustomMetadata oldSamplingConfig = event.previousState() + .projectState(projectId) + .metadata() + .custom(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME); + TransportPutSampleConfigAction.SamplingConfigCustomMetadata newSamplingConfig = event.state() + .projectState(projectId) + .metadata() + .custom(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME); + if (newSamplingConfig == null && oldSamplingConfig != null) { + samples.remove(oldSamplingConfig.indexName); + } else if (newSamplingConfig != null && newSamplingConfig.equals(oldSamplingConfig) == false) { + samples.computeIfPresent( + newSamplingConfig.indexName, + (s, sampleInfo) -> new SampleInfo(newSamplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) + ); + } + } + } + } + } + + private void checkTTLs(ProjectId projectId) { + long now = relativeNanoTimeSupplier.getAsLong(); + Set indices = samples.keySet(); + for (String index : indices) { + SampleInfo sampleInfo = samples.get(index); + if (sampleInfo.expiration < now) { + deleteSampleConfiguration(projectId, index); + } + } + } + private static final class SampleInfo { private final List samples; private final SampleStats stats; + private final long expiration; private volatile Script script; private volatile IngestConditionalScript.Factory factory; private volatile boolean compilationFailed = false; - SampleInfo() { + SampleInfo(TimeValue timeToLive, long relativeNowNanos) { this.samples = new ArrayList<>(); this.stats = new SampleStats(); + this.expiration = (timeToLive == null ? TimeValue.timeValueDays(5).nanos() : timeToLive.nanos()) + relativeNowNanos; } public List getSamples() { @@ -327,4 +400,20 @@ static class UpdateSampleConfigTask extends AckedBatchedClusterStateUpdateTask { this.condition = condition; } } + + static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask { + final ProjectId projectId; + final String indexName; + + DeleteSampleConfigTask( + ProjectId projectId, + String indexName, + TimeValue ackTimeout, + ActionListener listener + ) { + super(ackTimeout, listener); + this.projectId = projectId; + this.indexName = indexName; + } + } } diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index 418762ba6ee47..08888d24a1f5c 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -716,6 +716,7 @@ private void construct( FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); SamplingService samplingService = new SamplingService(scriptService, clusterService, System::nanoTime); modules.bindToInstance(SamplingService.class, samplingService); + clusterService.addListener(samplingService); FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); final IngestService ingestService = new IngestService( diff --git a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java index f8c97955bd6d4..ab92d13abe905 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportPutSampleConfigAction.java @@ -45,6 +45,7 @@ import java.io.IOException; import java.util.EnumSet; import java.util.Iterator; +import java.util.Objects; import static org.elasticsearch.cluster.metadata.Metadata.ALL_CONTEXTS; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; @@ -123,7 +124,6 @@ public static final class SamplingConfigCustomMetadata extends AbstractNamedDiff public static final ParseField TIME_TO_LIVE_FIELD = new ParseField("time_to_live"); public static final ParseField CONDITION_FIELD = new ParseField("condition"); - @SuppressWarnings("unchecked") private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( NAME, args -> new SamplingConfigCustomMetadata( @@ -244,5 +244,24 @@ public Iterator toXContentChunked(ToXContent.Params params public static SamplingConfigCustomMetadata fromXContent(XContentParser parser) throws IOException { return PARSER.apply(parser, null); } + + @Override + public boolean equals(Object other) { + if (other instanceof SamplingConfigCustomMetadata otherConfig) { + return Objects.equals(indexName, otherConfig.indexName) + && rate == otherConfig.rate + && Objects.equals(maxSamples, otherConfig.maxSamples) + && Objects.equals(maxSize, otherConfig.maxSize) + && Objects.equals(timeToLive, otherConfig.timeToLive) + && Objects.equals(condition, otherConfig.condition); + } else { + return false; + } + } + + @Override + public int hashCode() { + return Objects.hash(indexName, rate, maxSamples, maxSize, timeToLive, condition); + } } } From 383246753c42377daf1d327a96d61a7ceaf76d6c Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Tue, 9 Sep 2025 14:37:59 +0000 Subject: [PATCH 10/37] [CI] Auto commit changes from spotless --- .../java/org/elasticsearch/ingest/SamplingService.java | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index e0b5461809aa9..f15904c4202cf 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -405,12 +405,7 @@ static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask { final ProjectId projectId; final String indexName; - DeleteSampleConfigTask( - ProjectId projectId, - String indexName, - TimeValue ackTimeout, - ActionListener listener - ) { + DeleteSampleConfigTask(ProjectId projectId, String indexName, TimeValue ackTimeout, ActionListener listener) { super(ackTimeout, listener); this.projectId = projectId; this.indexName = indexName; From f1ac65a91bf3386f1c449d5ad7f58b7874412ba4 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 11:24:51 -0500 Subject: [PATCH 11/37] adding stats --- .../elasticsearch/action/ActionModule.java | 5 + .../elasticsearch/ingest/SamplingService.java | 95 ++++++++- .../sample/GetSampleStatsAction.java | 191 ++++++++++++++++++ .../sample/RestGetSampleStatsAction.java | 44 ++++ .../sample/TransportGetSampleAction.java | 1 + .../sample/TransportGetSampleStatsAction.java | 76 +++++++ 6 files changed, 405 insertions(+), 7 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java create mode 100644 server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java diff --git a/server/src/main/java/org/elasticsearch/action/ActionModule.java b/server/src/main/java/org/elasticsearch/action/ActionModule.java index a8bdd7bd4f3d2..5b23ffb8167db 100644 --- a/server/src/main/java/org/elasticsearch/action/ActionModule.java +++ b/server/src/main/java/org/elasticsearch/action/ActionModule.java @@ -406,10 +406,13 @@ import org.elasticsearch.rest.action.synonyms.RestPutSynonymRuleAction; import org.elasticsearch.rest.action.synonyms.RestPutSynonymsAction; import org.elasticsearch.sample.GetSampleAction; +import org.elasticsearch.sample.GetSampleStatsAction; import org.elasticsearch.sample.PutSampleConfigAction; import org.elasticsearch.sample.RestGetSampleAction; +import org.elasticsearch.sample.RestGetSampleStatsAction; import org.elasticsearch.sample.RestPutSampleConfigAction; import org.elasticsearch.sample.TransportGetSampleAction; +import org.elasticsearch.sample.TransportGetSampleStatsAction; import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.snapshots.TransportUpdateSnapshotStatusAction; import org.elasticsearch.tasks.Task; @@ -821,6 +824,7 @@ public void reg actions.register(PutSampleConfigAction.INSTANCE, TransportPutSampleConfigAction.class); actions.register(GetSampleAction.INSTANCE, TransportGetSampleAction.class); + actions.register(GetSampleStatsAction.INSTANCE, TransportGetSampleStatsAction.class); return unmodifiableMap(actions.getRegistry()); } @@ -1052,6 +1056,7 @@ public void initRestHandlers(Supplier nodesInCluster, Predicate< registerHandler.accept(new RestPutSampleConfigAction()); registerHandler.accept(new RestGetSampleAction()); + registerHandler.accept(new RestGetSampleStatsAction()); } @Override diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index e0b5461809aa9..4799dcdd0a88e 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -26,6 +26,9 @@ import org.elasticsearch.common.Priority; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.bytes.ReleasableBytesReference; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; @@ -40,6 +43,7 @@ import org.elasticsearch.script.IngestConditionalScript; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptService; +import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xcontent.XContentType; @@ -247,6 +251,10 @@ public List getSamples(String index) { return samples.get(index).getSamples(); } + public SampleStats getSampleStats(String index) { + return samples.get(index).stats; + } + private boolean evaluateCondition( IngestDocument ingestDocument, Script script, @@ -335,7 +343,7 @@ void setScript(Script script, IngestConditionalScript.Factory factory) { } } - private static final class SampleStats { + public static final class SampleStats implements Writeable, ToXContent { LongAdder potentialSamples = new LongAdder(); LongAdder samplesRejectedForSize = new LongAdder(); LongAdder samplesRejectedForCondition = new LongAdder(); @@ -347,6 +355,25 @@ private static final class SampleStats { LongAdder timeCompilingCondition = new LongAdder(); Exception lastException = null; + public SampleStats() {} + + public SampleStats(StreamInput in) throws IOException { + potentialSamples.add(in.readLong()); + samplesRejectedForSize.add(in.readLong()); + samplesRejectedForCondition.add(in.readLong()); + samplesRejectedForRate.add(in.readLong()); + samplesRejectedForException.add(in.readLong()); + samples.add(in.readLong()); + timeSampling.add(in.readLong()); + timeEvaluatingCondition.add(in.readLong()); + timeCompilingCondition.add(in.readLong()); + if (in.readBoolean()) { + lastException = in.readException(); + } else { + lastException = null; + } + } + @Override public String toString() { return "potentialSamples: " @@ -368,6 +395,65 @@ public String toString() { + ", timeCompilingCondition: " + (timeCompilingCondition.longValue() / 1000000); } + + public SampleStats combine(SampleStats other) { + SampleStats result = new SampleStats(); + result.potentialSamples.add(this.potentialSamples.longValue()); + result.potentialSamples.add(other.potentialSamples.longValue()); + result.samplesRejectedForSize.add(this.samplesRejectedForSize.longValue()); + result.samplesRejectedForSize.add(other.samplesRejectedForSize.longValue()); + result.samplesRejectedForCondition.add(this.samplesRejectedForCondition.longValue()); + result.samplesRejectedForCondition.add(other.samplesRejectedForCondition.longValue()); + result.samplesRejectedForRate.add(this.samplesRejectedForRate.longValue()); + result.samplesRejectedForRate.add(other.samplesRejectedForRate.longValue()); + result.samplesRejectedForException.add(this.samplesRejectedForException.longValue()); + result.samplesRejectedForException.add(other.samplesRejectedForException.longValue()); + result.samples.add(this.samples.longValue()); + result.samples.add(other.samples.longValue()); + result.timeSampling.add(this.timeSampling.longValue()); + result.timeSampling.add(other.timeSampling.longValue()); + result.timeEvaluatingCondition.add(this.timeEvaluatingCondition.longValue()); + result.timeEvaluatingCondition.add(other.timeEvaluatingCondition.longValue()); + result.timeCompilingCondition.add(this.timeCompilingCondition.longValue()); + result.timeCompilingCondition.add(other.timeCompilingCondition.longValue()); + result.lastException = this.lastException != null ? this.lastException : other.lastException; + return result; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field("potentialSamples", potentialSamples.longValue()); + builder.field("samplesRejectedForSize", samplesRejectedForSize.longValue()); + builder.field("samplesRejectedForCondition", samplesRejectedForCondition.longValue()); + builder.field("samplesRejectedForRate", samplesRejectedForRate.longValue()); + builder.field("samplesRejectedForException", samplesRejectedForException.longValue()); + builder.field("samples", samples.longValue()); + builder.field("timeSampling", (timeSampling.longValue() / 1000000)); + builder.field("timeEvaluatingCondition", (timeEvaluatingCondition.longValue() / 1000000)); + builder.field("timeCompilingCondition", (timeCompilingCondition.longValue() / 1000000)); + builder.endObject(); + return builder; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeLong(potentialSamples.longValue()); + out.writeLong(samplesRejectedForSize.longValue()); + out.writeLong(samplesRejectedForCondition.longValue()); + out.writeLong(samplesRejectedForRate.longValue()); + out.writeLong(samplesRejectedForException.longValue()); + out.writeLong(samples.longValue()); + out.writeLong(timeSampling.longValue()); + out.writeLong(timeEvaluatingCondition.longValue()); + out.writeLong(timeCompilingCondition.longValue()); + if (lastException == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeException(lastException); + } + } } static class UpdateSampleConfigTask extends AckedBatchedClusterStateUpdateTask { @@ -405,12 +491,7 @@ static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask { final ProjectId projectId; final String indexName; - DeleteSampleConfigTask( - ProjectId projectId, - String indexName, - TimeValue ackTimeout, - ActionListener listener - ) { + DeleteSampleConfigTask(ProjectId projectId, String indexName, TimeValue ackTimeout, ActionListener listener) { super(ackTimeout, listener); this.projectId = projectId; this.indexName = indexName; diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java new file mode 100644 index 0000000000000..b713b9ef84f00 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java @@ -0,0 +1,191 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.action.ActionType; +import org.elasticsearch.action.FailedNodeException; +import org.elasticsearch.action.IndicesRequest; +import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.action.support.nodes.BaseNodeResponse; +import org.elasticsearch.action.support.nodes.BaseNodesRequest; +import org.elasticsearch.action.support.nodes.BaseNodesResponse; +import org.elasticsearch.cluster.ClusterName; +import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.ingest.SamplingService; +import org.elasticsearch.tasks.CancellableTask; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.tasks.TaskId; +import org.elasticsearch.transport.AbstractTransportRequest; +import org.elasticsearch.xcontent.ToXContentObject; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class GetSampleStatsAction extends ActionType { + + public static final GetSampleStatsAction INSTANCE = new GetSampleStatsAction(); + public static final String NAME = "indices:admin/sample/stats"; + + private GetSampleStatsAction() { + super(NAME); + } + + public static class Response extends BaseNodesResponse implements Writeable, ToXContentObject { + + public Response(StreamInput in) throws IOException { + super(in); + } + + public Response(ClusterName clusterName, List nodes, List failures) { + super(clusterName, nodes, failures); + } + + public SamplingService.SampleStats getSampleStats() { + return getNodes().stream() + .map(n -> n.sampleStats) + .filter(Objects::nonNull) + .reduce(SamplingService.SampleStats::combine) + .orElse(new SamplingService.SampleStats()); + } + + @Override + protected List readNodesFrom(StreamInput in) throws IOException { + return in.readCollectionAsList(GetSampleStatsAction.NodeResponse::new); + } + + @Override + protected void writeNodesTo(StreamOutput out, List nodes) throws IOException { + out.writeCollection(nodes); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + GetSampleStatsAction.Response that = (GetSampleStatsAction.Response) o; + return Objects.equals(getNodes(), that.getNodes()) && Objects.equals(failures(), that.failures()); + } + + @Override + public int hashCode() { + return Objects.hash(getNodes(), failures()); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return getSampleStats().toXContent(builder, params); + } + } + + public static class NodeResponse extends BaseNodeResponse { + private final SamplingService.SampleStats sampleStats; + + protected NodeResponse(StreamInput in) throws IOException { + super(in); + sampleStats = new SamplingService.SampleStats(in); + } + + protected NodeResponse(DiscoveryNode node, SamplingService.SampleStats sampleStats) { + super(node); + this.sampleStats = sampleStats; + } + + public SamplingService.SampleStats getSampleStats() { + return sampleStats; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + sampleStats.writeTo(out); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + GetSampleStatsAction.NodeResponse that = (GetSampleStatsAction.NodeResponse) o; + return sampleStats.equals(that.sampleStats); // TODO + } + + @Override + public int hashCode() { + return Objects.hash(sampleStats); // TODO + } + } + + public static class Request extends BaseNodesRequest implements IndicesRequest.Replaceable { + private String[] names; + + public Request(String[] names) { + super((String[]) null); + this.names = names; + } + + @Override + public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { + return new CancellableTask(id, type, action, "", parentTaskId, headers); + } + + @Override + public ActionRequestValidationException validate() { + if (this.indices().length != 1) { + return new ActionRequestValidationException(); + } + return null; + } + + @Override + public IndicesRequest indices(String... indices) { + this.names = indices; + return this; + } + + @Override + public String[] indices() { + return names; + } + + @Override + public IndicesOptions indicesOptions() { + return IndicesOptions.DEFAULT; + } + } + + public static class NodeRequest extends AbstractTransportRequest { + private final String index; + + public NodeRequest(String index) { + this.index = index; + } + + public NodeRequest(StreamInput in) throws IOException { + super(in); + this.index = in.readString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(index); + } + + public String getIndex() { + return index; + } + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java new file mode 100644 index 0000000000000..dda05bde244a4 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java @@ -0,0 +1,44 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; + +import static org.elasticsearch.rest.RestRequest.Method.GET; + +public class RestGetSampleStatsAction extends BaseRestHandler { + @Override + public String getName() { + return "get_sample_stats"; + } + + @Override + public List routes() { + return List.of(new Route(GET, "/_sample/{name}/stats")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + GetSampleStatsAction.Request getSampleRequest = new GetSampleStatsAction.Request(new String[] { request.param("name") }); + // return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( + // GetSampleStatsAction.INSTANCE, + // getSampleRequest, + // (ActionListener) new RestToXContentListener<>(channel) + // ); + + return channel -> client.execute(GetSampleStatsAction.INSTANCE, getSampleRequest, new RestToXContentListener<>(channel)); + } +} diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java index cb027daa3b697..2f0d90ed7a109 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java @@ -55,6 +55,7 @@ public TransportGetSampleAction( @SuppressWarnings("checkstyle:LineLength") @Override protected Response newResponse(Request request, List nodeResponses, List failures) { + samplingService.getSampleConfig(request.projectId, request.indices()[0]); return new Response(clusterService.getClusterName(), nodeResponses, failures); } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java new file mode 100644 index 0000000000000..263c283d8e33b --- /dev/null +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java @@ -0,0 +1,76 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.sample; + +import org.elasticsearch.action.FailedNodeException; +import org.elasticsearch.action.support.ActionFilters; +import org.elasticsearch.action.support.nodes.TransportNodesAction; +import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.ingest.SamplingService; +import org.elasticsearch.injection.guice.Inject; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.TransportService; + +import java.io.IOException; +import java.util.List; + +import static org.elasticsearch.sample.GetSampleStatsAction.NodeRequest; +import static org.elasticsearch.sample.GetSampleStatsAction.NodeResponse; +import static org.elasticsearch.sample.GetSampleStatsAction.Request; +import static org.elasticsearch.sample.GetSampleStatsAction.Response; + +public class TransportGetSampleStatsAction extends TransportNodesAction { + private final SamplingService samplingService; + + @Inject + public TransportGetSampleStatsAction( + TransportService transportService, + ClusterService clusterService, + ThreadPool threadPool, + ActionFilters actionFilters, + SamplingService samplingService + ) { + super( + GetSampleStatsAction.NAME, + clusterService, + transportService, + actionFilters, + NodeRequest::new, + threadPool.executor(ThreadPool.Names.MANAGEMENT) + ); + this.samplingService = samplingService; + } + + @SuppressWarnings("checkstyle:LineLength") + @Override + protected Response newResponse(Request request, List nodeResponses, List failures) { + return new Response(clusterService.getClusterName(), nodeResponses, failures); + } + + @Override + protected NodeRequest newNodeRequest(Request request) { + return new NodeRequest(request.indices()[0]); + } + + @Override + protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throws IOException { + return new NodeResponse(in); + } + + @Override + protected NodeResponse nodeOperation(NodeRequest request, Task task) { + String index = request.getIndex(); + SamplingService.SampleStats sampleStats = samplingService.getSampleStats(index); + return new NodeResponse(transportService.getLocalNode(), sampleStats); + } +} From ca84960947ef091689da7735e676e4053f220634 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 14:49:56 -0500 Subject: [PATCH 12/37] incorporating projetId, limiting results --- .../elasticsearch/action/ActionModule.java | 4 +- .../elasticsearch/ingest/SamplingService.java | 45 +++++++++++------ .../elasticsearch/sample/GetSampleAction.java | 26 ++++++++-- .../sample/GetSampleStatsAction.java | 49 +++++++++++++++++-- .../sample/RestGetSampleAction.java | 12 ++++- .../sample/RestGetSampleStatsAction.java | 18 ++++--- .../sample/TransportGetSampleAction.java | 14 ++++-- .../sample/TransportGetSampleStatsAction.java | 12 +++-- 8 files changed, 139 insertions(+), 41 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/ActionModule.java b/server/src/main/java/org/elasticsearch/action/ActionModule.java index 5b23ffb8167db..23e7195f3ac1f 100644 --- a/server/src/main/java/org/elasticsearch/action/ActionModule.java +++ b/server/src/main/java/org/elasticsearch/action/ActionModule.java @@ -1055,8 +1055,8 @@ public void initRestHandlers(Supplier nodesInCluster, Predicate< registerHandler.accept(new RestDeleteSynonymRuleAction()); registerHandler.accept(new RestPutSampleConfigAction()); - registerHandler.accept(new RestGetSampleAction()); - registerHandler.accept(new RestGetSampleStatsAction()); + registerHandler.accept(new RestGetSampleAction(projectIdResolver)); + registerHandler.accept(new RestGetSampleStatsAction(projectIdResolver)); } @Override diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 4799dcdd0a88e..9f28a0fa31ef1 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -66,7 +66,7 @@ public class SamplingService implements ClusterStateListener { private final LongSupplier relativeNanoTimeSupplier; private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; private final MasterServiceTaskQueue deleteSamplingConfigTaskQueue; - private final Map samples = new HashMap<>(); + private final Map samples = new HashMap<>(); public SamplingService(ScriptService scriptService, ClusterService clusterService, LongSupplier relativeNanoTimeSupplier) { this.scriptService = scriptService; @@ -180,11 +180,12 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME ); + ProjectId projectId = projectMetadata.id(); if (samplingConfig != null) { String samplingIndex = samplingConfig.indexName; if (samplingIndex.equals(indexRequest.index())) { SampleInfo sampleInfo = samples.computeIfAbsent( - samplingIndex, + new ProjectIndex(projectId, samplingIndex), k -> new SampleInfo(samplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) ); SampleStats stats = sampleInfo.stats; @@ -244,15 +245,25 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque } } } - checkTTLs(projectMetadata.id()); // TODO make this happen less often? + checkTTLs(); // TODO make this happen less often? } - public List getSamples(String index) { - return samples.get(index).getSamples(); + public List getSamples(ProjectId projectId, String index) { + return samples.get(new ProjectIndex(projectId, index)).getSamples(); } - public SampleStats getSampleStats(String index) { - return samples.get(index).stats; + public SampleStats getSampleStats(ProjectId projectId, String index) { + return samples.get(new ProjectIndex(projectId, index)).stats; + } + + public TransportPutSampleConfigAction.SamplingConfigCustomMetadata getSampleConfig(ProjectMetadata projectMetadata, String index) { + TransportPutSampleConfigAction.SamplingConfigCustomMetadata sampleConfig = projectMetadata.custom( + TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME + ); + if (sampleConfig != null && sampleConfig.indexName.equals(index)) { + return sampleConfig; + } + return null; } private boolean evaluateCondition( @@ -296,10 +307,10 @@ public void clusterChanged(ClusterChangedEvent event) { .metadata() .custom(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME); if (newSamplingConfig == null && oldSamplingConfig != null) { - samples.remove(oldSamplingConfig.indexName); + samples.remove(new ProjectIndex(projectId, oldSamplingConfig.indexName)); } else if (newSamplingConfig != null && newSamplingConfig.equals(oldSamplingConfig) == false) { samples.computeIfPresent( - newSamplingConfig.indexName, + new ProjectIndex(projectId, newSamplingConfig.indexName), (s, sampleInfo) -> new SampleInfo(newSamplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) ); } @@ -308,13 +319,13 @@ public void clusterChanged(ClusterChangedEvent event) { } } - private void checkTTLs(ProjectId projectId) { + private void checkTTLs() { long now = relativeNanoTimeSupplier.getAsLong(); - Set indices = samples.keySet(); - for (String index : indices) { - SampleInfo sampleInfo = samples.get(index); + Set projectIndices = samples.keySet(); + for (ProjectIndex projectIndex : projectIndices) { + SampleInfo sampleInfo = samples.get(projectIndex); if (sampleInfo.expiration < now) { - deleteSampleConfiguration(projectId, index); + deleteSampleConfiguration(projectIndex.projectId, projectIndex.indexName); } } } @@ -345,11 +356,11 @@ void setScript(Script script, IngestConditionalScript.Factory factory) { public static final class SampleStats implements Writeable, ToXContent { LongAdder potentialSamples = new LongAdder(); - LongAdder samplesRejectedForSize = new LongAdder(); + public LongAdder samplesRejectedForSize = new LongAdder(); LongAdder samplesRejectedForCondition = new LongAdder(); LongAdder samplesRejectedForRate = new LongAdder(); LongAdder samplesRejectedForException = new LongAdder(); - LongAdder samples = new LongAdder(); + public LongAdder samples = new LongAdder(); LongAdder timeSampling = new LongAdder(); LongAdder timeEvaluatingCondition = new LongAdder(); LongAdder timeCompilingCondition = new LongAdder(); @@ -497,4 +508,6 @@ static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask { this.indexName = indexName; } } + + record ProjectIndex(ProjectId projectId, String indexName) {}; } diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java index 8aebf8ee7b206..2a39693a3e7ab 100644 --- a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java @@ -19,6 +19,7 @@ import org.elasticsearch.action.support.nodes.BaseNodesRequest; import org.elasticsearch.action.support.nodes.BaseNodesResponse; import org.elasticsearch.cluster.ClusterName; +import org.elasticsearch.cluster.metadata.ProjectId; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.collect.Iterators; import org.elasticsearch.common.io.stream.StreamInput; @@ -51,17 +52,20 @@ private GetSampleAction() { } public static class Response extends BaseNodesResponse implements Writeable, ChunkedToXContent { + private final int maxSize; public Response(StreamInput in) throws IOException { super(in); + maxSize = in.readInt(); } - public Response(ClusterName clusterName, List nodes, List failures) { + public Response(ClusterName clusterName, List nodes, List failures, int maxSize) { super(clusterName, nodes, failures); + this.maxSize = maxSize; } public List getSamples() { - return getNodes().stream().map(n -> n.samples).filter(Objects::nonNull).flatMap(Collection::stream).toList(); + return getNodes().stream().map(n -> n.samples).filter(Objects::nonNull).flatMap(Collection::stream).limit(maxSize).toList(); } @Override @@ -140,10 +144,12 @@ public int hashCode() { } public static class Request extends BaseNodesRequest implements IndicesRequest.Replaceable { + private final ProjectId projectId; private String[] names; - public Request(String[] names) { + public Request(ProjectId projectId, String[] names) { super((String[]) null); + this.projectId = projectId; this.names = names; } @@ -160,6 +166,10 @@ public ActionRequestValidationException validate() { return null; } + public ProjectId getProjectId() { + return projectId; + } + @Override public IndicesRequest indices(String... indices) { this.names = indices; @@ -178,23 +188,31 @@ public IndicesOptions indicesOptions() { } public static class NodeRequest extends AbstractTransportRequest { + private final ProjectId projectId; private final String index; - public NodeRequest(String index) { + public NodeRequest(ProjectId projectId, String index) { + this.projectId = projectId; this.index = index; } public NodeRequest(StreamInput in) throws IOException { super(in); + this.projectId = ProjectId.readFrom(in); this.index = in.readString(); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); + projectId.writeTo(out); out.writeString(index); } + public ProjectId getProjectId() { + return projectId; + } + public String getIndex() { return index; } diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java index b713b9ef84f00..caf667bad247d 100644 --- a/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java @@ -18,6 +18,7 @@ import org.elasticsearch.action.support.nodes.BaseNodesRequest; import org.elasticsearch.action.support.nodes.BaseNodesResponse; import org.elasticsearch.cluster.ClusterName; +import org.elasticsearch.cluster.metadata.ProjectId; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; @@ -45,16 +46,44 @@ private GetSampleStatsAction() { } public static class Response extends BaseNodesResponse implements Writeable, ToXContentObject { + private final int maxSize; public Response(StreamInput in) throws IOException { super(in); + maxSize = in.readInt(); } - public Response(ClusterName clusterName, List nodes, List failures) { + public Response( + ClusterName clusterName, + List nodes, + List failures, + int maxSize + ) { super(clusterName, nodes, failures); + this.maxSize = maxSize; } public SamplingService.SampleStats getSampleStats() { + SamplingService.SampleStats rawStats = getRawSampleStats(); + if (rawStats.samples.longValue() > maxSize) { + SamplingService.SampleStats filteredStats = new SamplingService.SampleStats().combine(rawStats); + System.out.println( + "**** samples: " + + filteredStats.samples.longValue() + + ", maxSize: " + + maxSize + + ", rawStats.samples: " + + rawStats.samples.longValue() + ); + filteredStats.samples.add(maxSize - rawStats.samples.longValue()); + filteredStats.samplesRejectedForSize.add(rawStats.samples.longValue() - maxSize); + return filteredStats; + } else { + return rawStats; + } + } + + private SamplingService.SampleStats getRawSampleStats() { return getNodes().stream() .map(n -> n.sampleStats) .filter(Objects::nonNull) @@ -129,10 +158,12 @@ public int hashCode() { } public static class Request extends BaseNodesRequest implements IndicesRequest.Replaceable { + private final ProjectId projectId; private String[] names; - public Request(String[] names) { + public Request(ProjectId projectId, String[] names) { super((String[]) null); + this.projectId = projectId; this.names = names; } @@ -149,6 +180,10 @@ public ActionRequestValidationException validate() { return null; } + public ProjectId getProjectId() { + return projectId; + } + @Override public IndicesRequest indices(String... indices) { this.names = indices; @@ -167,23 +202,31 @@ public IndicesOptions indicesOptions() { } public static class NodeRequest extends AbstractTransportRequest { + private final ProjectId projectId; private final String index; - public NodeRequest(String index) { + public NodeRequest(ProjectId projectId, String index) { + this.projectId = projectId; this.index = index; } public NodeRequest(StreamInput in) throws IOException { super(in); + this.projectId = ProjectId.readFrom(in); this.index = in.readString(); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); + projectId.writeTo(out); out.writeString(index); } + public ProjectId getProjectId() { + return projectId; + } + public String getIndex() { return index; } diff --git a/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java index abe16175c8765..ecfd783874f8e 100644 --- a/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/RestGetSampleAction.java @@ -10,6 +10,7 @@ package org.elasticsearch.sample; import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.cluster.project.ProjectIdResolver; import org.elasticsearch.rest.BaseRestHandler; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.action.RestCancellableNodeClient; @@ -21,6 +22,12 @@ import static org.elasticsearch.rest.RestRequest.Method.GET; public class RestGetSampleAction extends BaseRestHandler { + private final ProjectIdResolver projectIdResolver; + + public RestGetSampleAction(ProjectIdResolver projectIdResolver) { + this.projectIdResolver = projectIdResolver; + } + @Override public String getName() { return "get_sample"; @@ -33,7 +40,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - GetSampleAction.Request getSampleRequest = new GetSampleAction.Request(new String[] { request.param("name") }); + GetSampleAction.Request getSampleRequest = new GetSampleAction.Request( + projectIdResolver.getProjectId(), + new String[] { request.param("name") } + ); return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( GetSampleAction.INSTANCE, getSampleRequest, diff --git a/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java index dda05bde244a4..e432622340618 100644 --- a/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java +++ b/server/src/main/java/org/elasticsearch/sample/RestGetSampleStatsAction.java @@ -10,6 +10,7 @@ package org.elasticsearch.sample; import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.cluster.project.ProjectIdResolver; import org.elasticsearch.rest.BaseRestHandler; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.action.RestToXContentListener; @@ -20,6 +21,12 @@ import static org.elasticsearch.rest.RestRequest.Method.GET; public class RestGetSampleStatsAction extends BaseRestHandler { + private final ProjectIdResolver projectIdResolver; + + public RestGetSampleStatsAction(ProjectIdResolver projectIdResolver) { + this.projectIdResolver = projectIdResolver; + } + @Override public String getName() { return "get_sample_stats"; @@ -32,13 +39,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - GetSampleStatsAction.Request getSampleRequest = new GetSampleStatsAction.Request(new String[] { request.param("name") }); - // return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( - // GetSampleStatsAction.INSTANCE, - // getSampleRequest, - // (ActionListener) new RestToXContentListener<>(channel) - // ); - + GetSampleStatsAction.Request getSampleRequest = new GetSampleStatsAction.Request( + projectIdResolver.getProjectId(), + new String[] { request.param("name") } + ); return channel -> client.execute(GetSampleStatsAction.INSTANCE, getSampleRequest, new RestToXContentListener<>(channel)); } } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java index 2f0d90ed7a109..033e7b3c1a58a 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java @@ -13,6 +13,7 @@ import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.ActionFilters; import org.elasticsearch.action.support.nodes.TransportNodesAction; +import org.elasticsearch.cluster.metadata.ProjectId; import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.cluster.service.ClusterService; import org.elasticsearch.common.io.stream.StreamInput; @@ -55,13 +56,17 @@ public TransportGetSampleAction( @SuppressWarnings("checkstyle:LineLength") @Override protected Response newResponse(Request request, List nodeResponses, List failures) { - samplingService.getSampleConfig(request.projectId, request.indices()[0]); - return new Response(clusterService.getClusterName(), nodeResponses, failures); + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = samplingService.getSampleConfig( + clusterService.state().projectState(request.getProjectId()).metadata(), + request.indices()[0] + ); + int maxSamples = samplingConfig == null ? 0 : samplingConfig.maxSamples; + return new Response(clusterService.getClusterName(), nodeResponses, failures, maxSamples); } @Override protected NodeRequest newNodeRequest(Request request) { - return new NodeRequest(request.indices()[0]); + return new NodeRequest(request.getProjectId(), request.indices()[0]); } @Override @@ -71,8 +76,9 @@ protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throw @Override protected NodeResponse nodeOperation(NodeRequest request, Task task) { + ProjectId projectId = request.getProjectId(); String index = request.getIndex(); - List samples = samplingService.getSamples(index); + List samples = samplingService.getSamples(projectId, index); return new NodeResponse(transportService.getLocalNode(), samples == null ? List.of() : samples); } } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java index 263c283d8e33b..1c3c2abc7445f 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleStatsAction.java @@ -54,12 +54,17 @@ public TransportGetSampleStatsAction( @SuppressWarnings("checkstyle:LineLength") @Override protected Response newResponse(Request request, List nodeResponses, List failures) { - return new Response(clusterService.getClusterName(), nodeResponses, failures); + TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = samplingService.getSampleConfig( + clusterService.state().projectState(request.getProjectId()).metadata(), + request.indices()[0] + ); + int maxSamples = samplingConfig == null ? 0 : samplingConfig.maxSamples; + return new Response(clusterService.getClusterName(), nodeResponses, failures, maxSamples); } @Override protected NodeRequest newNodeRequest(Request request) { - return new NodeRequest(request.indices()[0]); + return new NodeRequest(request.getProjectId(), request.indices()[0]); } @Override @@ -69,8 +74,7 @@ protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throw @Override protected NodeResponse nodeOperation(NodeRequest request, Task task) { - String index = request.getIndex(); - SamplingService.SampleStats sampleStats = samplingService.getSampleStats(index); + SamplingService.SampleStats sampleStats = samplingService.getSampleStats(request.getProjectId(), request.getIndex()); return new NodeResponse(transportService.getLocalNode(), sampleStats); } } From f906bf5273099332cf99f2459813ff615d7abf7e Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 14:52:10 -0500 Subject: [PATCH 13/37] fixing merge --- .../main/java/org/elasticsearch/TransportVersions.java | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index 5517a2075425f..1f57b8475ca5d 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -341,16 +341,6 @@ static TransportVersion def(int id) { public static final TransportVersion PROJECT_STATE_REGISTRY_ENTRY = def(9_124_0_00); public static final TransportVersion ML_INFERENCE_LLAMA_ADDED = def(9_125_0_00); public static final TransportVersion SHARD_WRITE_LOAD_IN_CLUSTER_INFO = def(9_126_0_00); -<<<<<<< HEAD - public static final TransportVersion ESQL_SAMPLE_OPERATOR_STATUS = def(9_127_0_00); - public static final TransportVersion PROJECT_RESERVED_STATE_MOVE_TO_REGISTRY = def(9_147_0_00); - public static final TransportVersion STREAMS_ENDPOINT_PARAM_RESTRICTIONS = def(9_148_0_00); - public static final TransportVersion RESOLVE_INDEX_MODE_FILTER = def(9_149_0_00); - public static final TransportVersion SEMANTIC_QUERY_MULTIPLE_INFERENCE_IDS = def(9_150_0_00); - public static final TransportVersion ESQL_LOOKUP_JOIN_PRE_JOIN_FILTER = def(9_151_0_00); - public static final TransportVersion INFERENCE_API_DISABLE_EIS_RATE_LIMITING = def(9_152_0_00); - public static final TransportVersion GEMINI_THINKING_BUDGET_ADDED = def(9_153_0_00); - public static final TransportVersion VISIT_PERCENTAGE = def(9_154_0_00); public static final TransportVersion TIME_SERIES_TELEMETRY = def(9_155_0_00); public static final TransportVersion INFERENCE_API_EIS_DIAGNOSTICS = def(9_156_0_00); public static final TransportVersion RANDOM_SAMPLING = def(9_157_0_00); From 2961cc5b71b549648d77331d086c55bf33e47960 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 15:04:29 -0500 Subject: [PATCH 14/37] removing some dead code --- .../bulk/TransportAbstractBulkAction.java | 30 ++----------------- .../elasticsearch/ingest/IngestService.java | 25 +++++----------- .../sample/RestPutSampleConfigAction.java | 6 ++-- 3 files changed, 14 insertions(+), 47 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java index 6c1ca0c5d6b11..d01b18813f5fb 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java +++ b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java @@ -51,7 +51,6 @@ import java.io.IOException; import java.util.HashMap; import java.util.HashSet; -import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; @@ -65,7 +64,6 @@ */ public abstract class TransportAbstractBulkAction extends HandledTransportAction { private static final Logger logger = LogManager.getLogger(TransportAbstractBulkAction.class); - private final Map> samples = new HashMap<>(); public static final Set STREAMS_ALLOWED_PARAMS = new HashSet<>(9) { { add("error_trace"); @@ -314,7 +312,7 @@ private boolean applyPipelines( }); return true; } else if (samplingService != null && firstTime) { - // else sample, but only if this is the first time through? + // else sample, but only if this is the first time through. Otherwise we had pipelines and sampled in IngestService for (DocWriteRequest actionRequest : bulkRequest.requests) { if (actionRequest instanceof IndexRequest ir) { samplingService.maybeSample(project, ir); @@ -334,10 +332,6 @@ private void processBulkIndexIngestRequest( final long ingestStartTimeInNanos = relativeTimeNanos(); final BulkRequestModifier bulkRequestModifier = new BulkRequestModifier(original); final Thread originalThread = Thread.currentThread(); - // Function rawDocSaver = indexRequest -> { - // indexRequest.source(); - // return null; - // }; getIngestService(original).executeBulkRequest( metadata.id(), original.numberOfActions(), @@ -455,22 +449,8 @@ private void applyPipelinesAndDoInternalExecute( } var wrappedListener = bulkRequestModifier.wrapActionListenerIfNeeded(listener); - boolean noPipelinesRemaining; - try { - noPipelinesRemaining = applyPipelines( - task, - bulkRequestModifier.getBulkRequest(), - executor, - wrappedListener, - firstTime - ) == false; - } catch (Exception e) { - maybeSample(); - throw e; - } - if (noPipelinesRemaining) { - // we have run all pipelines, so maybe sample - maybeSample(); + + if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, firstTime) == false) { doInternalExecute(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, relativeStartTimeNanos); } } @@ -527,10 +507,6 @@ private boolean streamsRestrictedParamsUsed(BulkRequest bulkRequest) { return Sets.difference(bulkRequest.requestParamsUsed(), STREAMS_ALLOWED_PARAMS).isEmpty() == false; } - private void maybeSample() { - // Is sampling enabled for this index? - } - /** * This method creates any missing resources and actually applies the BulkRequest to the relevant indices * @param task The task in which this work is being done diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 07458a97d60bc..406f35a6b97da 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -1362,24 +1362,15 @@ private void executePipelines( if (newPipelines.hasNext()) { executePipelines(newPipelines, indexRequest, ingestDocument, resolveFailureStore, listener, originalDocumentMetadata); } else { - // update the index request's source and (potentially) cache the timestamp for TSDB - // Here is where we finally overwrite source. Somehow sample before this? - // But also somewhere else if there are no pipelines - // Can't do this in the listener though b/c the source is already gone - // byte[] originalBytes; - // BytesReference sourceBytesRef = indexRequest.source(); - // if (sourceBytesRef.hasArray()) { - // originalBytes = Arrays.copyOfRange(sourceBytesRef.array(), sourceBytesRef.arrayOffset(), sourceBytesRef.arrayOffset() - // + sourceBytesRef.length()); - // } else { - // originalBytes = sourceBytesRef.toBytesRef().bytes; - // } - // - // if sampling enabled for this index AND condition is met AND sample THEN try { - IndexRequest sample = copyIndexRequest(indexRequest); - updateIndexRequestMetadata(sample, originalDocumentMetadata); - samplingService.maybeSample(project, sample, ingestDocument); + /* + * 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. + */ + IndexRequest original = copyIndexRequest(indexRequest); + updateIndexRequestMetadata(original, originalDocumentMetadata); + samplingService.maybeSample(project, original, ingestDocument); } catch (IOException ex) { logger.warn("unable to sample data"); } diff --git a/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java b/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java index 6592ec3643930..2adb036362304 100644 --- a/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java +++ b/server/src/main/java/org/elasticsearch/sample/RestPutSampleConfigAction.java @@ -60,13 +60,13 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).execute( PutSampleConfigAction.INSTANCE, putSampleConfigRequest, - new ReindexDataStreamRestToXContentListener(channel) + new RestToXContentListener(channel) ); } - static class ReindexDataStreamRestToXContentListener extends RestBuilderListener { + static class RestToXContentListener extends RestBuilderListener { - ReindexDataStreamRestToXContentListener(RestChannel channel) { + RestToXContentListener(RestChannel channel) { super(channel); } From 7ab11884461fc80423629989921db4a748bced55 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 15:32:05 -0500 Subject: [PATCH 15/37] fixing ingestion --- .../src/main/java/org/elasticsearch/ingest/IngestService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 406f35a6b97da..b3c64eaae350a 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -1369,7 +1369,7 @@ private void executePipelines( * before overwriting it here. We can discard it after sampling. */ IndexRequest original = copyIndexRequest(indexRequest); - updateIndexRequestMetadata(original, originalDocumentMetadata); + updateIndexRequestMetadata(indexRequest, originalDocumentMetadata); samplingService.maybeSample(project, original, ingestDocument); } catch (IOException ex) { logger.warn("unable to sample data"); From 47da33603ba0d9261755f4f8276edf2f381ba9e0 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 9 Sep 2025 15:34:21 -0500 Subject: [PATCH 16/37] fixing build --- .../org/elasticsearch/sample/GetSampleStatsAction.java | 8 -------- .../action/TransportGetTrainedModelsStatsActionTests.java | 3 ++- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java index caf667bad247d..151dc0c8d6e6c 100644 --- a/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleStatsAction.java @@ -67,14 +67,6 @@ public SamplingService.SampleStats getSampleStats() { SamplingService.SampleStats rawStats = getRawSampleStats(); if (rawStats.samples.longValue() > maxSize) { SamplingService.SampleStats filteredStats = new SamplingService.SampleStats().combine(rawStats); - System.out.println( - "**** samples: " - + filteredStats.samples.longValue() - + ", maxSize: " - + maxSize - + ", rawStats.samples: " - + rawStats.samples.longValue() - ); filteredStats.samples.add(maxSize - rawStats.samples.longValue()); filteredStats.samplesRejectedForSize.add(rawStats.samples.longValue() - maxSize); return filteredStats; diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java index 06ba7ba113d4e..fc39db5115338 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java @@ -159,7 +159,8 @@ public void setUpVariables() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + null ); } From 917b9c6b7c8c09b5a0e49c2921805693d077246b Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 10 Sep 2025 11:06:53 -0500 Subject: [PATCH 17/37] Fixing IndexRequest cloning --- .../elasticsearch/ingest/IngestService.java | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index b3c64eaae350a..f182266f07fe8 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -16,7 +16,6 @@ import org.elasticsearch.ElasticsearchParseException; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.ResourceNotFoundException; -import org.elasticsearch.TransportVersion; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.DocWriteRequest; import org.elasticsearch.action.admin.cluster.node.info.NodeInfo; @@ -53,9 +52,8 @@ 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.io.stream.BytesStreamOutput; -import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.common.regex.Regex; @@ -1369,7 +1367,7 @@ private void executePipelines( * before overwriting it here. We can discard it after sampling. */ IndexRequest original = copyIndexRequest(indexRequest); - updateIndexRequestMetadata(indexRequest, originalDocumentMetadata); + updateIndexRequestMetadata(original, originalDocumentMetadata); samplingService.maybeSample(project, original, ingestDocument); } catch (IOException ex) { logger.warn("unable to sample data"); @@ -1392,16 +1390,24 @@ private void executePipelines( } private IndexRequest copyIndexRequest(IndexRequest original) throws IOException { - // This makes a whole new copy of the source, and removes it from the BytesReference. I think this is what we want here? That way - // if we have a huge bulk request and only one thing is sampled, we don't keep it all. - try (BytesStreamOutput output = new BytesStreamOutput()) { - output.setTransportVersion(TransportVersion.current()); - original.writeTo(output); - try (StreamInput in = output.bytes().streamInput()) { - in.setTransportVersion(TransportVersion.current()); - return new IndexRequest(in); - } - } + 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.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()); + BytesReference source = original.source(); + if (source != null) { + clonedRequest.source(source, original.getContentType()); + } + return clonedRequest; } private static void executePipeline( From 7d0db7221fc4d108b1b6262ca3d650241d8403cd Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 10 Sep 2025 11:10:29 -0500 Subject: [PATCH 18/37] fixing OperatorPrivilegesIT --- .../org/elasticsearch/xpack/security/operator/Constants.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java b/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java index da9a81898de60..4339fcd38ecaf 100644 --- a/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java +++ b/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java @@ -655,6 +655,9 @@ public class Constants { "indices:admin/index/create_from_source", "indices:admin/index/copy_lifecycle_index_metadata", "internal:admin/repository/verify", - "internal:admin/repository/verify/coordinate" + "internal:admin/repository/verify/coordinate", + "indices:admin/sample/stats", + "indices:admin/sample/config/update", + "indices:admin/sample" ).filter(Objects::nonNull).collect(Collectors.toUnmodifiableSet()); } From 0ef634ce6f28cd1f948bd274d29466d3565c2821 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 10 Sep 2025 11:48:57 -0500 Subject: [PATCH 19/37] using SoftReferences --- .../elasticsearch/ingest/SamplingService.java | 155 ++++++++++-------- 1 file changed, 83 insertions(+), 72 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 9f28a0fa31ef1..82eeb8b243a62 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -50,6 +50,7 @@ import org.elasticsearch.xcontent.json.JsonXContent; import java.io.IOException; +import java.lang.ref.SoftReference; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -66,7 +67,7 @@ public class SamplingService implements ClusterStateListener { private final LongSupplier relativeNanoTimeSupplier; private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; private final MasterServiceTaskQueue deleteSamplingConfigTaskQueue; - private final Map samples = new HashMap<>(); + private final Map> samples = new HashMap<>(); public SamplingService(ScriptService scriptService, ClusterService clusterService, LongSupplier relativeNanoTimeSupplier) { this.scriptService = scriptService; @@ -184,64 +185,69 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque if (samplingConfig != null) { String samplingIndex = samplingConfig.indexName; if (samplingIndex.equals(indexRequest.index())) { - SampleInfo sampleInfo = samples.computeIfAbsent( + SoftReference sampleInfoReference = samples.compute( new ProjectIndex(projectId, samplingIndex), - k -> new SampleInfo(samplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) + (k, v) -> v == null || v.get() == null + ? new SoftReference<>(new SampleInfo(samplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong())) + : v ); - SampleStats stats = sampleInfo.stats; - stats.potentialSamples.increment(); - try { - if (sampleInfo.getSamples().size() < samplingConfig.maxSamples) { - String condition = samplingConfig.condition; - if (condition != null) { - if (sampleInfo.script == null || sampleInfo.factory == null) { - // We don't want to pay for synchronization because worst case, we compile the script twice - long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); - try { - if (sampleInfo.compilationFailed) { - // we don't want to waste time - stats.samplesRejectedForException.increment(); - return; - } else { - Script script = getScript(condition); - sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); + SampleInfo sampleInfo = sampleInfoReference.get(); + if (sampleInfo != null) { + SampleStats stats = sampleInfo.stats; + stats.potentialSamples.increment(); + try { + if (sampleInfo.getSamples().size() < samplingConfig.maxSamples) { + String condition = samplingConfig.condition; + if (condition != null) { + if (sampleInfo.script == null || sampleInfo.factory == null) { + // We don't want to pay for synchronization because worst case, we compile the script twice + long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); + try { + if (sampleInfo.compilationFailed) { + // we don't want to waste time + stats.samplesRejectedForException.increment(); + return; + } else { + Script script = getScript(condition); + sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); + } + } catch (Exception e) { + sampleInfo.compilationFailed = true; + throw e; + } finally { + stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); } - } catch (Exception e) { - sampleInfo.compilationFailed = true; - throw e; - } finally { - stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); } } - } - long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); - if (condition == null - || evaluateCondition(ingestDocument, sampleInfo.script, sampleInfo.factory, sampleInfo.stats)) { - stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); - if (Math.random() < samplingConfig.rate) { - indexRequest.incRef(); - if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) { - releaseableSource.incRef(); + long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); + if (condition == null + || evaluateCondition(ingestDocument, sampleInfo.script, sampleInfo.factory, sampleInfo.stats)) { + stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); + if (Math.random() < samplingConfig.rate) { + indexRequest.incRef(); + if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) { + releaseableSource.incRef(); + } + sampleInfo.getSamples().add(indexRequest); + stats.samples.increment(); + logger.info("Sampling " + indexRequest); + } else { + stats.samplesRejectedForRate.increment(); } - sampleInfo.getSamples().add(indexRequest); - stats.samples.increment(); - logger.info("Sampling " + indexRequest); } else { - stats.samplesRejectedForRate.increment(); + stats.samplesRejectedForCondition.increment(); } } else { - stats.samplesRejectedForCondition.increment(); + stats.samplesRejectedForSize.increment(); } - } else { - stats.samplesRejectedForSize.increment(); + } catch (Exception e) { + stats.samplesRejectedForException.increment(); + stats.lastException = e; + logger.info("Error performing sampling for " + samplingIndex, e); + } finally { + stats.timeSampling.add((relativeNanoTimeSupplier.getAsLong() - startTime)); + logger.info("********* Stats: " + stats); } - } catch (Exception e) { - stats.samplesRejectedForException.increment(); - stats.lastException = e; - logger.info("Error performing sampling for " + samplingIndex, e); - } finally { - stats.timeSampling.add((relativeNanoTimeSupplier.getAsLong() - startTime)); - logger.info("********* Stats: " + stats); } } } @@ -249,11 +255,15 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque } public List getSamples(ProjectId projectId, String index) { - return samples.get(new ProjectIndex(projectId, index)).getSamples(); + SoftReference sampleInfoReference = samples.get(new ProjectIndex(projectId, index)); + SampleInfo sampleInfo = sampleInfoReference.get(); + return sampleInfo == null ? List.of() : sampleInfo.getSamples(); } public SampleStats getSampleStats(ProjectId projectId, String index) { - return samples.get(new ProjectIndex(projectId, index)).stats; + SoftReference sampleInfoReference = samples.get(new ProjectIndex(projectId, index)); + SampleInfo sampleInfo = sampleInfoReference.get(); + return sampleInfo == null ? new SampleStats() : sampleInfo.stats; } public TransportPutSampleConfigAction.SamplingConfigCustomMetadata getSampleConfig(ProjectMetadata projectMetadata, String index) { @@ -309,9 +319,9 @@ public void clusterChanged(ClusterChangedEvent event) { if (newSamplingConfig == null && oldSamplingConfig != null) { samples.remove(new ProjectIndex(projectId, oldSamplingConfig.indexName)); } else if (newSamplingConfig != null && newSamplingConfig.equals(oldSamplingConfig) == false) { - samples.computeIfPresent( + samples.put( new ProjectIndex(projectId, newSamplingConfig.indexName), - (s, sampleInfo) -> new SampleInfo(newSamplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong()) + new SoftReference<>(new SampleInfo(newSamplingConfig.timeToLive, relativeNanoTimeSupplier.getAsLong())) ); } } @@ -323,8 +333,9 @@ private void checkTTLs() { long now = relativeNanoTimeSupplier.getAsLong(); Set projectIndices = samples.keySet(); for (ProjectIndex projectIndex : projectIndices) { - SampleInfo sampleInfo = samples.get(projectIndex); - if (sampleInfo.expiration < now) { + SoftReference sampleInfoReference = samples.get(projectIndex); + SampleInfo sampleInfo = sampleInfoReference.get(); + if (sampleInfo != null && sampleInfo.expiration < now) { deleteSampleConfiguration(projectIndex.projectId, projectIndex.indexName); } } @@ -387,23 +398,23 @@ public SampleStats(StreamInput in) throws IOException { @Override public String toString() { - return "potentialSamples: " + return "potential_samples: " + potentialSamples - + ", samplesRejectedForSize: " + + ", samples_rejected_for_size: " + samplesRejectedForSize - + ", samplesRejectedForCondition: " + + ", samples_rejected_for_condition: " + samplesRejectedForCondition - + ", samplesRejectedForRate: " + + ", samples_rejected_for_rate: " + samplesRejectedForRate - + ", samplesRejectedForException: " + + ", samples_rejected_for_exception: " + samplesRejectedForException - + ", samples: " + + ", samples_accepted: " + samples - + ", timeSampling: " + + ", time_sampling: " + (timeSampling.longValue() / 1000000) - + ", timeEvaluatingCondition: " + + ", time_evaluating_condition: " + (timeEvaluatingCondition.longValue() / 1000000) - + ", timeCompilingCondition: " + + ", time_compiling_condition: " + (timeCompilingCondition.longValue() / 1000000); } @@ -434,15 +445,15 @@ public SampleStats combine(SampleStats other) { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); - builder.field("potentialSamples", potentialSamples.longValue()); - builder.field("samplesRejectedForSize", samplesRejectedForSize.longValue()); - builder.field("samplesRejectedForCondition", samplesRejectedForCondition.longValue()); - builder.field("samplesRejectedForRate", samplesRejectedForRate.longValue()); - builder.field("samplesRejectedForException", samplesRejectedForException.longValue()); - builder.field("samples", samples.longValue()); - builder.field("timeSampling", (timeSampling.longValue() / 1000000)); - builder.field("timeEvaluatingCondition", (timeEvaluatingCondition.longValue() / 1000000)); - builder.field("timeCompilingCondition", (timeCompilingCondition.longValue() / 1000000)); + builder.field("potential_samples", potentialSamples.longValue()); + builder.field("samples_rejected_for_size", samplesRejectedForSize.longValue()); + builder.field("samples_rejected_for_condition", samplesRejectedForCondition.longValue()); + builder.field("samples_rejected_for_rate", samplesRejectedForRate.longValue()); + builder.field("samples_rejected_for_exception", samplesRejectedForException.longValue()); + builder.field("samples_accepted", samples.longValue()); + builder.field("time_sampling", (timeSampling.longValue() / 1000000)); + builder.field("time_evaluating_condition", (timeEvaluatingCondition.longValue() / 1000000)); + builder.field("time_compiling_condition", (timeCompilingCondition.longValue() / 1000000)); builder.endObject(); return builder; } From c7e948737d6e20603aed642c4e0b2a82a3912b97 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Thu, 11 Sep 2025 15:45:46 -0500 Subject: [PATCH 20/37] handling bad source --- .../java/org/elasticsearch/ingest/SamplingService.java | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 82eeb8b243a62..0b6f5cbbad3ed 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -45,6 +45,7 @@ import org.elasticsearch.script.ScriptService; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParseException; import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xcontent.json.JsonXContent; @@ -162,6 +163,13 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) { } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { + Map sourceAsMap; + try { + sourceAsMap = indexRequest.sourceAsMap(XContentParserDecorator.NOOP); + } catch (XContentParseException e) { + sourceAsMap = Map.of(); + logger.trace("Invalid index request source, attempting to sample anyway"); + } maybeSample( projectMetadata, indexRequest, @@ -171,7 +179,7 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque indexRequest.version(), indexRequest.routing(), indexRequest.versionType(), - indexRequest.sourceAsMap(XContentParserDecorator.NOOP) + sourceAsMap ) ); } From 518615ba0f99c36cb74b1078a88b56b3cb79a7b3 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Thu, 11 Sep 2025 16:23:02 -0500 Subject: [PATCH 21/37] using less of ConditionalProcessor --- .../elasticsearch/ingest/ConditionalProcessor.java | 14 +++++++++----- .../org/elasticsearch/ingest/SamplingService.java | 9 ++------- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java b/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java index 20db72db39f91..f1f53766009ba 100644 --- a/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java +++ b/server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java @@ -40,7 +40,7 @@ public class ConditionalProcessor extends AbstractProcessor implements WrappingProcessor { private static final DeprecationLogger deprecationLogger = DeprecationLogger.getLogger(DynamicMap.class); - public static final Map> FUNCTIONS = Map.of("_type", value -> { + private static final Map> FUNCTIONS = Map.of("_type", value -> { deprecationLogger.warn( DeprecationCategory.INDICES, "conditional-processor__type", @@ -146,7 +146,7 @@ boolean evaluate(IngestDocument ingestDocument) { } return factory.newInstance( condition.getParams(), - new UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)) + wrapUnmodifiableMap(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)) ).execute(); } @@ -171,8 +171,8 @@ public String getCondition() { private static Object wrapUnmodifiable(Object raw) { // Wraps all mutable types that the JSON parser can create by immutable wrappers. // Any inputs not wrapped are assumed to be immutable - if (raw instanceof Map) { - return new UnmodifiableIngestData((Map) raw); + if (raw instanceof Map rawMap) { + return wrapUnmodifiableMap((Map) rawMap); } else if (raw instanceof List) { return new UnmodifiableIngestList((List) raw); } else if (raw instanceof byte[] bytes) { @@ -185,7 +185,11 @@ private static UnsupportedOperationException unmodifiableException() { return new UnsupportedOperationException("Mutating ingest documents in conditionals is not supported"); } - public static final class UnmodifiableIngestData implements Map { + public static Map wrapUnmodifiableMap(Map map) { + return new UnmodifiableIngestData(map); + } + + private static final class UnmodifiableIngestData implements Map { private final Map data; diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 0b6f5cbbad3ed..b2bea62972f15 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -39,7 +39,6 @@ import org.elasticsearch.logging.Logger; import org.elasticsearch.plugins.internal.XContentParserDecorator; import org.elasticsearch.sample.TransportPutSampleConfigAction; -import org.elasticsearch.script.DynamicMap; import org.elasticsearch.script.IngestConditionalScript; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptService; @@ -60,8 +59,6 @@ import java.util.concurrent.atomic.LongAdder; import java.util.function.LongSupplier; -import static org.elasticsearch.ingest.ConditionalProcessor.FUNCTIONS; - public class SamplingService implements ClusterStateListener { private static final Logger logger = LogManager.getLogger(SamplingService.class); private final ScriptService scriptService; @@ -290,10 +287,8 @@ private boolean evaluateCondition( IngestConditionalScript.Factory factory, SampleStats stats ) { - return factory.newInstance( - script.getParams(), - new ConditionalProcessor.UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)) - ).execute(); + return factory.newInstance(script.getParams(), ConditionalProcessor.wrapUnmodifiableMap(ingestDocument.getSourceAndMetadata())) + .execute(); } private static Script getScript(String conditional) throws IOException { From 6335456f31f04186b2cea5c0d7a2d434efec4f80 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 12 Sep 2025 10:38:49 -0500 Subject: [PATCH 22/37] fixing compilation error after merge --- .../main/java/org/elasticsearch/ingest/SamplingService.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index b2bea62972f15..a085dbe9d4072 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -37,7 +37,6 @@ import org.elasticsearch.core.Tuple; import org.elasticsearch.logging.LogManager; import org.elasticsearch.logging.Logger; -import org.elasticsearch.plugins.internal.XContentParserDecorator; import org.elasticsearch.sample.TransportPutSampleConfigAction; import org.elasticsearch.script.IngestConditionalScript; import org.elasticsearch.script.Script; @@ -162,7 +161,7 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) { public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { Map sourceAsMap; try { - sourceAsMap = indexRequest.sourceAsMap(XContentParserDecorator.NOOP); + sourceAsMap = indexRequest.sourceAsMap(); } catch (XContentParseException e) { sourceAsMap = Map.of(); logger.trace("Invalid index request source, attempting to sample anyway"); From ef920e5d0f4207ec8ff13ca2333cfba4de75865f Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 12 Sep 2025 12:46:32 -0500 Subject: [PATCH 23/37] moving TTL logic to master node only --- .../src/main/groovy/elasticsearch.run.gradle | 2 +- .../elasticsearch/ingest/SamplingService.java | 82 ++++++++++++++++++- .../elasticsearch/node/NodeConstruction.java | 3 +- 3 files changed, 81 insertions(+), 6 deletions(-) diff --git a/build-tools-internal/src/main/groovy/elasticsearch.run.gradle b/build-tools-internal/src/main/groovy/elasticsearch.run.gradle index 90185179ef6ad..f0c8a9b52cba8 100644 --- a/build-tools-internal/src/main/groovy/elasticsearch.run.gradle +++ b/build-tools-internal/src/main/groovy/elasticsearch.run.gradle @@ -31,7 +31,7 @@ testClusters.register("runTask") { setting 'xpack.security.enabled', 'true' keystore 'bootstrap.password', 'password' user username: 'elastic-admin', password: 'elastic-password', role: '_es_test_root' - numberOfNodes = 1 + numberOfNodes = 3 } } diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index a085dbe9d4072..4d37642045cb7 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -9,6 +9,7 @@ package org.elasticsearch.ingest; +import org.apache.lucene.util.SetOnce; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.master.AcknowledgedResponse; @@ -26,9 +27,13 @@ import org.elasticsearch.common.Priority; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.bytes.ReleasableBytesReference; +import org.elasticsearch.common.component.Lifecycle; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.common.scheduler.SchedulerEngine; +import org.elasticsearch.common.scheduler.TimeValueSchedule; +import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; @@ -50,6 +55,7 @@ import java.io.IOException; import java.lang.ref.SoftReference; +import java.time.Clock; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -58,17 +64,33 @@ import java.util.concurrent.atomic.LongAdder; import java.util.function.LongSupplier; -public class SamplingService implements ClusterStateListener { +public class SamplingService implements ClusterStateListener, SchedulerEngine.Listener { private static final Logger logger = LogManager.getLogger(SamplingService.class); private final ScriptService scriptService; + private final ClusterService clusterService; private final LongSupplier relativeNanoTimeSupplier; private final MasterServiceTaskQueue updateSamplingConfigTaskQueue; private final MasterServiceTaskQueue deleteSamplingConfigTaskQueue; private final Map> samples = new HashMap<>(); - - public SamplingService(ScriptService scriptService, ClusterService clusterService, LongSupplier relativeNanoTimeSupplier) { + private volatile boolean isMaster = false; + private final SetOnce scheduler = new SetOnce<>(); + private SchedulerEngine.Job scheduledJob; + private final Clock clock; + private final Settings settings; + private volatile TimeValue pollInterval = TimeValue.timeValueMinutes(30); + + public SamplingService( + ScriptService scriptService, + ClusterService clusterService, + LongSupplier relativeNanoTimeSupplier, + Clock clock, + Settings settings + ) { this.scriptService = scriptService; + this.clusterService = clusterService; this.relativeNanoTimeSupplier = relativeNanoTimeSupplier; + this.clock = clock; + this.settings = settings; ClusterStateTaskExecutor updateSampleConfigExecutor = new SimpleBatchedAckListenerTaskExecutor<>() { @Override @@ -118,6 +140,7 @@ public Tuple executeTask( ClusterState updatedClusterState = ClusterState.builder(clusterState) .putProjectMetadata(projectMetadataBuilder) .build(); + logger.info("Removing sampling config " + samplingConfig.indexName + " from cluster state"); return new Tuple<>(updatedClusterState, deleteSamplingConfigTask); } else { return null; // someone beat us to it. This seems like a bad plan TODO @@ -151,6 +174,8 @@ public void updateSampleConfiguration( } public void deleteSampleConfiguration(ProjectId projectId, String index) { + logger.info("Calling deleteSampleConfiguration"); + // to be called by the master node deleteSamplingConfigTaskQueue.submitTask( "deleting sampling config", new DeleteSampleConfigTask(projectId, index, TimeValue.THIRTY_SECONDS, ActionListener.noop()), @@ -255,7 +280,7 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque } } } - checkTTLs(); // TODO make this happen less often? + // checkTTLs(); // TODO make this happen less often? } public List getSamples(ProjectId projectId, String index) { @@ -304,8 +329,49 @@ private static Script getScript(String conditional) throws IOException { } } + private synchronized void maybeScheduleJob() { + if (this.isMaster) { + if (scheduler.get() == null) { + // don't create scheduler if the node is shutting down + if (isClusterServiceStoppedOrClosed() == false) { + scheduler.set(new SchedulerEngine(settings, clock)); + scheduler.get().register(this); + } + } + + // scheduler could be null if the node might be shutting down + if (scheduler.get() != null) { + scheduledJob = new SchedulerEngine.Job("sampling_ttl", new TimeValueSchedule(pollInterval)); + scheduler.get().add(scheduledJob); + } + } + } + + private void cancelJob() { + if (scheduler.get() != null) { + scheduler.get().remove("sampling_ttl"); + scheduledJob = null; + } + } + + private boolean isClusterServiceStoppedOrClosed() { + final Lifecycle.State state = clusterService.lifecycleState(); + return state == Lifecycle.State.STOPPED || state == Lifecycle.State.CLOSED; + } + @Override public void clusterChanged(ClusterChangedEvent event) { + final boolean prevIsMaster = this.isMaster; + if (prevIsMaster != event.localNodeMaster()) { + this.isMaster = event.localNodeMaster(); + if (this.isMaster) { + // we weren't the master, and now we are + maybeScheduleJob(); + } else { + // we were the master, and now we aren't + cancelJob(); + } + } if (event.metadataChanged()) { for (ProjectMetadata projectMetadata : event.state().metadata().projects().values()) { ProjectId projectId = projectMetadata.id(); @@ -319,6 +385,7 @@ public void clusterChanged(ClusterChangedEvent event) { .metadata() .custom(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME); if (newSamplingConfig == null && oldSamplingConfig != null) { + logger.info("Removing sampling config info from buffer because it has been deleted from cluster state"); samples.remove(new ProjectIndex(projectId, oldSamplingConfig.indexName)); } else if (newSamplingConfig != null && newSamplingConfig.equals(oldSamplingConfig) == false) { samples.put( @@ -339,10 +406,17 @@ private void checkTTLs() { SampleInfo sampleInfo = sampleInfoReference.get(); if (sampleInfo != null && sampleInfo.expiration < now) { deleteSampleConfiguration(projectIndex.projectId, projectIndex.indexName); + // samples.remove(new ProjectIndex(projectIndex.projectId, projectIndex.indexName)); } } } + @Override + public void triggered(SchedulerEngine.Event event) { + logger.info("job triggered: {}, {}, {}", event.jobName(), event.scheduledTime(), event.triggeredTime()); + checkTTLs(); + } + private static final class SampleInfo { private final List samples; private final SampleStats stats; diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index 2f42801d76bf7..e29588ccb257f 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -235,6 +235,7 @@ import java.io.Closeable; import java.io.IOException; import java.io.UncheckedIOException; +import java.time.Clock; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -716,7 +717,7 @@ private void construct( modules.bindToInstance(DocumentParsingProvider.class, documentParsingProvider); FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); - SamplingService samplingService = new SamplingService(scriptService, clusterService, System::nanoTime); + SamplingService samplingService = new SamplingService(scriptService, clusterService, System::nanoTime, Clock.systemUTC(), settings); modules.bindToInstance(SamplingService.class, samplingService); clusterService.addListener(samplingService); From 6918c5877bce5d33b95b7da890c5cfe4a0241eee Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 12 Sep 2025 12:48:47 -0500 Subject: [PATCH 24/37] removing accidental commit --- build-tools-internal/src/main/groovy/elasticsearch.run.gradle | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build-tools-internal/src/main/groovy/elasticsearch.run.gradle b/build-tools-internal/src/main/groovy/elasticsearch.run.gradle index f0c8a9b52cba8..90185179ef6ad 100644 --- a/build-tools-internal/src/main/groovy/elasticsearch.run.gradle +++ b/build-tools-internal/src/main/groovy/elasticsearch.run.gradle @@ -31,7 +31,7 @@ testClusters.register("runTask") { setting 'xpack.security.enabled', 'true' keystore 'bootstrap.password', 'password' user username: 'elastic-admin', password: 'elastic-password', role: '_es_test_root' - numberOfNodes = 3 + numberOfNodes = 1 } } From 42b3b8acc97941330e18bad2828d02579561209e Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 12 Sep 2025 14:48:51 -0500 Subject: [PATCH 25/37] Sampling on pipeline exception --- .../org/elasticsearch/ingest/IngestService.java | 13 ++++++++++++- .../org/elasticsearch/ingest/SamplingService.java | 4 ++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index a5361068bcaba..cffba7f9a02cc 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -101,6 +101,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; @@ -1197,7 +1198,7 @@ private void executePipelines( listener.onFailure(e); } }; - + AtomicBoolean haveAttemptedSampling = new AtomicBoolean(false); try { if (pipeline == null) { throw new IllegalArgumentException("pipeline with id [" + pipelineId + "] does not exist"); @@ -1360,6 +1361,7 @@ private void executePipelines( * 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); @@ -1373,6 +1375,15 @@ private void executePipelines( } }); } 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"); + } // 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( diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 4d37642045cb7..321aa544bf6c0 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -183,7 +183,7 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) { ); } - public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) throws IOException { + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) { Map sourceAsMap; try { sourceAsMap = indexRequest.sourceAsMap(); @@ -205,7 +205,7 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque ); } - public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) throws IOException { + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) { long startTime = relativeNanoTimeSupplier.getAsLong(); TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME From 5349303d4526bc11561ed0e683ba7067de0bb886 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 12 Sep 2025 14:59:08 -0500 Subject: [PATCH 26/37] fixing a bad merge --- server/src/main/java/org/elasticsearch/TransportVersions.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index 0f87495f8ae2b..bb63e7d141f79 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -326,7 +326,7 @@ static TransportVersion def(int id) { public static final TransportVersion INDEX_SOURCE = def(9_158_0_00); public static final TransportVersion MAX_HEAP_SIZE_PER_NODE_IN_CLUSTER_INFO = def(9_159_0_00); public static final TransportVersion TIMESERIES_DEFAULT_LIMIT = def(9_160_0_00); - public static final TransportVersion RANDOM_SAMPLING = def(9_160_0_00); + public static final TransportVersion RANDOM_SAMPLING = def(9_161_0_00); /* * STOP! READ THIS FIRST! No, really, From 255f45ba1e0aad779bbb4c78b34ab5f72bd1d03c Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Mon, 15 Sep 2025 14:48:19 -0500 Subject: [PATCH 27/37] performance improvements --- .../elasticsearch/ingest/SamplingService.java | 91 +++++++++++-------- 1 file changed, 52 insertions(+), 39 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 321aa544bf6c0..43608f28b9ff6 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -63,6 +63,7 @@ import java.util.Set; import java.util.concurrent.atomic.LongAdder; import java.util.function.LongSupplier; +import java.util.function.Supplier; public class SamplingService implements ClusterStateListener, SchedulerEngine.Listener { private static final Logger logger = LogManager.getLogger(SamplingService.class); @@ -184,28 +185,30 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) { } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) { - Map sourceAsMap; - try { - sourceAsMap = indexRequest.sourceAsMap(); - } catch (XContentParseException e) { - sourceAsMap = Map.of(); - logger.trace("Invalid index request source, attempting to sample anyway"); - } - maybeSample( - projectMetadata, - indexRequest, - new IngestDocument( + maybeSample(projectMetadata, indexRequest, () -> { + Map sourceAsMap; + try { + sourceAsMap = indexRequest.sourceAsMap(); + } catch (XContentParseException e) { + sourceAsMap = Map.of(); + logger.trace("Invalid index request source, attempting to sample anyway"); + } + return new IngestDocument( indexRequest.index(), indexRequest.id(), indexRequest.version(), indexRequest.routing(), indexRequest.versionType(), sourceAsMap - ) - ); + ); + }); } public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) { + maybeSample(projectMetadata, indexRequest, () -> ingestDocument); + } + + private void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, Supplier ingestDocumentSupplier) { long startTime = relativeNanoTimeSupplier.getAsLong(); TransportPutSampleConfigAction.SamplingConfigCustomMetadata samplingConfig = projectMetadata.custom( TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME @@ -226,33 +229,43 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque stats.potentialSamples.increment(); try { if (sampleInfo.getSamples().size() < samplingConfig.maxSamples) { - String condition = samplingConfig.condition; - if (condition != null) { - if (sampleInfo.script == null || sampleInfo.factory == null) { - // We don't want to pay for synchronization because worst case, we compile the script twice - long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); - try { - if (sampleInfo.compilationFailed) { - // we don't want to waste time - stats.samplesRejectedForException.increment(); - return; - } else { - Script script = getScript(condition); - sampleInfo.setScript(script, scriptService.compile(script, IngestConditionalScript.CONTEXT)); + if (Math.random() < samplingConfig.rate) { + String condition = samplingConfig.condition; + if (condition != null) { + if (sampleInfo.script == null || sampleInfo.factory == null) { + // We don't want to pay for synchronization because worst case, we compile the script twice + long compileScriptStartTime = relativeNanoTimeSupplier.getAsLong(); + try { + if (sampleInfo.compilationFailed) { + // we don't want to waste time + stats.samplesRejectedForException.increment(); + return; + } else { + Script script = getScript(condition); + sampleInfo.setScript( + script, + scriptService.compile(script, IngestConditionalScript.CONTEXT) + ); + } + } catch (Exception e) { + sampleInfo.compilationFailed = true; + throw e; + } finally { + stats.timeCompilingCondition.add( + (relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime) + ); } - } catch (Exception e) { - sampleInfo.compilationFailed = true; - throw e; - } finally { - stats.timeCompilingCondition.add((relativeNanoTimeSupplier.getAsLong() - compileScriptStartTime)); } } - } - long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); - if (condition == null - || evaluateCondition(ingestDocument, sampleInfo.script, sampleInfo.factory, sampleInfo.stats)) { - stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); - if (Math.random() < samplingConfig.rate) { + long conditionStartTime = relativeNanoTimeSupplier.getAsLong(); + if (condition == null + || evaluateCondition( + ingestDocumentSupplier.get(), + sampleInfo.script, + sampleInfo.factory, + sampleInfo.stats + )) { + stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); indexRequest.incRef(); if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) { releaseableSource.incRef(); @@ -261,10 +274,10 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque stats.samples.increment(); logger.info("Sampling " + indexRequest); } else { - stats.samplesRejectedForRate.increment(); + stats.samplesRejectedForCondition.increment(); } } else { - stats.samplesRejectedForCondition.increment(); + stats.samplesRejectedForRate.increment(); } } else { stats.samplesRejectedForSize.increment(); From 8da7c26e7e50ec65b7cc7ae2216de0510de21dae Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Mon, 15 Sep 2025 17:47:56 -0500 Subject: [PATCH 28/37] Adding the shell of the SamplingService --- .../bulk/TransportAbstractBulkAction.java | 31 +++++-- .../action/bulk/TransportBulkAction.java | 19 +++-- .../bulk/TransportSimulateBulkAction.java | 3 +- .../elasticsearch/ingest/IngestService.java | 83 ++++++++++++++++-- .../elasticsearch/ingest/SamplingService.java | 72 ++++++++++++++++ .../elasticsearch/node/NodeConstruction.java | 8 +- .../bulk/TransportBulkActionIngestTests.java | 10 ++- .../action/bulk/TransportBulkActionTests.java | 26 +++++- .../bulk/TransportBulkActionTookTests.java | 5 +- .../ingest/IngestServiceTests.java | 85 +++++++++++++++++-- .../ingest/SimulateIngestServiceTests.java | 3 +- .../snapshots/SnapshotResiliencyTests.java | 7 +- 12 files changed, 315 insertions(+), 37 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/ingest/SamplingService.java diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java index 9183f77448a2f..33fb86678430e 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java +++ b/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 bulkAction; protected final FeatureService featureService; + protected final SamplingService samplingService; public TransportAbstractBulkAction( ActionType 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, true); } }); } - private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor executor, ActionListener listener) - throws IOException { + private boolean applyPipelines( + Task task, + BulkRequest bulkRequest, + Executor executor, + ActionListener listener, + boolean firstTime + ) throws IOException { boolean hasIndexRequestsWithPipelines = false; ClusterState state = clusterService.state(); ProjectId projectId = projectResolver.getProjectId(); @@ -303,6 +312,13 @@ private boolean applyPipelines(Task task, BulkRequest bulkRequest, Executor exec } }); return true; + } else if (firstTime && samplingService != null && samplingService.atLeastOneSampleConfigured()) { + // else sample, but only if this is the first time through. Otherwise we had pipelines and sampled in IngestService + for (DocWriteRequest actionRequest : bulkRequest.requests) { + if (actionRequest instanceof IndexRequest ir) { + samplingService.maybeSample(project, ir); + } + } } return false; } @@ -338,7 +354,7 @@ private void processBulkIndexIngestRequest( ActionRunnable runnable = new ActionRunnable<>(actionListener) { @Override protected void doRun() throws IOException { - applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener); + applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, false); } @Override @@ -416,7 +432,8 @@ private void applyPipelinesAndDoInternalExecute( Task task, BulkRequest bulkRequest, Executor executor, - ActionListener listener + ActionListener listener, + boolean firstTime ) throws IOException { final long relativeStartTimeNanos = relativeTimeNanos(); @@ -434,7 +451,7 @@ private void applyPipelinesAndDoInternalExecute( var wrappedListener = bulkRequestModifier.wrapActionListenerIfNeeded(listener); - if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener) == false) { + if (applyPipelines(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, firstTime) == false) { doInternalExecute(task, bulkRequestModifier.getBulkRequest(), executor, wrappedListener, relativeStartTimeNanos); } } diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java index 7e443e055cc90..81d60886b7bab 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportBulkAction.java +++ b/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); diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java index b52f5447b9311..6712920b3bf85 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportSimulateBulkAction.java +++ b/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; diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 2c0359e367086..080e44cedc9ac 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/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,10 +78,12 @@ 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; +import java.io.IOException; import java.time.Instant; import java.time.InstantSource; import java.util.ArrayList; @@ -98,6 +101,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 +155,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> nodeInfoListener; private static BiFunction createScheduler(ThreadPool threadPool) { @@ -252,6 +257,7 @@ public IngestService( FailureStoreMetrics failureStoreMetrics, ProjectResolver projectResolver, FeatureService featureService, + SamplingService samplingService, Consumer> nodeInfoListener ) { this.clusterService = clusterService; @@ -276,6 +282,7 @@ public IngestService( this.failureStoreMetrics = failureStoreMetrics; this.projectResolver = projectResolver; this.featureService = featureService; + this.samplingService = samplingService; this.nodeInfoListener = nodeInfoListener; } @@ -290,7 +297,8 @@ public IngestService( MatcherWatchdog matcherWatchdog, FailureStoreMetrics failureStoreMetrics, ProjectResolver projectResolver, - FeatureService featureService + FeatureService featureService, + SamplingService samplingService ) { this( clusterService, @@ -304,6 +312,7 @@ public IngestService( failureStoreMetrics, projectResolver, featureService, + samplingService, createNodeInfoListener(client) ); } @@ -324,6 +333,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 +981,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 +994,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 documentListener = ActionListener.runAfter( @@ -1030,7 +1041,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 +1167,8 @@ private void executePipelines( final IndexRequest indexRequest, final IngestDocument ingestDocument, final Function resolveFailureStore, - final ActionListener listener + final ActionListener listener, + final Metadata originalDocumentMetadata ) { assert pipelines.hasNext(); PipelineSlot slot = pipelines.next(); @@ -1180,12 +1199,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 +1354,21 @@ 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) { + 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 +1377,48 @@ private void executePipelines( } } + private void attemptToSampleData( + ProjectMetadata projectMetadata, + IndexRequest indexRequest, + IngestDocument ingestDocument, + Metadata originalDocumentMetadata + ) { + if (samplingService != null && samplingService.atLeastOneSampleConfigured()) { + try { + /* + * 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. + */ + IndexRequest original = copyIndexRequest(indexRequest); + updateIndexRequestMetadata(original, originalDocumentMetadata); + samplingService.maybeSample(projectMetadata, original, ingestDocument); + } catch (IOException ex) { + logger.warn("unable to sample data"); + } + } + } + + private IndexRequest copyIndexRequest(IndexRequest original) throws IOException { + 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.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()); + BytesReference source = original.source(); + if (source != null) { + clonedRequest.source(source, original.getContentType()); + } + return clonedRequest; + } + private static void executePipeline( final IngestDocument ingestDocument, final Pipeline pipeline, diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java new file mode 100644 index 0000000000000..1d119f0694928 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -0,0 +1,72 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.ingest; + +import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.cluster.ClusterChangedEvent; +import org.elasticsearch.cluster.ClusterStateListener; +import org.elasticsearch.cluster.metadata.ProjectMetadata; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.logging.LogManager; +import org.elasticsearch.logging.Logger; +import org.elasticsearch.script.ScriptService; +import org.elasticsearch.xcontent.XContentParseException; + +import java.util.Map; +import java.util.function.Supplier; + +public class SamplingService implements ClusterStateListener { + private static final Logger logger = LogManager.getLogger(SamplingService.class); + private final ScriptService scriptService; + private final ClusterService clusterService; + + public SamplingService(ScriptService scriptService, ClusterService clusterService) { + this.scriptService = scriptService; + this.clusterService = clusterService; + } + + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) { + maybeSample(projectMetadata, indexRequest, () -> { + Map sourceAsMap; + try { + sourceAsMap = indexRequest.sourceAsMap(); + } catch (XContentParseException e) { + sourceAsMap = Map.of(); + logger.trace("Invalid index request source, attempting to sample anyway"); + } + return new IngestDocument( + indexRequest.index(), + indexRequest.id(), + indexRequest.version(), + indexRequest.routing(), + indexRequest.versionType(), + sourceAsMap + ); + }); + } + + public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) { + maybeSample(projectMetadata, indexRequest, () -> ingestDocument); + } + + private void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, Supplier ingestDocumentSupplier) { + // TODO Sampling logic to go here in the near future + } + + 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) { + // TODO: React to sampling config changes + } + +} diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index 72c48c303ffa9..57f3dda579819 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -142,6 +142,7 @@ import org.elasticsearch.indices.recovery.plan.RecoveryPlannerService; import org.elasticsearch.indices.recovery.plan.ShardSnapshotsService; import org.elasticsearch.ingest.IngestService; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.injection.guice.Injector; import org.elasticsearch.injection.guice.Key; import org.elasticsearch.injection.guice.Module; @@ -716,6 +717,10 @@ private void construct( FeatureService featureService = new FeatureService(pluginsService.loadServiceProviders(FeatureSpecification.class)); + SamplingService samplingService = new SamplingService(scriptService, clusterService); + modules.bindToInstance(SamplingService.class, samplingService); + clusterService.addListener(samplingService); + FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); final IngestService ingestService = new IngestService( clusterService, @@ -728,7 +733,8 @@ private void construct( IngestService.createGrokThreadWatchdog(environment, threadPool), failureStoreMetrics, projectResolver, - featureService + featureService, + samplingService ); SystemIndices systemIndices = createSystemIndices(settings); diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java index a54cd08c3738a..119385319e52f 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java +++ b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionIngestTests.java @@ -54,6 +54,7 @@ import org.elasticsearch.indices.EmptySystemIndices; import org.elasticsearch.indices.TestIndexNameExpressionResolver; import org.elasticsearch.ingest.IngestService; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.tasks.Task; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; @@ -172,10 +173,17 @@ class TestTransportBulkAction extends TransportBulkAction { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + initializeSamplingService() ); } + private static SamplingService initializeSamplingService() { + SamplingService samplingService = mock(SamplingService.class); + when(samplingService.atLeastOneSampleConfigured()).thenReturn(true); + return samplingService; + } + @Override void executeBulk( Task task, diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java index 481fdf5ea3530..4242c44e29808 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java +++ b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTests.java @@ -63,6 +63,7 @@ import org.elasticsearch.index.VersionType; import org.elasticsearch.indices.SystemIndexDescriptorUtils; import org.elasticsearch.indices.SystemIndices; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.VersionUtils; import org.elasticsearch.test.index.IndexVersionUtils; @@ -81,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; @@ -92,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 { @@ -152,10 +156,17 @@ public ProjectId getProjectId() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + initializeSamplingService() ); } + private static SamplingService initializeSamplingService() { + SamplingService samplingService = mock(SamplingService.class); + when(samplingService.atLeastOneSampleConfigured()).thenReturn(true); + return samplingService; + } + @Override void createIndex(CreateIndexRequest createIndexRequest, ActionListener listener) { indexCreated = true; @@ -733,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 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 indices) { BulkRequest request = new BulkRequest(); for (String index : indices) { diff --git a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java b/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java index 0077f739bf7b6..fb2e8963614a3 100644 --- a/server/src/test/java/org/elasticsearch/action/bulk/TransportBulkActionTookTests.java +++ b/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 { @@ -267,7 +269,8 @@ static class TestTransportBulkAction extends TransportBulkAction { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } } diff --git a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java index df3211331c7f5..309588c0831a5 100644 --- a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java @@ -171,7 +171,8 @@ public void testIngestPlugin() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); Map factories = ingestService.getProcessorFactories(); assertTrue(factories.containsKey("foo")); @@ -198,7 +199,8 @@ public void testIngestPluginDuplicate() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ) ); assertTrue(e.getMessage(), e.getMessage().contains("already registered")); @@ -222,7 +224,8 @@ public void testExecuteIndexPipelineDoesNotExist() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); final IndexRequest indexRequest = new IndexRequest("_index").id("_id") .source(Map.of()) @@ -1833,7 +1836,8 @@ public void testFailureRedirectionWithoutNodeFeatureEnabled() throws Exception { "set", (factories, tag, description, config, projectId) -> new FakeProcessor("set", "", "", (ingestDocument) -> fail()) ), - Predicates.never() + Predicates.never(), + mock(SamplingService.class) ); PutPipelineRequest putRequest1 = putJsonPipelineRequest("_id1", "{\"processors\": [{\"mock\" : {}}]}"); // given that set -> fail() above, it's a failure if a document executes against this pipeline @@ -2567,7 +2571,8 @@ public Map getProcessors(Processor.Parameters paramet public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); ingestService.addIngestClusterStateListener(ingestClusterStateListener); @@ -3075,6 +3080,7 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } }, + mock(SamplingService.class), consumer ); ingestService.applyClusterState(new ClusterChangedEvent("", clusterState, clusterState)); @@ -3337,6 +3343,64 @@ public void testResolvePipelinesWithNonePipeline() { } } + public void testSampling() { + SamplingService samplingService = mock(SamplingService.class); + IngestService ingestService = createWithProcessors( + Map.of("mock", (factories, tag, description, config, projectId) -> mockCompoundProcessor()), + Predicates.never(), + samplingService + ); + when(samplingService.atLeastOneSampleConfigured()).thenReturn(true); + PutPipelineRequest putRequest = putJsonPipelineRequest("_id", "{\"processors\": [{\"mock\" : {}}]}"); + var projectId = randomProjectIdOrDefault(); + ClusterState clusterState = ClusterState.builder(ClusterName.DEFAULT) + .putProjectMetadata(ProjectMetadata.builder(projectId).build()) + .build(); + ClusterState previousClusterState = clusterState; + clusterState = executePut(projectId, putRequest, clusterState); + ingestService.applyClusterState(new ClusterChangedEvent("", clusterState, previousClusterState)); + + BulkRequest bulkRequest = new BulkRequest(); + + IndexRequest indexRequest1 = new IndexRequest("_index").id("_id1").source(Map.of()).setPipeline("_none").setFinalPipeline("_none"); + bulkRequest.add(indexRequest1); + IndexRequest indexRequest2 = new IndexRequest("_index").id("_id2").source(Map.of()).setPipeline("_id").setFinalPipeline("_none"); + bulkRequest.add(indexRequest2); + IndexRequest indexRequest3 = new IndexRequest("_index").id("_id3") + .source(Map.of()) + .setPipeline("does_not_exist") + .setFinalPipeline("_none"); + bulkRequest.add(indexRequest3); + @SuppressWarnings("unchecked") + TriConsumer failureHandler = mock(TriConsumer.class); + @SuppressWarnings("unchecked") + final ActionListener listener = mock(ActionListener.class); + + Boolean noRedirect = randomBoolean() ? false : null; + IndexDocFailureStoreStatus fsStatus = IndexDocFailureStoreStatus.NOT_APPLICABLE_OR_UNKNOWN; + + ingestService.executeBulkRequest( + projectId, + bulkRequest.numberOfActions(), + bulkRequest.requests(), + indexReq -> {}, + (s) -> noRedirect, + (slot, targetIndex, e) -> fail("Should not be redirecting failures"), + failureHandler, + listener + ); + verify(failureHandler, times(1)).apply( + argThat(item -> item == 2), + argThat(iae -> "pipeline with id [does_not_exist] does not exist".equals(iae.getMessage())), + argThat(fsStatus::equals) + ); + verify(listener, times(1)).onResponse(null); + // In the case where there is a pipeline, or there is a pipeline failure, there will be an IngestDocument so this verion is called: + verify(samplingService, times(2)).maybeSample(any(), any(), any()); + // When there is no pipeline, we have no IngestDocument, and the maybeSample that does not require an IngestDocument is called: + verify(samplingService, times(1)).maybeSample(any(), any()); + } + private static Tuple randomMapEntry() { return tuple(randomAlphaOfLength(5), randomObject()); } @@ -3377,10 +3441,14 @@ private static IngestService createWithProcessors() { } private static IngestService createWithProcessors(Map processors) { - return createWithProcessors(processors, DataStream.DATA_STREAM_FAILURE_STORE_FEATURE::equals); + return createWithProcessors(processors, DataStream.DATA_STREAM_FAILURE_STORE_FEATURE::equals, mock(SamplingService.class)); } - private static IngestService createWithProcessors(Map processors, Predicate featureTest) { + private static IngestService createWithProcessors( + Map processors, + Predicate featureTest, + SamplingService samplingService + ) { Client client = mock(Client.class); ThreadPool threadPool = mock(ThreadPool.class); when(threadPool.generic()).thenReturn(EsExecutors.DIRECT_EXECUTOR_SERVICE); @@ -3406,7 +3474,8 @@ public Map getProcessors(final Processor.Parameters p public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return featureTest.test(feature); } - } + }, + samplingService ); if (randomBoolean()) { /* diff --git a/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java index 6733ff4c44e6e..79b4e76932fb3 100644 --- a/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/SimulateIngestServiceTests.java @@ -173,7 +173,8 @@ public Map getProcessors(final Processor.Parameters p public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } } diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java index 6a459cab07328..d6a57ba4587c5 100644 --- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotResiliencyTests.java @@ -166,6 +166,7 @@ import org.elasticsearch.indices.recovery.SnapshotFilesProvider; import org.elasticsearch.indices.recovery.plan.PeerOnlyRecoveryPlannerService; import org.elasticsearch.ingest.IngestService; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.monitor.StatusInfo; import org.elasticsearch.node.ResponseCollectorService; import org.elasticsearch.plugins.PluginsService; @@ -2666,7 +2667,8 @@ public RecyclerBytesStreamOutput newNetworkBytesStream() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ), client, actionFilters, @@ -2681,7 +2683,8 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ) ); final TransportShardBulkAction transportShardBulkAction = new TransportShardBulkAction( From 79c77fa2d09159ccc0e4e3a86f96ad5c1a015354 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 16 Sep 2025 09:16:22 -0500 Subject: [PATCH 29/37] fixing tests --- .../org/elasticsearch/ingest/IngestServiceTests.java | 9 ++++++++- .../TransportGetTrainedModelsStatsActionTests.java | 4 +++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java index 309588c0831a5..4ff762595d354 100644 --- a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java @@ -246,11 +246,18 @@ public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { assertThat(status, equalTo(fsStatus)); }; + // This is due to a quirk of IngestService. It uses a cluster state from the most recent cluster change event: + ProjectId projectId = randomProjectIdOrDefault(); + ClusterState clusterState = ClusterState.builder(ClusterName.DEFAULT) + .putProjectMetadata(ProjectMetadata.builder(projectId).build()) + .build(); + ingestService.applyClusterState(new ClusterChangedEvent("", clusterState, clusterState)); + @SuppressWarnings("unchecked") final ActionListener listener = mock(ActionListener.class); ingestService.executeBulkRequest( - randomProjectIdOrDefault(), + projectId, 1, List.of(indexRequest), indexReq -> {}, diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java index 06ba7ba113d4e..85f0aa5c40a35 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/action/TransportGetTrainedModelsStatsActionTests.java @@ -31,6 +31,7 @@ import org.elasticsearch.ingest.IngestService; import org.elasticsearch.ingest.IngestStats; import org.elasticsearch.ingest.Processor; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.license.MockLicenseState; import org.elasticsearch.plugins.IngestPlugin; import org.elasticsearch.test.ESTestCase; @@ -159,7 +160,8 @@ public void setUpVariables() { public boolean clusterHasFeature(ClusterState state, NodeFeature feature) { return DataStream.DATA_STREAM_FAILURE_STORE_FEATURE.equals(feature); } - } + }, + mock(SamplingService.class) ); } From 44bdca4a9da5267fd6107c7ad2c44ea052d83fdc Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Tue, 16 Sep 2025 16:03:55 -0500 Subject: [PATCH 30/37] do not sample twice on exception --- .../main/java/org/elasticsearch/ingest/IngestService.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 080e44cedc9ac..2495bc61348b3 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -1368,7 +1368,10 @@ private void executePipelines( } }); } catch (Exception e) { - attemptToSampleData(project, indexRequest, ingestDocument, originalDocumentMetadata); + 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 From 7af6e27c964c7b9439e79bcd92452ebac707ee70 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 17 Sep 2025 14:38:32 -0500 Subject: [PATCH 31/37] renaming parameter --- .../bulk/TransportAbstractBulkAction.java | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java index 33fb86678430e..f7460dd3de47d 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java +++ b/server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java @@ -208,7 +208,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); } }); } @@ -218,7 +218,7 @@ private boolean applyPipelines( BulkRequest bulkRequest, Executor executor, ActionListener listener, - boolean firstTime + boolean haveRunIngestService ) throws IOException { boolean hasIndexRequestsWithPipelines = false; ClusterState state = clusterService.state(); @@ -312,8 +312,11 @@ private boolean applyPipelines( } }); return true; - } else if (firstTime && samplingService != null && samplingService.atLeastOneSampleConfigured()) { - // 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); @@ -354,7 +357,7 @@ private void processBulkIndexIngestRequest( ActionRunnable runnable = new ActionRunnable<>(actionListener) { @Override protected void doRun() throws IOException { - applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, false); + applyPipelinesAndDoInternalExecute(task, bulkRequest, executor, actionListener, true); } @Override @@ -433,7 +436,7 @@ private void applyPipelinesAndDoInternalExecute( BulkRequest bulkRequest, Executor executor, ActionListener listener, - boolean firstTime + boolean haveRunIngestService ) throws IOException { final long relativeStartTimeNanos = relativeTimeNanos(); @@ -451,7 +454,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); } } From 6ba15dd265acfff959f1c6638960004f989a8572 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 17 Sep 2025 14:44:43 -0500 Subject: [PATCH 32/37] reducing logging --- .../src/main/java/org/elasticsearch/ingest/IngestService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 2495bc61348b3..667d5ee41d2eb 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -1396,7 +1396,7 @@ private void attemptToSampleData( updateIndexRequestMetadata(original, originalDocumentMetadata); samplingService.maybeSample(projectMetadata, original, ingestDocument); } catch (IOException ex) { - logger.warn("unable to sample data"); + logger.debug("Error attempting to sample data", ex); } } } From 342cbefb3a29f30ee47bfc7526004f557a228c83 Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Wed, 17 Sep 2025 17:48:45 -0500 Subject: [PATCH 33/37] do not copy an IndexRequest unless necessary --- .../elasticsearch/ingest/IngestService.java | 31 ++++++++++++------- .../elasticsearch/ingest/SamplingService.java | 29 ++++++++++++++--- .../ingest/IngestServiceTests.java | 2 +- 3 files changed, 45 insertions(+), 17 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 667d5ee41d2eb..2265a0343576c 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/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; @@ -1387,27 +1386,32 @@ private void attemptToSampleData( Metadata originalDocumentMetadata ) { if (samplingService != null && samplingService.atLeastOneSampleConfigured()) { - try { - /* - * 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. - */ - IndexRequest original = copyIndexRequest(indexRequest); + /* + * 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); - samplingService.maybeSample(projectMetadata, original, ingestDocument); - } catch (IOException ex) { - logger.debug("Error attempting to sample data", ex); - } + return original; + }, ingestDocument); + } } - private IndexRequest copyIndexRequest(IndexRequest original) throws IOException { + /** + * 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()); @@ -1415,6 +1419,9 @@ private IndexRequest copyIndexRequest(IndexRequest original) throws IOException 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()); diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 1d119f0694928..477ef12a5c042 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -32,8 +32,13 @@ public SamplingService(ScriptService scriptService, ClusterService clusterServic this.clusterService = clusterService; } + /** + * 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 sourceAsMap; try { sourceAsMap = indexRequest.sourceAsMap(); @@ -52,11 +57,27 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque }); } - public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, IngestDocument ingestDocument) { - maybeSample(projectMetadata, indexRequest, () -> ingestDocument); + /** + * + * @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 indexRequestSupplier, + IngestDocument ingestDocument + ) { + maybeSample(projectMetadata, indexName, indexRequestSupplier, () -> ingestDocument); } - private void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest, Supplier ingestDocumentSupplier) { + private void maybeSample( + ProjectMetadata projectMetadata, + String indexName, + Supplier indexRequest, + Supplier ingestDocumentSupplier + ) { // TODO Sampling logic to go here in the near future } diff --git a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java index 4ff762595d354..395407c897a63 100644 --- a/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/IngestServiceTests.java @@ -3403,7 +3403,7 @@ public void testSampling() { ); verify(listener, times(1)).onResponse(null); // In the case where there is a pipeline, or there is a pipeline failure, there will be an IngestDocument so this verion is called: - verify(samplingService, times(2)).maybeSample(any(), any(), any()); + verify(samplingService, times(2)).maybeSample(any(), any(), any(), any()); // When there is no pipeline, we have no IngestDocument, and the maybeSample that does not require an IngestDocument is called: verify(samplingService, times(1)).maybeSample(any(), any()); } From 6e524bfedfcd9ce4964f29095a3ea037ea696c2f Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Thu, 18 Sep 2025 13:20:55 -0500 Subject: [PATCH 34/37] fixing bad merge --- .../main/java/org/elasticsearch/node/NodeConstruction.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java index ad4cccac30c12..e29588ccb257f 100644 --- a/server/src/main/java/org/elasticsearch/node/NodeConstruction.java +++ b/server/src/main/java/org/elasticsearch/node/NodeConstruction.java @@ -721,10 +721,6 @@ private void construct( modules.bindToInstance(SamplingService.class, samplingService); clusterService.addListener(samplingService); - SamplingService samplingService = new SamplingService(scriptService, clusterService); - modules.bindToInstance(SamplingService.class, samplingService); - clusterService.addListener(samplingService); - FailureStoreMetrics failureStoreMetrics = new FailureStoreMetrics(telemetryProvider.getMeterRegistry()); final IngestService ingestService = new IngestService( clusterService, From bc68e3d318ad15431097fec3c01504c82f9ff46a Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 19 Sep 2025 08:39:29 -0500 Subject: [PATCH 35/37] getting the prototype to work again --- .../src/main/java/org/elasticsearch/ingest/IngestService.java | 2 +- .../src/main/java/org/elasticsearch/ingest/SamplingService.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 2265a0343576c..918407fff741b 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -1390,7 +1390,7 @@ private void attemptToSampleData( * 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(), () -> { + samplingService.maybeSample(projectMetadata, originalDocumentMetadata.getIndex(), () -> { IndexRequest original = copyIndexRequestForSampling(indexRequest); updateIndexRequestMetadata(original, originalDocumentMetadata); return original; diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 43cb257d762e0..1350db9581222 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -394,7 +394,7 @@ public void maybeSample( } public boolean atLeastOneSampleConfigured() { - return false; // TODO Return true if there is at least one sample in the cluster state + return true; // TODO Return true if there is at least one sample in the cluster state } @Override From 86171a8812fc1d7c2b24aaeb4234b25101cf434e Mon Sep 17 00:00:00 2001 From: Keith Massey Date: Fri, 19 Sep 2025 10:21:07 -0500 Subject: [PATCH 36/37] avoiding storing IndexRequest --- .../elasticsearch/ingest/IngestService.java | 38 +------------ .../elasticsearch/ingest/SamplingService.java | 57 ++++++++++++------- .../elasticsearch/sample/GetSampleAction.java | 23 +++++--- .../sample/TransportGetSampleAction.java | 3 +- 4 files changed, 54 insertions(+), 67 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/ingest/IngestService.java b/server/src/main/java/org/elasticsearch/ingest/IngestService.java index 918407fff741b..d0726276fda73 100644 --- a/server/src/main/java/org/elasticsearch/ingest/IngestService.java +++ b/server/src/main/java/org/elasticsearch/ingest/IngestService.java @@ -52,7 +52,6 @@ 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; @@ -1390,45 +1389,10 @@ private void attemptToSampleData( * 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, originalDocumentMetadata.getIndex(), () -> { - IndexRequest original = copyIndexRequestForSampling(indexRequest); - updateIndexRequestMetadata(original, originalDocumentMetadata); - return original; - }, ingestDocument); - + samplingService.maybeSample(projectMetadata, originalDocumentMetadata.getIndex(), indexRequest, 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, diff --git a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java index 1350db9581222..1bd170f73503d 100644 --- a/server/src/main/java/org/elasticsearch/ingest/SamplingService.java +++ b/server/src/main/java/org/elasticsearch/ingest/SamplingService.java @@ -26,7 +26,6 @@ import org.elasticsearch.cluster.service.MasterServiceTaskQueue; import org.elasticsearch.common.Priority; import org.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.component.Lifecycle; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; @@ -190,7 +189,7 @@ public void deleteSampleConfiguration(ProjectId projectId, String index) { * @param indexRequest The raw request to potentially sample */ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexRequest) { - maybeSample(projectMetadata, indexRequest.index(), () -> indexRequest, () -> { + maybeSample(projectMetadata, indexRequest.index(), indexRequest, () -> { Map sourceAsMap; try { sourceAsMap = indexRequest.sourceAsMap(); @@ -212,7 +211,7 @@ public void maybeSample(ProjectMetadata projectMetadata, IndexRequest indexReque private void maybeSample( ProjectMetadata projectMetadata, String indexName, - Supplier indexRequestSupplier, + IndexRequest indexRequest, Supplier ingestDocumentSupplier ) { long startTime = relativeNanoTimeSupplier.getAsLong(); @@ -272,12 +271,8 @@ private void maybeSample( sampleInfo.stats )) { stats.timeEvaluatingCondition.add((relativeNanoTimeSupplier.getAsLong() - conditionStartTime)); - IndexRequest indexRequest = indexRequestSupplier.get(); - indexRequest.incRef(); - if (indexRequest.source() instanceof ReleasableBytesReference releaseableSource) { - releaseableSource.incRef(); - } - sampleInfo.getSamples().add(indexRequest); + Sample sample = getSampleForIndexRequest(projectId, indexName, indexRequest); + sampleInfo.getSamples().add(sample); stats.samples.increment(); logger.info("Sampling " + indexRequest); } else { @@ -303,7 +298,7 @@ private void maybeSample( // checkTTLs(); // TODO make this happen less often? } - public List getSamples(ProjectId projectId, String index) { + public List getSamples(ProjectId projectId, String index) { SoftReference sampleInfoReference = samples.get(new ProjectIndex(projectId, index)); SampleInfo sampleInfo = sampleInfoReference.get(); return sampleInfo == null ? List.of() : sampleInfo.getSamples(); @@ -381,16 +376,11 @@ private boolean isClusterServiceStoppedOrClosed() { /** * * @param projectMetadata Used to get the sampling configuration - * @param indexRequestSupplier A supplier for the raw request to potentially sample + * @param indexRequest 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 indexRequestSupplier, - IngestDocument ingestDocument - ) { - maybeSample(projectMetadata, indexName, indexRequestSupplier, () -> ingestDocument); + public void maybeSample(ProjectMetadata projectMetadata, String indexName, IndexRequest indexRequest, IngestDocument ingestDocument) { + maybeSample(projectMetadata, indexName, indexRequest, () -> ingestDocument); } public boolean atLeastOneSampleConfigured() { @@ -456,7 +446,7 @@ public void triggered(SchedulerEngine.Event event) { } private static final class SampleInfo { - private final List samples; + private final List samples; private final SampleStats stats; private final long expiration; private volatile Script script; @@ -469,7 +459,7 @@ private static final class SampleInfo { this.expiration = (timeToLive == null ? TimeValue.timeValueDays(5).nanos() : timeToLive.nanos()) + relativeNowNanos; } - public List getSamples() { + public List getSamples() { return samples; } @@ -636,4 +626,31 @@ static class DeleteSampleConfigTask extends AckedBatchedClusterStateUpdateTask { record ProjectIndex(ProjectId projectId, String indexName) {}; + public record Sample(ProjectId projectId, String indexName, byte[] source, XContentType contentType) implements Writeable { + + public Sample(StreamInput in) throws IOException { + this(ProjectId.readFrom(in), in.readString(), in.readByteArray(), in.readEnum(XContentType.class)); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + projectId.writeTo(out); + out.writeString(indexName); + out.writeBytes(source); + out.writeEnum(contentType); + } + } + + private Sample getSampleForIndexRequest(ProjectId projectId, String indexName, IndexRequest indexRequest) { + BytesReference sourceReference = indexRequest.source(); + final byte[] sourceCopy; + if (sourceReference == null) { + sourceCopy = null; + } else { + byte[] source = sourceReference.array(); + sourceCopy = new byte[sourceReference.length()]; + System.arraycopy(source, sourceReference.arrayOffset(), sourceCopy, 0, sourceReference.length()); + } + return new Sample(projectId, indexName, sourceCopy, indexRequest.getContentType()); + } } diff --git a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java index 2a39693a3e7ab..cd7d36b2537ea 100644 --- a/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/GetSampleAction.java @@ -13,7 +13,6 @@ import org.elasticsearch.action.ActionType; import org.elasticsearch.action.FailedNodeException; import org.elasticsearch.action.IndicesRequest; -import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.IndicesOptions; import org.elasticsearch.action.support.nodes.BaseNodeResponse; import org.elasticsearch.action.support.nodes.BaseNodesRequest; @@ -26,6 +25,8 @@ import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.xcontent.ChunkedToXContent; +import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.ingest.SamplingService; import org.elasticsearch.tasks.CancellableTask; import org.elasticsearch.tasks.Task; import org.elasticsearch.tasks.TaskId; @@ -64,7 +65,7 @@ public Response(ClusterName clusterName, List nodes, List getSamples() { + public List getSamples() { return getNodes().stream().map(n -> n.samples).filter(Objects::nonNull).flatMap(Collection::stream).limit(maxSize).toList(); } @@ -83,8 +84,14 @@ public Iterator toXContentChunked(ToXContent.Params params return Iterators.concat( chunk((builder, p) -> builder.startObject().startArray("samples")), Iterators.flatMap(getSamples().iterator(), sample -> single((builder, params1) -> { - Map source = sample.sourceAsMap(); - builder.value(source); + Map sourceAsMap = XContentHelper.convertToMap( + sample.contentType().xContent(), + sample.source(), + 0, + sample.source().length, + false + ); + builder.value(sourceAsMap); return builder; })), chunk((builder, p) -> builder.endArray().endObject()) @@ -107,19 +114,19 @@ public int hashCode() { } public static class NodeResponse extends BaseNodeResponse { - private final List samples; + private final List samples; protected NodeResponse(StreamInput in) throws IOException { super(in); - samples = in.readCollectionAsList(IndexRequest::new); + samples = in.readCollectionAsList(SamplingService.Sample::new); } - protected NodeResponse(DiscoveryNode node, List samples) { + protected NodeResponse(DiscoveryNode node, List samples) { super(node); this.samples = samples; } - public List getSamples() { + public List getSamples() { return samples; } diff --git a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java index 033e7b3c1a58a..580e0df7d41fb 100644 --- a/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java +++ b/server/src/main/java/org/elasticsearch/sample/TransportGetSampleAction.java @@ -10,7 +10,6 @@ package org.elasticsearch.sample; import org.elasticsearch.action.FailedNodeException; -import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.support.ActionFilters; import org.elasticsearch.action.support.nodes.TransportNodesAction; import org.elasticsearch.cluster.metadata.ProjectId; @@ -78,7 +77,7 @@ protected NodeResponse newNodeResponse(StreamInput in, DiscoveryNode node) throw protected NodeResponse nodeOperation(NodeRequest request, Task task) { ProjectId projectId = request.getProjectId(); String index = request.getIndex(); - List samples = samplingService.getSamples(projectId, index); + List samples = samplingService.getSamples(projectId, index); return new NodeResponse(transportService.getLocalNode(), samples == null ? List.of() : samples); } } From dc6ff3fc92713867e3e8f8eb5659fb9edaf2d8ce Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Thu, 2 Oct 2025 07:46:25 +0000 Subject: [PATCH 37/37] [CI] Update transport version definitions --- server/src/main/resources/transport/upper_bounds/8.18.csv | 2 +- server/src/main/resources/transport/upper_bounds/8.19.csv | 2 +- server/src/main/resources/transport/upper_bounds/9.0.csv | 2 +- server/src/main/resources/transport/upper_bounds/9.1.csv | 2 +- server/src/main/resources/transport/upper_bounds/9.2.csv | 2 +- server/src/main/resources/transport/upper_bounds/9.3.csv | 1 + 6 files changed, 6 insertions(+), 5 deletions(-) create mode 100644 server/src/main/resources/transport/upper_bounds/9.3.csv diff --git a/server/src/main/resources/transport/upper_bounds/8.18.csv b/server/src/main/resources/transport/upper_bounds/8.18.csv index 4eb5140004ea6..266bfbbd3bf78 100644 --- a/server/src/main/resources/transport/upper_bounds/8.18.csv +++ b/server/src/main/resources/transport/upper_bounds/8.18.csv @@ -1 +1 @@ -initial_elasticsearch_8_18_6,8840008 +transform_check_for_dangling_tasks,8840011 diff --git a/server/src/main/resources/transport/upper_bounds/8.19.csv b/server/src/main/resources/transport/upper_bounds/8.19.csv index 476468b203875..3600b3f8c633a 100644 --- a/server/src/main/resources/transport/upper_bounds/8.19.csv +++ b/server/src/main/resources/transport/upper_bounds/8.19.csv @@ -1 +1 @@ -initial_elasticsearch_8_19_3,8841067 +transform_check_for_dangling_tasks,8841070 diff --git a/server/src/main/resources/transport/upper_bounds/9.0.csv b/server/src/main/resources/transport/upper_bounds/9.0.csv index f8f50cc6d7839..c11e6837bb813 100644 --- a/server/src/main/resources/transport/upper_bounds/9.0.csv +++ b/server/src/main/resources/transport/upper_bounds/9.0.csv @@ -1 +1 @@ -initial_elasticsearch_9_0_6,9000015 +transform_check_for_dangling_tasks,9000018 diff --git a/server/src/main/resources/transport/upper_bounds/9.1.csv b/server/src/main/resources/transport/upper_bounds/9.1.csv index 5a65f2e578156..80b97d85f7511 100644 --- a/server/src/main/resources/transport/upper_bounds/9.1.csv +++ b/server/src/main/resources/transport/upper_bounds/9.1.csv @@ -1 +1 @@ -initial_elasticsearch_9_1_4,9112007 +transform_check_for_dangling_tasks,9112009 diff --git a/server/src/main/resources/transport/upper_bounds/9.2.csv b/server/src/main/resources/transport/upper_bounds/9.2.csv index e24f914a1d1ca..2147eab66c207 100644 --- a/server/src/main/resources/transport/upper_bounds/9.2.csv +++ b/server/src/main/resources/transport/upper_bounds/9.2.csv @@ -1 +1 @@ -ml_inference_endpoint_cache,9157000 +initial_9.2.0,9185000 diff --git a/server/src/main/resources/transport/upper_bounds/9.3.csv b/server/src/main/resources/transport/upper_bounds/9.3.csv new file mode 100644 index 0000000000000..2147eab66c207 --- /dev/null +++ b/server/src/main/resources/transport/upper_bounds/9.3.csv @@ -0,0 +1 @@ +initial_9.2.0,9185000