Make stale segments cleanup logic depend on segments map size as well (#20976)

Signed-off-by: rayshrey <rayshrey@amazon.com>

Author: rayshrey

CHANGELOG.md

@@ -63,6 +63,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Delegate getMin/getMax methods for ExitableTerms ([#20775](https://github.com/opensearch-project/OpenSearch/pull/20775))
 - Fix terms lookup subquery fetch limit reading from non-existent index setting instead of cluster `max_clause_count` ([#20823](https://github.com/opensearch-project/OpenSearch/pull/20823))
 - Fix array_index_out_of_bounds_exception with wildcard and aggregations ([#20842](https://github.com/opensearch-project/OpenSearch/pull/20842))
+- Fix stale segment cleanup logic for remote store ([#20976](https://github.com/opensearch-project/OpenSearch/pull/20976))
 - Ensure that transient ThreadContext headers with propagators survive restore ([#169373](https://github.com/opensearch-project/OpenSearch/pull/20854))
 - Handle dependencies between analyzers ([#19248](https://github.com/opensearch-project/OpenSearch/pull/19248))
 - Fix `_field_caps` returning empty results and corrupted field names for `disable_objects: true` mappings ([#20800](https://github.com/opensearch-project/OpenSearch/pull/20800))

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

@@ -819,6 +819,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.CONCURRENT_SEGMENT_SEARCH_PARTITION_MIN_SEGMENT_SIZE,
 
                 RemoteStoreSettings.CLUSTER_REMOTE_INDEX_SEGMENT_METADATA_RETENTION_MAX_COUNT_SETTING,
+                RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING,
                 RemoteStoreSettings.CLUSTER_REMOTE_TRANSLOG_BUFFER_INTERVAL_SETTING,
                 RemoteStoreSettings.CLUSTER_REMOTE_TRANSLOG_TRANSFER_TIMEOUT_SETTING,
                 RemoteStoreSettings.CLUSTER_REMOTE_SEGMENT_TRANSFER_TIMEOUT_SETTING,

server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java

@@ -244,7 +244,9 @@ private boolean syncSegments() {
                 // if a new segments_N file is present in local that is not uploaded to remote store yet, it
                 // is considered as a first refresh post commit. A cleanup of stale commit files is triggered.
                 // This is done to avoid delete post each refresh.
-                if (isRefreshAfterCommit()) {
+                // Also trigger cleanup if the uploaded segments map exceeds the configured threshold,
+                // to prevent unbounded memory growth when flushes do not happen.
+                if (isRefreshAfterCommit() || uploadedSegmentsMapExceedsThreshold()) {
                     remoteDirectory.deleteStaleSegmentsAsync(indexShard.getRemoteStoreSettings().getMinRemoteSegmentMetadataFiles());
                 }
 
@@ -449,6 +451,11 @@ private boolean isRefreshAfterCommitSafe() {
         return false;
     }
 
+    private boolean uploadedSegmentsMapExceedsThreshold() {
+        int threshold = indexShard.getRemoteStoreSettings().getUploadedSegmentsCleanupThreshold();
+        return threshold != -1 && remoteDirectory.getSegmentsUploadedToRemoteStoreSize() > threshold;
+    }
+
     void uploadMetadata(Collection<String> localSegmentsPostRefresh, SegmentInfos segmentInfos, ReplicationCheckpoint replicationCheckpoint)
         throws IOException {
         final long maxSeqNo = indexShard.getIndexer().currentOngoingRefreshCheckpoint();

server/src/main/java/org/opensearch/index/store/RemoteSegmentStoreDirectory.java

@@ -893,6 +893,10 @@ public Map<String, UploadedSegmentMetadata> getSegmentsUploadedToRemoteStore() {
         return Collections.unmodifiableMap(this.segmentsUploadedToRemoteStore);
     }
 
+    public int getSegmentsUploadedToRemoteStoreSize() {
+        return segmentsUploadedToRemoteStore.size();
+    }
+
     // Visible for testing
     Set<String> getMetadataFilesToFilterActiveSegments(
         final int lastNMetadataFilesToKeep,

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

@@ -26,6 +26,9 @@
 @PublicApi(since = "2.14.0")
 public class RemoteStoreSettings {
     private static final int MIN_CLUSTER_REMOTE_MAX_TRANSLOG_READERS = 100;
+    private static final int MIN_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD = 100;
+    private static final int MAX_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD = 100000;
+    private static final int DEFAULT_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD = 1000;
 
     /**
      * Used to specify the default translog buffer interval for remote store backed indexes.
@@ -174,6 +177,29 @@ public class RemoteStoreSettings {
         Property.Final
     );
 
+    /**
+     * Controls the threshold for the number of segments uploaded to remote store map.
+     * When the map size exceeds this threshold, stale segment cleanup is triggered even without a flush/commit.
+     * {@code -1} disables threshold-based cleanup.
+     */
+    public static final Setting<Integer> CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING = Setting.intSetting(
+        "cluster.remote_store.uploaded_segments_cleanup_threshold",
+        DEFAULT_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD,
+        -1,
+        v -> {
+            if (v != -1 && (v < MIN_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD || v > MAX_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD)) {
+                throw new IllegalArgumentException(
+                    "Value must be -1 or between "
+                        + MIN_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD
+                        + " and "
+                        + MAX_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD
+                );
+            }
+        },
+        Property.NodeScope,
+        Property.Dynamic
+    );
+
     /**
      * Controls the fixed prefix for the segments path on remote store.
      */
@@ -208,6 +234,7 @@ public class RemoteStoreSettings {
     private static volatile TimeValue pinnedTimestampsLookbackInterval;
     private final String translogPathFixedPrefix;
     private final String segmentsPathFixedPrefix;
+    private volatile int uploadedSegmentsCleanupThreshold;
 
     public RemoteStoreSettings(Settings settings, ClusterSettings clusterSettings) {
         clusterRemoteTranslogBufferInterval = CLUSTER_REMOTE_TRANSLOG_BUFFER_INTERVAL_SETTING.get(settings);
@@ -255,6 +282,12 @@ public RemoteStoreSettings(Settings settings, ClusterSettings clusterSettings) {
 
         translogPathFixedPrefix = CLUSTER_REMOTE_STORE_TRANSLOG_PATH_PREFIX.get(settings);
         segmentsPathFixedPrefix = CLUSTER_REMOTE_STORE_SEGMENTS_PATH_PREFIX.get(settings);
+
+        uploadedSegmentsCleanupThreshold = CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.get(settings);
+        clusterSettings.addSettingsUpdateConsumer(
+            CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING,
+            this::setUploadedSegmentsCleanupThreshold
+        );
     }
 
     public TimeValue getClusterRemoteTranslogBufferInterval() {
@@ -355,4 +388,12 @@ public String getTranslogPathFixedPrefix() {
     public String getSegmentsPathFixedPrefix() {
         return segmentsPathFixedPrefix;
     }
+
+    public int getUploadedSegmentsCleanupThreshold() {
+        return uploadedSegmentsCleanupThreshold;
+    }
+
+    private void setUploadedSegmentsCleanupThreshold(int uploadedSegmentsCleanupThreshold) {
+        this.uploadedSegmentsCleanupThreshold = uploadedSegmentsCleanupThreshold;
+    }
 }

server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java

@@ -906,4 +906,78 @@ public void testRemoteSegmentStoreNotInSync() throws IOException {
         }
     }
 
+    public void testCleanupTriggeredWhenMapExceedsThreshold() throws IOException {
+        int threshold = 10;
+        RemoteSegmentStoreDirectory remoteSegmentStoreDirectory = setupDirectoryWithThreshold(threshold);
+
+        indexAndRefreshWithoutFlush(100);
+
+        int mapSize = remoteSegmentStoreDirectory.getSegmentsUploadedToRemoteStoreSize();
+        assertTrue("Map size should be bounded by threshold cleanup, but was: " + mapSize, mapSize < 100);
+    }
+
+    public void testCleanupNotTriggeredWhenThresholdDisabled() throws IOException {
+        RemoteSegmentStoreDirectory remoteSegmentStoreDirectory = setupDirectoryWithThreshold(-1);
+        int initialMapSize = remoteSegmentStoreDirectory.getSegmentsUploadedToRemoteStoreSize();
+
+        indexAndRefreshWithoutFlush(100);
+
+        int finalMapSize = remoteSegmentStoreDirectory.getSegmentsUploadedToRemoteStoreSize();
+        assertTrue(
+            "Map size should have grown with threshold disabled, initial=" + initialMapSize + " final=" + finalMapSize,
+            finalMapSize > initialMapSize
+        );
+    }
+
+    private RemoteSegmentStoreDirectory setupDirectoryWithThreshold(int threshold) throws IOException {
+        indexShard = newStartedShard(
+            true,
+            Settings.builder()
+                .put(IndexMetadata.SETTING_REMOTE_STORE_ENABLED, true)
+                .put(IndexMetadata.SETTING_REMOTE_SEGMENT_STORE_REPOSITORY, "temp-fs")
+                .put(IndexMetadata.SETTING_REMOTE_TRANSLOG_STORE_REPOSITORY, "temp-fs")
+                .put(SETTING_REPLICATION_TYPE, ReplicationType.SEGMENT)
+                .build(),
+            new InternalEngineFactory()
+        );
+
+        indexDocs(1, 3);
+        indexShard.refresh("test");
+
+        clusterService = ClusterServiceUtils.createClusterService(
+            Settings.EMPTY,
+            new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS),
+            threadPool
+        );
+        remoteStoreStatsTrackerFactory = new RemoteStoreStatsTrackerFactory(clusterService, Settings.EMPTY);
+        remoteStoreStatsTrackerFactory.afterIndexShardCreated(indexShard);
+        RemoteSegmentTransferTracker tracker = remoteStoreStatsTrackerFactory.getRemoteSegmentTransferTracker(indexShard.shardId());
+
+        RemoteStoreSettings mockSettings = mock(RemoteStoreSettings.class);
+        when(mockSettings.getUploadedSegmentsCleanupThreshold()).thenReturn(threshold);
+        when(mockSettings.getMinRemoteSegmentMetadataFiles()).thenReturn(10);
+        when(mockSettings.getClusterRemoteSegmentTransferTimeout()).thenReturn(TimeValue.timeValueMinutes(30));
+
+        IndexShard spyShard = spy(indexShard);
+        when(spyShard.getRemoteStoreSettings()).thenReturn(mockSettings);
+
+        remoteStoreRefreshListener = new RemoteStoreRefreshListener(
+            spyShard,
+            SegmentReplicationCheckpointPublisher.EMPTY,
+            tracker,
+            mockSettings
+        );
+
+        return (RemoteSegmentStoreDirectory) ((FilterDirectory) ((FilterDirectory) indexShard.remoteStore().directory()).getDelegate())
+            .getDelegate();
+    }
+
+    private void indexAndRefreshWithoutFlush(int iterations) throws IOException {
+        for (int i = 0; i < iterations; i++) {
+            indexDocs(10 + (i * 5), 5);
+            indexShard.refresh("test");
+            remoteStoreRefreshListener.afterRefresh(true);
+        }
+    }
+
 }

server/src/test/java/org/opensearch/indices/RemoteStoreSettingsDynamicUpdateTests.java

@@ -127,4 +127,67 @@ public void testDisableMaxRemoteReferencedTranslogFiles() {
         );
         assertEquals(-1, remoteStoreSettings.getMaxRemoteTranslogReaders());
     }
+
+    public void testUploadedSegmentsCleanupThreshold() {
+        // Test default value
+        assertEquals(1000, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        // Test override with valid value
+        clusterSettings.applySettings(
+            Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 5000).build()
+        );
+        assertEquals(5000, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        // Test disable with -1
+        clusterSettings.applySettings(
+            Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), -1).build()
+        );
+        assertEquals(-1, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        // Test value below -1 should fail
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> clusterSettings.applySettings(
+                Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), -5).build()
+            )
+        );
+        assertEquals(-1, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        // Test value below minimum (but not -1) should fail
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> clusterSettings.applySettings(
+                Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 50).build()
+            )
+        );
+
+        // Test value above maximum should fail
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> clusterSettings.applySettings(
+                Settings.builder()
+                    .put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 200000)
+                    .build()
+            )
+        );
+
+        // Test boundary values
+        clusterSettings.applySettings(
+            Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 100).build()
+        );
+        assertEquals(100, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        clusterSettings.applySettings(
+            Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 100000).build()
+        );
+        assertEquals(100000, remoteStoreSettings.getUploadedSegmentsCleanupThreshold());
+
+        // Test 0 should fail (not -1 and below minimum)
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> clusterSettings.applySettings(
+                Settings.builder().put(RemoteStoreSettings.CLUSTER_REMOTE_UPLOADED_SEGMENTS_CLEANUP_THRESHOLD_SETTING.getKey(), 0).build()
+            )
+        );
+    }
 }