Implementing time to live for random sampling (#136644)

Author: masseyke

server/src/internalClusterTest/java/org/elasticsearch/ingest/SamplingServiceIT.java

@@ -0,0 +1,77 @@
+/*
+ * 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.ResourceNotFoundException;
+import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
+import org.elasticsearch.action.admin.indices.create.TransportCreateIndexAction;
+import org.elasticsearch.action.admin.indices.sampling.GetSampleAction;
+import org.elasticsearch.action.admin.indices.sampling.PutSampleConfigurationAction;
+import org.elasticsearch.action.admin.indices.sampling.SamplingConfiguration;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.bulk.TransportBulkAction;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.test.ESIntegTestCase;
+import org.junit.After;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.elasticsearch.ingest.SamplingService.RANDOM_SAMPLING_FEATURE_FLAG;
+import static org.elasticsearch.ingest.SamplingService.TTL_POLL_INTERVAL_SETTING;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.hamcrest.Matchers.equalTo;
+
+public class SamplingServiceIT extends ESIntegTestCase {
+    public void testTTL() throws Exception {
+        assumeTrue("Requires sampling feature flag", RANDOM_SAMPLING_FEATURE_FLAG);
+        assertAcked(
+            clusterAdmin().prepareUpdateSettings(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT)
+                .setPersistentSettings(Map.of(TTL_POLL_INTERVAL_SETTING.getKey(), TimeValue.timeValueSeconds(1)))
+        );
+        String indexName = randomIdentifier();
+        client().execute(TransportCreateIndexAction.TYPE, new CreateIndexRequest(indexName)).actionGet();
+        ensureYellow(indexName);
+        PutSampleConfigurationAction.Request putSampleConfigRequest = new PutSampleConfigurationAction.Request(
+            new SamplingConfiguration(1.0d, 10, null, TimeValue.timeValueSeconds(1), null),
+            TimeValue.THIRTY_SECONDS,
+            TimeValue.THIRTY_SECONDS
+        ).indices(indexName);
+        client().execute(PutSampleConfigurationAction.INSTANCE, putSampleConfigRequest).actionGet();
+        BulkRequest bulkRequest = new BulkRequest();
+        for (int i = 0; i < 20; i++) {
+            IndexRequest indexRequest = new IndexRequest(indexName);
+            indexRequest.source(Map.of("foo", randomBoolean() ? 3L : randomLong(), "bar", randomBoolean()));
+            bulkRequest.add(indexRequest);
+        }
+        BulkResponse bulkResponse = client().execute(TransportBulkAction.TYPE, bulkRequest).actionGet();
+        assertThat(bulkResponse.hasFailures(), equalTo(false));
+        GetSampleAction.Response getSampleResponse = client().execute(GetSampleAction.INSTANCE, new GetSampleAction.Request(indexName))
+            .actionGet();
+        assertThat(getSampleResponse.getSample().size(), equalTo(10));
+        assertBusy(() -> {
+            assertThrows(
+                ResourceNotFoundException.class,
+                () -> client().execute(GetSampleAction.INSTANCE, new GetSampleAction.Request(indexName)).actionGet()
+            );
+        });
+    }
+
+    @After
+    public void cleanup() {
+        Map<String, Object> clearedSettings = new HashMap<>();
+        clearedSettings.put(TTL_POLL_INTERVAL_SETTING.getKey(), null);
+        assertAcked(
+            clusterAdmin().prepareUpdateSettings(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT).setPersistentSettings(clearedSettings)
+        );
+    }
+}

server/src/main/java/org/elasticsearch/common/settings/ClusterSettings.java

@@ -106,6 +106,7 @@
 import org.elasticsearch.indices.recovery.RecoverySettings;
 import org.elasticsearch.indices.store.IndicesStore;
 import org.elasticsearch.ingest.IngestSettings;
+import org.elasticsearch.ingest.SamplingService;
 import org.elasticsearch.monitor.fs.FsHealthService;
 import org.elasticsearch.monitor.fs.FsService;
 import org.elasticsearch.monitor.jvm.JvmGcMonitorService;
@@ -650,6 +651,7 @@ public void apply(Settings value, Settings current, Settings previous) {
         WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_HIGH_UTILIZATION_THRESHOLD_SETTING,
         WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_HIGH_UTILIZATION_DURATION_SETTING,
         WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_QUEUE_LATENCY_THRESHOLD_SETTING,
-        WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_REROUTE_INTERVAL_SETTING
+        WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_REROUTE_INTERVAL_SETTING,
+        SamplingService.TTL_POLL_INTERVAL_SETTING
     );
 }

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.ResourceNotFoundException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.indices.sampling.SamplingConfiguration;
@@ -29,10 +30,17 @@
 import org.elasticsearch.cluster.service.MasterServiceTaskQueue;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.component.AbstractLifecycleComponent;
+import org.elasticsearch.common.component.Lifecycle;
+import org.elasticsearch.common.component.LifecycleListener;
 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.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.FeatureFlag;
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
@@ -53,6 +61,10 @@
 
 import java.io.IOException;
 import java.lang.ref.SoftReference;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -67,13 +79,20 @@
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
-public class SamplingService implements ClusterStateListener {
+public class SamplingService extends AbstractLifecycleComponent implements ClusterStateListener, SchedulerEngine.Listener {
     public static final boolean RANDOM_SAMPLING_FEATURE_FLAG = new FeatureFlag("random_sampling").isEnabled();
+    public static final Setting<TimeValue> TTL_POLL_INTERVAL_SETTING = Setting.timeSetting(
+        "random_sampling.ttl_poll_interval",
+        TimeValue.timeValueMinutes(30),
+        TimeValue.timeValueSeconds(1),
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
     private static final Logger logger = LogManager.getLogger(SamplingService.class);
+    private static final String TTL_JOB_ID = "sampling_ttl";
     private final ScriptService scriptService;
     private final ClusterService clusterService;
     private final ProjectResolver projectResolver;
-    private final LongSupplier relativeMillisTimeSupplier;
     private final LongSupplier statsTimeSupplier = System::nanoTime;
     private final MasterServiceTaskQueue<UpdateSamplingConfigurationTask> updateSamplingConfigurationTaskQueue;
     private final MasterServiceTaskQueue<DeleteSampleConfigurationTask> deleteSamplingConfigurationTaskQueue;
@@ -85,24 +104,41 @@ public class SamplingService implements ClusterStateListener {
         Setting.Property.NodeScope,
         Setting.Property.Dynamic
     );
-
+    private final SetOnce<SchedulerEngine> scheduler = new SetOnce<>();
+    private SchedulerEngine.Job scheduledJob;
+    private volatile TimeValue pollInterval;
+    private final Settings settings;
+    private final Clock clock = Clock.systemUTC();
     /*
      * This Map contains the samples that exist on this node. They are not persisted to disk. They are stored as SoftReferences so that
      * sampling does not contribute to a node running out of memory. The idea is that access to samples is desirable, but not critical. We
      * make a best effort to keep them around, but do not worry about the complexity or cost of making them durable.
      */
     private final Map<ProjectIndex, SoftReference<SampleInfo>> samples = new ConcurrentHashMap<>();
 
-    public SamplingService(
+    /*
+     * This creates a new SamplingService, and configures various listeners on it.
+     */
+    public static SamplingService create(
         ScriptService scriptService,
         ClusterService clusterService,
         ProjectResolver projectResolver,
-        LongSupplier relativeMillisTimeSupplier
+        Settings settings
+    ) {
+        SamplingService samplingService = new SamplingService(scriptService, clusterService, projectResolver, settings);
+        samplingService.configureListeners();
+        return samplingService;
+    }
+
+    private SamplingService(
+        ScriptService scriptService,
+        ClusterService clusterService,
+        ProjectResolver projectResolver,
+        Settings settings
     ) {
         this.scriptService = scriptService;
         this.clusterService = clusterService;
         this.projectResolver = projectResolver;
-        this.relativeMillisTimeSupplier = relativeMillisTimeSupplier;
         this.updateSamplingConfigurationTaskQueue = clusterService.createTaskQueue(
             "update-sampling-configuration",
             Priority.NORMAL,
@@ -113,6 +149,24 @@ public SamplingService(
             Priority.NORMAL,
             new DeleteSampleConfigurationExecutor()
         );
+        this.settings = settings;
+        this.pollInterval = TTL_POLL_INTERVAL_SETTING.get(settings);
+    }
+
+    private void configureListeners() {
+        ClusterSettings clusterSettings = clusterService.getClusterSettings();
+        clusterSettings.addSettingsUpdateConsumer(TTL_POLL_INTERVAL_SETTING, (v) -> {
+            pollInterval = v;
+            if (clusterService.state().nodes().isLocalNodeElectedMaster()) {
+                maybeScheduleJob();
+            }
+        });
+        this.addLifecycleListener(new LifecycleListener() {
+            @Override
+            public void afterStop() {
+                cancelJob();
+            }
+        });
     }
 
     /**
@@ -172,11 +226,7 @@ private void maybeSample(
         }
         SoftReference<SampleInfo> sampleInfoReference = samples.compute(
             new ProjectIndex(projectMetadata.id(), indexName),
-            (k, v) -> v == null || v.get() == null
-                ? new SoftReference<>(
-                    new SampleInfo(samplingConfig.maxSamples(), samplingConfig.timeToLive(), relativeMillisTimeSupplier.getAsLong())
-                )
-                : v
+            (k, v) -> v == null || v.get() == null ? new SoftReference<>(new SampleInfo(samplingConfig.maxSamples())) : v
         );
         SampleInfo sampleInfo = sampleInfoReference.get();
         if (sampleInfo == null) {
@@ -345,11 +395,35 @@ public void deleteSampleConfiguration(
         );
     }
 
+    /*
+     * This version is meant to be used by background processes, not user requests.
+     */
+    private void deleteSampleConfiguration(ProjectId projectId, String index) {
+        deleteSampleConfiguration(projectId, index, TimeValue.MAX_VALUE, TimeValue.MAX_VALUE, ActionListener.wrap(response -> {
+            if (response.isAcknowledged()) {
+                logger.debug("Deleted sampling configuration for {}", index);
+            } else {
+                logger.warn("Deletion of sampling configuration for {} not acknowledged", index);
+            }
+        }, e -> logger.warn("Failed to delete sample configuration for " + index, e)));
+    }
+
     @Override
     public void clusterChanged(ClusterChangedEvent event) {
         if (RANDOM_SAMPLING_FEATURE_FLAG == false) {
             return;
         }
+        final boolean isMaster = event.localNodeMaster();
+        final boolean wasMaster = event.previousState().nodes().isLocalNodeElectedMaster();
+        if (wasMaster != isMaster) {
+            if (isMaster) {
+                // we weren't the master, and now we are
+                maybeScheduleJob();
+            } else {
+                // we were the master, and now we aren't
+                cancelJob();
+            }
+        }
         if (samples.isEmpty()) {
             return;
         }
@@ -407,6 +481,31 @@ public void clusterChanged(ClusterChangedEvent event) {
         }
     }
 
+    private void maybeScheduleJob() {
+        if (isClusterServiceStoppedOrClosed()) {
+            // don't create scheduler if the node is shutting down
+            return;
+        }
+        if (scheduler.get() == null) {
+            scheduler.set(new SchedulerEngine(settings, clock));
+            scheduler.get().register(this);
+        }
+        scheduledJob = new SchedulerEngine.Job(TTL_JOB_ID, new TimeValueSchedule(pollInterval));
+        scheduler.get().add(scheduledJob);
+    }
+
+    private void cancelJob() {
+        if (scheduler.get() != null) {
+            scheduler.get().remove(TTL_JOB_ID);
+            scheduledJob = null;
+        }
+    }
+
+    private boolean isClusterServiceStoppedOrClosed() {
+        final Lifecycle.State state = clusterService.lifecycleState();
+        return state == Lifecycle.State.STOPPED || state == Lifecycle.State.CLOSED;
+    }
+
     private boolean evaluateCondition(
         Supplier<IngestDocument> ingestDocumentSupplier,
         Script script,
@@ -457,6 +556,53 @@ private static boolean checkMaxConfigLimitBreached(ProjectId projectId, String i
         return false;
     }
 
+    @Override
+    public void triggered(SchedulerEngine.Event event) {
+        logger.debug("job triggered: {}, {}, {}", event.jobName(), event.scheduledTime(), event.triggeredTime());
+        checkTTLs();
+    }
+
+    @Override
+    protected void doStart() {}
+
+    @Override
+    protected void doStop() {
+        clusterService.removeListener(this);
+        logger.debug("Sampling service is stopping.");
+    }
+
+    @Override
+    protected void doClose() throws IOException {
+        logger.debug("Sampling service is closing.");
+        SchedulerEngine engine = scheduler.get();
+        if (engine != null) {
+            engine.stop();
+        }
+    }
+
+    private void checkTTLs() {
+        long now = clock.instant().toEpochMilli();
+        for (ProjectMetadata projectMetadata : clusterService.state().metadata().projects().values()) {
+            SamplingMetadata samplingMetadata = projectMetadata.custom(SamplingMetadata.TYPE);
+            if (samplingMetadata != null) {
+                for (Map.Entry<String, SamplingConfiguration> entry : samplingMetadata.getIndexToSamplingConfigMap().entrySet()) {
+                    SamplingConfiguration samplingConfiguration = entry.getValue();
+                    if (samplingConfiguration.creationTime() + samplingConfiguration.timeToLive().millis() < now) {
+                        String indexName = entry.getKey();
+                        logger.debug(
+                            "Deleting configuration for {} created at {} because it is older than {} now at {}",
+                            indexName,
+                            ZonedDateTime.ofInstant(Instant.ofEpochMilli(samplingConfiguration.creationTime()), ZoneOffset.UTC),
+                            samplingConfiguration.timeToLive(),
+                            ZonedDateTime.ofInstant(Instant.ofEpochMilli(now), ZoneOffset.UTC)
+                        );
+                        deleteSampleConfiguration(projectMetadata.id(), indexName);
+                    }
+                }
+            }
+        }
+    }
+
     /*
      * This represents a raw document as the user sent it to us in an IndexRequest. It only holds onto the information needed for the
      * sampling API, rather than holding all of the fields a user might send in an IndexRequest.
@@ -835,18 +981,14 @@ private static final class SampleInfo {
          */
         private volatile Tuple<Integer, Long> sizeInBytesAtIndex = Tuple.tuple(-1, 0L);
         private final SampleStats stats;
-        private final long expiration;
-        private final TimeValue timeToLive;
         private volatile Script script;
         private volatile IngestConditionalScript.Factory factory;
         private volatile boolean compilationFailed = false;
         private volatile boolean isFull = false;
 
-        SampleInfo(int maxSamples, TimeValue timeToLive, long relativeNowMillis) {
-            this.timeToLive = timeToLive;
+        SampleInfo(int maxSamples) {
             this.rawDocuments = new RawDocument[maxSamples];
             this.stats = new SampleStats();
-            this.expiration = (timeToLive == null ? TimeValue.timeValueDays(5).millis() : timeToLive.millis()) + relativeNowMillis;
         }
 
         /*