Add interface for the Multi format merge flow (#20908)

Signed-off-by: Sagar Darji <darsaga@amazon.com>
Co-authored-by: Sagar Darji <darsaga@amazon.com>

Author: darjisagar7

CHANGELOG.md

@@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add ref_path support for package-based hunspell dictionary loading ([#20840](https://github.com/opensearch-project/OpenSearch/pull/20840))
 - Add support for enabling pluggable data formats, starting with phase-1 of decoupling shard from engine, and introducing basic abstractions ([#20675](https://github.com/opensearch-project/OpenSearch/pull/20675))
 - Add concurrent queue in libs and composite engine sandbox plugin ([#20909](https://github.com/opensearch-project/OpenSearch/pull/20909))
+- Add interface for the Multi format merge flow ([#20908](https://github.com/opensearch-project/OpenSearch/pull/20908))
 
 - Add warmup phase to wait for lag to catch up in pull-based ingestion before serving ([#20526](https://github.com/opensearch-project/OpenSearch/pull/20526))
 - Add a new static method to IndicesOptions API to expose `STRICT_EXPAND_OPEN_HIDDEN_FORBID_CLOSED` index option ([#20980](https://github.com/opensearch-project/OpenSearch/pull/20980))

sandbox/plugins/composite-engine/src/main/java/org/opensearch/composite/CompositeIndexingExecutionEngine.java

@@ -61,6 +61,7 @@ public class CompositeIndexingExecutionEngine implements IndexingExecutionEngine
     private final Set<IndexingExecutionEngine<?, ?>> secondaryEngines;
     private final CompositeDataFormat compositeDataFormat;
     private final LockablePool<CompositeWriter> writerPool;
+    private final AtomicLong writerGenerationCounter;
 
     /**
      * Constructs a CompositeIndexingExecutionEngine by reading index settings to
@@ -113,7 +114,7 @@ public CompositeIndexingExecutionEngine(
         this.compositeDataFormat = new CompositeDataFormat(allFormats);
 
         // Create the writer pool internally, matching the reference code pattern
-        AtomicLong writerGenerationCounter = new AtomicLong(0);
+        writerGenerationCounter = new AtomicLong(0);
         this.writerPool = new LockablePool<>(
             () -> new CompositeWriter(this, writerGenerationCounter.getAndIncrement()),
             LinkedList::new,
@@ -232,6 +233,11 @@ public RefreshResult refresh(RefreshInput refreshInput) throws IOException {
         return new RefreshResult(refreshedSegments);
     }
 
+    @Override
+    public long getNextWriterGeneration() {
+        return writerGenerationCounter.getAndIncrement();
+    }
+
     @Override
     public CompositeDataFormat getDataFormat() {
         return compositeDataFormat;

sandbox/plugins/composite-engine/src/test/java/org/opensearch/composite/CompositeTestHelper.java

@@ -31,6 +31,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * Shared test utilities for composite engine tests.
@@ -134,6 +135,7 @@ public String toString() {
     static class StubIndexingExecutionEngine implements IndexingExecutionEngine<DataFormat, DocumentInput<?>> {
 
         private final DataFormat dataFormat;
+        private final AtomicLong writerGeneration = new AtomicLong(0);
 
         StubIndexingExecutionEngine(DataFormat dataFormat) {
             this.dataFormat = dataFormat;
@@ -162,6 +164,11 @@ public DataFormat getDataFormat() {
         @Override
         public void deleteFiles(Map<String, Collection<String>> filesToDelete) {}
 
+        @Override
+        public long getNextWriterGeneration() {
+            return writerGeneration.getAndIncrement();
+        }
+
         @Override
         public DocumentInput<?> newDocumentInput() {
             return new StubDocumentInput();

server/src/main/java/org/opensearch/index/engine/dataformat/IndexingExecutionEngine.java

@@ -49,6 +49,15 @@ public interface IndexingExecutionEngine<T extends DataFormat, P extends Documen
      */
     RefreshResult refresh(RefreshInput refreshInput) throws IOException;
 
+    /**
+     * Returns the next writer generation number to be used when creating a new writer.
+     * Each writer is associated with a monotonically increasing generation number
+     * that uniquely identifies it within this engine's lifecycle.
+     *
+     * @return the next writer generation number
+     */
+    long getNextWriterGeneration();
+
     /**
      * Returns the data format handled by this engine.
      *

server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java

@@ -0,0 +1,163 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.engine.dataformat.merge;
+
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.common.concurrent.GatedCloseable;
+import org.opensearch.common.logging.Loggers;
+import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.index.engine.dataformat.MergeResult;
+import org.opensearch.index.engine.exec.CatalogSnapshot;
+import org.opensearch.index.engine.exec.Indexer;
+import org.opensearch.index.engine.exec.Segment;
+
+import java.util.ArrayDeque;
+import java.util.Collection;
+import java.util.Deque;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Abstract handler responsible for managing segment merge operations.
+ * <p>
+ * Subclasses define the merge policy by implementing {@link #findMerges()} and
+ * {@link #findForceMerges(int)}, while this base class manages the pending merge
+ * queue and lifecycle callbacks.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public abstract class MergeHandler {
+
+    private final Deque<OneMerge> mergingSegments = new ArrayDeque<>();
+    private final Set<Segment> currentlyMergingSegments = new HashSet<>();
+    private final Indexer indexer;
+    private final Logger logger;
+
+    public MergeHandler(Indexer indexer, ShardId shardId) {
+        this.logger = Loggers.getLogger(getClass(), shardId);
+        this.indexer = indexer;
+    }
+
+    /**
+     * Finds merges that should be executed based on the current segment state.
+     *
+     * @return a collection of merges to execute, or an empty collection if none are needed
+     */
+    public abstract Collection<OneMerge> findMerges();
+
+    /**
+     * Finds merges required to reduce the number of segments to at most {@code maxSegmentCount}.
+     *
+     * @param maxSegmentCount the maximum number of segments allowed after merging
+     * @return a collection of merges to execute
+     */
+    public abstract Collection<OneMerge> findForceMerges(int maxSegmentCount);
+
+    /**
+     * Updates the set of pending merges. Called to refresh the merge queue
+     * when the segment state changes.
+     */
+    public synchronized void updatePendingMerges() {
+        Collection<OneMerge> oneMerges = findMerges();
+        for (OneMerge oneMerge : oneMerges) {
+            boolean isValidMerge = true;
+            for (Segment segment : oneMerge.getSegmentsToMerge()) {
+                if (currentlyMergingSegments.contains(segment)) {
+                    isValidMerge = false;
+                    break;
+                }
+            }
+            if (isValidMerge) {
+                registerMerge(oneMerge);
+            }
+        }
+    }
+
+    /**
+     * Registers a merge to be executed.
+     *
+     * @param merge the merge to register
+     */
+    public synchronized void registerMerge(OneMerge merge) {
+        try (GatedCloseable<CatalogSnapshot> catalogSnapshotReleasableRef = indexer.acquireSnapshot()) {
+            // Validate segments exist in catalog
+            List<Segment> catalogSegments = catalogSnapshotReleasableRef.get().getSegments();
+            for (Segment mergeSegment : merge.getSegmentsToMerge()) {
+                if (!catalogSegments.contains(mergeSegment)) {
+                    return;
+                }
+            }
+        } catch (Exception e) {
+            logger.warn("Failed to acquire snapshots", e);
+            throw new RuntimeException(e);
+        }
+        mergingSegments.add(merge);
+        currentlyMergingSegments.addAll(merge.getSegmentsToMerge());
+        logger.debug(() -> new ParameterizedMessage("Registered merge [{}], mergingSegments: [{}]", merge, mergingSegments));
+    }
+
+    /**
+     * Returns whether there are any pending merges in the queue.
+     *
+     * @return {@code true} if there are pending merges
+     */
+    public synchronized boolean hasPendingMerges() {
+        return !mergingSegments.isEmpty();
+    }
+
+    /**
+     * Retrieves and removes the next pending merge from the queue.
+     *
+     * @return the next merge to execute, or {@code null} if the queue is empty
+     */
+    public synchronized OneMerge getNextMerge() {
+        if (mergingSegments.isEmpty()) {
+            return null;
+        }
+        return mergingSegments.removeFirst();
+    }
+
+    /**
+     * Callback invoked when a merge completes successfully.
+     *
+     * @param oneMerge the merge that finished
+     */
+    public synchronized void onMergeFinished(OneMerge oneMerge) {
+        removeMergingSegments(oneMerge);
+        updatePendingMerges();
+    }
+
+    /**
+     * Callback invoked when a merge fails.
+     *
+     * @param oneMerge the merge that failed
+     */
+    public synchronized void onMergeFailure(OneMerge oneMerge) {
+        removeMergingSegments(oneMerge);
+        logger.warn(() -> new ParameterizedMessage("Merge failed for OneMerge [{}]", oneMerge));
+    }
+
+    /**
+     * Executes the given merge operation.
+     *
+     * @param oneMerge the merge to execute
+     * @return the result of the merge
+     */
+    public abstract MergeResult doMerge(OneMerge oneMerge);
+
+    private synchronized void removeMergingSegments(OneMerge oneMerge) {
+        mergingSegments.remove(oneMerge);
+        oneMerge.getSegmentsToMerge().forEach(currentlyMergingSegments::remove);
+    }
+
+}

server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeScheduler.java

@@ -0,0 +1,133 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.engine.dataformat.merge;
+
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.common.logging.Loggers;
+import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.index.IndexSettings;
+import org.opensearch.index.MergeSchedulerConfig;
+import org.opensearch.index.merge.MergeStats;
+
+/**
+ * Schedules and coordinates segment merge operations for a shard.
+ * <p>
+ * This scheduler delegates merge selection to a {@link MergeHandler} and controls
+ * concurrency via configurable thread and merge count limits sourced from
+ * {@link MergeSchedulerConfig}.
+ *
+ * @opensearch.experimental
+ */
+@ExperimentalApi
+public class MergeScheduler {
+
+    private final Logger logger;
+    private volatile int maxConcurrentMerges;
+    private volatile int maxMergeCount;
+    private final MergeSchedulerConfig mergeSchedulerConfig;
+
+    /** true if we should rate-limit writes for each merge */
+    private boolean doAutoIOThrottle = false;
+
+    /** Initial value for IO write rate limit when doAutoIOThrottle is true */
+    private static final double START_MB_PER_SEC = 20.0;
+
+    /** Current IO writes throttle rate */
+    protected double targetMBPerSec = START_MB_PER_SEC;
+
+    /**
+     * Creates a new merge scheduler.
+     *
+     * @param mergeHandler   the handler that selects and executes merges
+     * @param shardId        the shard this scheduler is associated with
+     * @param indexSettings  the index settings providing merge scheduler configuration
+     */
+    public MergeScheduler(MergeHandler mergeHandler, ShardId shardId, IndexSettings indexSettings) {
+        logger = Loggers.getLogger(getClass(), shardId);
+        this.mergeSchedulerConfig = indexSettings.getMergeSchedulerConfig();
+        refreshConfig();
+    }
+
+    /**
+     * Refreshes the max concurrent merge thread count and max merge count from
+     * the current {@link MergeSchedulerConfig}. No-op if the values have not changed.
+     */
+    public synchronized void refreshConfig() {
+        int newMaxThreadCount = mergeSchedulerConfig.getMaxThreadCount();
+        int newMaxMergeCount = mergeSchedulerConfig.getMaxMergeCount();
+
+        if (newMaxThreadCount == this.maxConcurrentMerges && newMaxMergeCount == this.maxMergeCount) {
+            return;
+        }
+
+        logger.info(
+            () -> new ParameterizedMessage(
+                "Updating from merge scheduler config: maxThreadCount {} -> {}, " + "maxMergeCount {} -> {}",
+                this.maxConcurrentMerges,
+                newMaxThreadCount,
+                this.maxMergeCount,
+                newMaxMergeCount
+            )
+        );
+
+        this.maxConcurrentMerges = newMaxThreadCount;
+        this.maxMergeCount = newMaxMergeCount;
+    }
+
+    /**
+     * Triggers pending merge operations. Merges are selected by the
+     * underlying {@link MergeHandler} and executed up to the configured
+     * concurrency limits.
+     */
+    public void triggerMerges() {
+
+    }
+
+    /**
+     * Forces a merge down to at most {@code maxNumSegment} segments.
+     *
+     * @param maxNumSegment the maximum number of segments after the force merge
+     */
+    public void forceMerge(int maxNumSegment) {
+
+    }
+
+    /**
+     * Turn on dynamic IO throttling, to adaptively rate limit writes bytes/sec to the minimal rate
+     * necessary so merges do not fall behind. By default, this is disabled and writes are not
+     * rate-limited.
+     */
+    public synchronized void enableAutoIOThrottle() {
+        doAutoIOThrottle = true;
+        targetMBPerSec = START_MB_PER_SEC;
+    }
+
+    /**
+     * Returns the currently set per-merge IO writes rate limit, if {@link #enableAutoIOThrottle} was
+     * called, else {@code Double.POSITIVE_INFINITY}.
+     */
+    public synchronized double getIORateLimitMBPerSec() {
+        if (doAutoIOThrottle) {
+            return targetMBPerSec;
+        }
+
+        return Double.POSITIVE_INFINITY;
+    }
+
+    /**
+     * Returns the current merge statistics for this scheduler.
+     *
+     * @return the merge stats
+     */
+    public MergeStats stats() {
+        return new MergeStats();
+    }
+}