Enable a custom ReplicationTracker override per index (#92353)

Necessary to disable replication via a plugin for stateless.

Author: original-brownbear

server/src/internalClusterTest/java/org/elasticsearch/index/shard/IndexShardIT.java

@@ -51,6 +51,7 @@
 import org.elasticsearch.index.engine.NoOpEngine;
 import org.elasticsearch.index.flush.FlushStats;
 import org.elasticsearch.index.mapper.SourceToParse;
+import org.elasticsearch.index.seqno.ReplicationTracker;
 import org.elasticsearch.index.seqno.RetentionLeaseSyncer;
 import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.translog.TestTranslog;
@@ -661,7 +662,8 @@ public static final IndexShard newIndexShard(
             cbs,
             IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER,
             System::nanoTime,
-            null
+            null,
+            ReplicationTracker.DEFAULT_FACTORY
         );
     }
 

server/src/main/java/org/elasticsearch/index/IndexModule.java

@@ -44,6 +44,7 @@
 import org.elasticsearch.index.mapper.IdFieldMapper;
 import org.elasticsearch.index.mapper.MapperRegistry;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.seqno.ReplicationTracker;
 import org.elasticsearch.index.shard.IndexEventListener;
 import org.elasticsearch.index.shard.IndexingOperationListener;
 import org.elasticsearch.index.shard.SearchOperationListener;
@@ -167,6 +168,8 @@ public interface DirectoryWrapper extends CheckedFunction<Directory, Directory,
     private final Map<String, IndexStorePlugin.RecoveryStateFactory> recoveryStateFactories;
     private final SetOnce<Engine.IndexCommitListener> indexCommitListener = new SetOnce<>();
 
+    private final SetOnce<ReplicationTracker.Factory> replicationTrackerFactory = new SetOnce<>();
+
     /**
      * Construct the index module for the index with the specified index settings. The index module contains extension points for plugins
      * via {@link org.elasticsearch.plugins.Plugin#onIndexModule(IndexModule)}.
@@ -376,6 +379,11 @@ public void setIndexCommitListener(Engine.IndexCommitListener listener) {
         this.indexCommitListener.set(Objects.requireNonNull(listener));
     }
 
+    public void setReplicationTrackerFactory(ReplicationTracker.Factory factory) {
+        ensureNotFrozen();
+        this.replicationTrackerFactory.set(factory);
+    }
+
     IndexEventListener freeze() { // pkg private for testing
         if (this.frozen.compareAndSet(false, true)) {
             return new CompositeIndexEventListener(indexSettings, indexEventListeners);
@@ -524,7 +532,8 @@ public IndexService newIndexService(
                 recoveryStateFactory,
                 indexFoldersDeletionListener,
                 snapshotCommitSupplier,
-                indexCommitListener.get()
+                indexCommitListener.get(),
+                Objects.requireNonNullElse(replicationTrackerFactory.get(), ReplicationTracker.DEFAULT_FACTORY)
             );
             success = true;
             return indexService;

server/src/main/java/org/elasticsearch/index/IndexService.java

@@ -54,6 +54,7 @@
 import org.elasticsearch.index.mapper.NodeMappingStats;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.query.SearchIndexNameMatcher;
+import org.elasticsearch.index.seqno.ReplicationTracker;
 import org.elasticsearch.index.seqno.RetentionLeaseSyncer;
 import org.elasticsearch.index.shard.IndexEventListener;
 import org.elasticsearch.index.shard.IndexShard;
@@ -146,6 +147,8 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final Supplier<Sort> indexSortSupplier;
     private final ValuesSourceRegistry valuesSourceRegistry;
 
+    private final ReplicationTracker.Factory replicationTrackerFactory;
+
     public IndexService(
         IndexSettings indexSettings,
         IndexCreationContext indexCreationContext,
@@ -177,7 +180,8 @@ public IndexService(
         IndexStorePlugin.RecoveryStateFactory recoveryStateFactory,
         IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
         IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier,
-        Engine.IndexCommitListener indexCommitListener
+        Engine.IndexCommitListener indexCommitListener,
+        ReplicationTracker.Factory replicationTrackerFactory
     ) {
         super(indexSettings);
         this.allowExpensiveQueries = allowExpensiveQueries;
@@ -252,6 +256,7 @@ public IndexService(
             this.globalCheckpointTask = new AsyncGlobalCheckpointTask(this);
             this.retentionLeaseSyncTask = new AsyncRetentionLeaseSyncTask(this);
         }
+        this.replicationTrackerFactory = replicationTrackerFactory;
         updateFsyncTaskIfNecessary();
     }
 
@@ -520,7 +525,8 @@ public synchronized IndexShard createShard(
                 circuitBreakerService,
                 snapshotCommitSupplier,
                 System::nanoTime,
-                indexCommitListener
+                indexCommitListener,
+                replicationTrackerFactory
             );
             eventListener.indexShardStateChanged(indexShard, null, indexShard.state(), "shard created");
             eventListener.afterIndexShardCreated(indexShard);

server/src/main/java/org/elasticsearch/index/seqno/ReplicationTracker.java

@@ -52,6 +52,8 @@
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
 
+import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO;
+
 /**
  * This class is responsible for tracking the replication group with its progress and safety markers (local and global checkpoints).
  *
@@ -65,6 +67,28 @@
  */
 public class ReplicationTracker extends AbstractIndexShardComponent implements LongSupplier {
 
+    public static final ReplicationTracker.Factory DEFAULT_FACTORY = (
+        shardId,
+        allocationId,
+        indexSettings,
+        operationPrimaryTerm,
+        onGlobalCheckpointUpdated,
+        currentTimeMillisSupplier,
+        onSyncRetentionLeases,
+        safeCommitInfoSupplier,
+        onReplicationGroupUpdated) -> new ReplicationTracker(
+            shardId,
+            allocationId,
+            indexSettings,
+            operationPrimaryTerm,
+            UNASSIGNED_SEQ_NO,
+            onGlobalCheckpointUpdated,
+            currentTimeMillisSupplier,
+            onSyncRetentionLeases,
+            safeCommitInfoSupplier,
+            onReplicationGroupUpdated
+        );
+
     /**
      * The allocation ID for the shard to which this tracker is a component of.
      */
@@ -1627,4 +1651,22 @@ public int hashCode() {
             return result;
         }
     }
+
+    /**
+     * Factory interface used by {@link org.elasticsearch.index.IndexModule#setReplicationTrackerFactory(Factory)} to enable custom
+     * overrides of this class.
+     */
+    public interface Factory {
+        ReplicationTracker create(
+            ShardId shardId,
+            String allocationId,
+            IndexSettings indexSettings,
+            long operationPrimaryTerm,
+            LongConsumer onGlobalCheckpointUpdated,
+            LongSupplier currentTimeMillisSupplier,
+            BiConsumer<RetentionLeases, ActionListener<ReplicationResponse>> onSyncRetentionLeases,
+            Supplier<SafeCommitInfo> safeCommitInfoSupplier,
+            Consumer<ReplicationGroup> onReplicationGroupUpdated
+        );
+    }
 }

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -313,7 +313,8 @@ public IndexShard(
         final CircuitBreakerService circuitBreakerService,
         final IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier,
         final LongSupplier relativeTimeInNanosSupplier,
-        final Engine.IndexCommitListener indexCommitListener
+        final Engine.IndexCommitListener indexCommitListener,
+        final ReplicationTracker.Factory replicationTrackerFactory
     ) throws IOException {
         super(shardRouting.shardId(), indexSettings);
         assert shardRouting.initializing();
@@ -361,12 +362,11 @@ public IndexShard(
         this.pendingPrimaryTerm = primaryTerm;
         this.globalCheckpointListeners = new GlobalCheckpointListeners(shardId, threadPool.scheduler(), logger);
         this.pendingReplicationActions = new PendingReplicationActions(shardId, threadPool);
-        this.replicationTracker = new ReplicationTracker(
+        this.replicationTracker = replicationTrackerFactory.create(
             shardId,
             aId,
             indexSettings,
             primaryTerm,
-            UNASSIGNED_SEQ_NO,
             globalCheckpointListeners::globalCheckpointUpdated,
             threadPool::absoluteTimeInMillis,
             (retentionLeases, listener) -> retentionLeaseSyncer.sync(shardId, aId, getPendingPrimaryTerm(), retentionLeases, listener),
@@ -3774,7 +3774,7 @@ EngineFactory getEngineFactory() {
     }
 
     // for tests
-    ReplicationTracker getReplicationTracker() {
+    public ReplicationTracker getReplicationTracker() {
         return replicationTracker;
     }
 

server/src/test/java/org/elasticsearch/index/IndexModuleTests.java

@@ -24,8 +24,10 @@
 import org.apache.lucene.tests.index.AssertingDirectoryReader;
 import org.apache.lucene.util.SetOnce.AlreadySetException;
 import org.elasticsearch.Version;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.indices.flush.FlushRequest;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.replication.ReplicationResponse;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -58,14 +60,18 @@
 import org.elasticsearch.index.engine.EngineTestCase;
 import org.elasticsearch.index.engine.InternalEngine;
 import org.elasticsearch.index.engine.InternalEngineFactory;
+import org.elasticsearch.index.engine.SafeCommitInfo;
 import org.elasticsearch.index.fielddata.IndexFieldDataCache;
 import org.elasticsearch.index.mapper.MapperRegistry;
 import org.elasticsearch.index.mapper.ParsedDocument;
 import org.elasticsearch.index.mapper.Uid;
+import org.elasticsearch.index.seqno.ReplicationTracker;
 import org.elasticsearch.index.seqno.RetentionLeaseSyncer;
+import org.elasticsearch.index.seqno.RetentionLeases;
 import org.elasticsearch.index.shard.IndexEventListener;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.IndexingOperationListener;
+import org.elasticsearch.index.shard.ReplicationGroup;
 import org.elasticsearch.index.shard.SearchOperationListener;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.shard.ShardPath;
@@ -106,10 +112,16 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.LongConsumer;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.index.IndexService.IndexCreationContext.CREATE_INDEX;
+import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_SEQ_NO;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
@@ -432,6 +444,7 @@ public void testFrozen() {
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.forceQueryCacheProvider(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setDirectoryWrapper(null)).getMessage());
         assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setIndexCommitListener(null)).getMessage());
+        assertEquals(msg, expectThrows(IllegalStateException.class, () -> module.setReplicationTrackerFactory(null)).getMessage());
     }
 
     public void testSetupUnknownSimilarity() {
@@ -718,6 +731,79 @@ public void onIndexCommitDelete(ShardId shardId, IndexCommit deletedCommit) {
         }
     }
 
+    public void testCustomReplicationTrackerFactory() throws IOException {
+        IndexModule module = new IndexModule(
+            indexSettings,
+            emptyAnalysisRegistry,
+            InternalEngine::new,
+            Collections.emptyMap(),
+            () -> true,
+            indexNameExpressionResolver,
+            Collections.emptyMap()
+        );
+
+        class CustomReplicationTracker extends ReplicationTracker {
+            CustomReplicationTracker(
+                ShardId shardId,
+                String allocationId,
+                IndexSettings indexSettings,
+                long operationPrimaryTerm,
+                long globalCheckpoint,
+                LongConsumer onGlobalCheckpointUpdated,
+                LongSupplier currentTimeMillisSupplier,
+                BiConsumer<RetentionLeases, ActionListener<ReplicationResponse>> onSyncRetentionLeases,
+                Supplier<SafeCommitInfo> safeCommitInfoSupplier,
+                Consumer<ReplicationGroup> onReplicationGroupUpdated
+            ) {
+                super(
+                    shardId,
+                    allocationId,
+                    indexSettings,
+                    operationPrimaryTerm,
+                    globalCheckpoint,
+                    onGlobalCheckpointUpdated,
+                    currentTimeMillisSupplier,
+                    onSyncRetentionLeases,
+                    safeCommitInfoSupplier,
+                    onReplicationGroupUpdated
+                );
+            }
+        }
+        module.setReplicationTrackerFactory(
+            (
+                shardId,
+                allocationId,
+                indexSettings,
+                operationPrimaryTerm,
+                onGlobalCheckpointUpdated,
+                currentTimeMillisSupplier,
+                onSyncRetentionLeases,
+                safeCommitInfoSupplier,
+                onReplicationGroupUpdated) -> new CustomReplicationTracker(
+                    shardId,
+                    allocationId,
+                    indexSettings,
+                    operationPrimaryTerm,
+                    UNASSIGNED_SEQ_NO,
+                    onGlobalCheckpointUpdated,
+                    currentTimeMillisSupplier,
+                    onSyncRetentionLeases,
+                    safeCommitInfoSupplier,
+                    onReplicationGroupUpdated
+                )
+        );
+        final IndexService indexService = newIndexService(module);
+        ShardId shardId = new ShardId("index", UUIDs.randomBase64UUID(random()), 0);
+        ShardRouting shardRouting = ShardRouting.newUnassigned(
+            shardId,
+            true,
+            RecoverySource.EmptyStoreRecoverySource.INSTANCE,
+            new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, null)
+        ).initialize("_node_id", null, -1);
+        IndexShard indexShard = indexService.createShard(shardRouting, s -> {}, RetentionLeaseSyncer.EMPTY);
+        assertThat(indexShard.getReplicationTracker(), instanceOf(CustomReplicationTracker.class));
+    }
+
     private ShardRouting createInitializedShardRouting() {
         ShardRouting shard = ShardRouting.newUnassigned(
             new ShardId("test", "_na_", 0),

test/framework/src/main/java/org/elasticsearch/index/shard/IndexShardTestCase.java

@@ -504,7 +504,8 @@ protected IndexShard newShard(
                 breakerService,
                 IndexModule.DEFAULT_SNAPSHOT_COMMIT_SUPPLIER,
                 relativeTimeSupplier,
-                null
+                null,
+                ReplicationTracker.DEFAULT_FACTORY
             );
             indexShard.addShardFailureCallback(DEFAULT_SHARD_FAILURE_HANDLER);
             success = true;