[Backport 2.x] Support shard promotion with Segment Replication. (#4135…

…) (#4325)

* Support shard promotion with Segment Replication.

This change adds basic failover support with segment replication.  Once selected, a replica will commit and reopen a writeable engine.

Signed-off-by: Marc Handalian <handalm@amazon.com>

* Add check to ensure a closed shard does not publish checkpoints.

Signed-off-by: Marc Handalian <handalm@amazon.com>

* Clean up in SegmentReplicationIT.

Signed-off-by: Marc Handalian <handalm@amazon.com>

* PR feedback.

Signed-off-by: Marc Handalian <handalm@amazon.com>

* Fix merge conflict.

Signed-off-by: Marc Handalian <handalm@amazon.com>

Signed-off-by: Marc Handalian <handalm@amazon.com>

Signed-off-by: Marc Handalian <handalm@amazon.com>

CHANGELOG.md

@@ -4,6 +4,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased]
 ### Added
 - Github workflow for changelog verification ([#4085](https://github.com/opensearch-project/OpenSearch/pull/4085))
+- Add failover support with Segment Replication enabled. ([#4325](https://github.com/opensearch-project/OpenSearch/pull/4325)
 
 ### Changed
  - Dependency updates (httpcore, mockito, slf4j, httpasyncclient, commons-codec) ([#4308](https://github.com/opensearch-project/OpenSearch/pull/4308))

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationIT.java

@@ -19,7 +19,10 @@
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.routing.IndexShardRoutingTable;
 import org.opensearch.cluster.routing.ShardRouting;
+import org.opensearch.cluster.routing.allocation.command.CancelAllocationCommand;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.Index;
@@ -30,6 +33,7 @@
 import org.opensearch.indices.IndicesService;
 import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.test.BackgroundIndexer;
+import org.opensearch.test.InternalTestCluster;
 import org.opensearch.test.OpenSearchIntegTestCase;
 
 import java.io.IOException;
@@ -73,6 +77,109 @@ protected boolean addMockInternalEngine() {
  return false;
  }
 
+ public void testPrimaryStopped_ReplicaPromoted() throws Exception {
+ final String primary = internalCluster().startNode();
+ createIndex(INDEX_NAME);
+ ensureYellowAndNoInitializingShards(INDEX_NAME);
+ final String replica = internalCluster().startNode();
+ ensureGreen(INDEX_NAME);
+
+ client().prepareIndex(INDEX_NAME).setId("1").setSource("foo", "bar").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
+ refresh(INDEX_NAME);
+
+ waitForReplicaUpdate();
+ assertHitCount(client(primary).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 1);
+ assertHitCount(client(replica).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 1);
+
+ // index another doc but don't refresh, we will ensure this is searchable once replica is promoted.
+ client().prepareIndex(INDEX_NAME).setId("2").setSource("bar", "baz").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
+
+ // stop the primary node - we only have one shard on here.
+ internalCluster().stopRandomNode(InternalTestCluster.nameFilter(primary));
+ ensureYellowAndNoInitializingShards(INDEX_NAME);
+
+ final ShardRouting replicaShardRouting = getShardRoutingForNodeName(replica);
+ assertNotNull(replicaShardRouting);
+ assertTrue(replicaShardRouting + " should be promoted as a primary", replicaShardRouting.primary());
+ assertHitCount(client(replica).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 2);
+
+ // assert we can index into the new primary.
+ client().prepareIndex(INDEX_NAME).setId("3").setSource("bar", "baz").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
+ assertHitCount(client(replica).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 3);
+
+ // start another node, index another doc and replicate.
+ String nodeC = internalCluster().startNode();
+ ensureGreen(INDEX_NAME);
+ client().prepareIndex(INDEX_NAME).setId("4").setSource("baz", "baz").get();
+ refresh(INDEX_NAME);
+ waitForReplicaUpdate();
+ assertHitCount(client(nodeC).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 4);
+ assertHitCount(client(replica).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 4);
+ assertSegmentStats(REPLICA_COUNT);
+ }
+
+ public void testRestartPrimary() throws Exception {
+ final String primary = internalCluster().startNode();
+ createIndex(INDEX_NAME);
+ ensureYellowAndNoInitializingShards(INDEX_NAME);
+ final String replica = internalCluster().startNode();
+ ensureGreen(INDEX_NAME);
+
+ assertEquals(getNodeContainingPrimaryShard().getName(), primary);
+
+ final int initialDocCount = 1;
+ client().prepareIndex(INDEX_NAME).setId("1").setSource("foo", "bar").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
+ refresh(INDEX_NAME);
+
+ waitForReplicaUpdate();
+ assertDocCounts(initialDocCount, replica, primary);
+
+ internalCluster().restartNode(primary);
+ ensureGreen(INDEX_NAME);
+
+ assertEquals(getNodeContainingPrimaryShard().getName(), replica);
+
+ flushAndRefresh(INDEX_NAME);
+ waitForReplicaUpdate();
+
+ assertDocCounts(initialDocCount, replica, primary);
+ assertSegmentStats(REPLICA_COUNT);
+ }
+
+ public void testCancelPrimaryAllocation() throws Exception {
+ // this test cancels allocation on the primary - promoting the new replica and recreating the former primary as a replica.
+ final String primary = internalCluster().startNode();
+ createIndex(INDEX_NAME);
+ ensureYellowAndNoInitializingShards(INDEX_NAME);
+ final String replica = internalCluster().startNode();
+ ensureGreen(INDEX_NAME);
+
+ final int initialDocCount = 1;
+
+ client().prepareIndex(INDEX_NAME).setId("1").setSource("foo", "bar").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
+ refresh(INDEX_NAME);
+
+ waitForReplicaUpdate();
+ assertDocCounts(initialDocCount, replica, primary);
+
+ final IndexShard indexShard = getIndexShard(primary);
+ client().admin()
+ .cluster()
+ .prepareReroute()
+ .add(new CancelAllocationCommand(INDEX_NAME, indexShard.shardId().id(), primary, true))
+ .execute()
+ .actionGet();
+ ensureGreen(INDEX_NAME);
+
+ assertEquals(getNodeContainingPrimaryShard().getName(), replica);
+
+ flushAndRefresh(INDEX_NAME);
+ waitForReplicaUpdate();
+
+ assertDocCounts(initialDocCount, replica, primary);
+ assertSegmentStats(REPLICA_COUNT);
+ }
+
  public void testReplicationAfterPrimaryRefreshAndFlush() throws Exception {
  final String nodeA = internalCluster().startNode();
  final String nodeB = internalCluster().startNode();
@@ -240,9 +347,8 @@ public void testStartReplicaAfterPrimaryIndexesDocs() throws Exception {
  assertHitCount(client(primaryNode).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 3);
  assertHitCount(client(replicaNode).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), 3);
 
- final Index index = resolveIndex(INDEX_NAME);
- IndexShard primaryShard = getIndexShard(index, primaryNode);
- IndexShard replicaShard = getIndexShard(index, replicaNode);
+ IndexShard primaryShard = getIndexShard(primaryNode);
+ IndexShard replicaShard = getIndexShard(replicaNode);
  assertEquals(
  primaryShard.translogStats().estimatedNumberOfOperations(),
  replicaShard.translogStats().estimatedNumberOfOperations()
@@ -351,8 +457,7 @@ private void assertSegmentStats(int numberOfReplicas) throws IOException {
  final ShardRouting replicaShardRouting = shardSegment.getShardRouting();
  ClusterState state = client(internalCluster().getMasterName()).admin().cluster().prepareState().get().getState();
  final DiscoveryNode replicaNode = state.nodes().resolveNode(replicaShardRouting.currentNodeId());
- final Index index = resolveIndex(INDEX_NAME);
- IndexShard indexShard = getIndexShard(index, replicaNode.getName());
+ IndexShard indexShard = getIndexShard(replicaNode.getName());
  final String lastCommitSegmentsFileName = SegmentInfos.getLastCommitSegmentsFileName(indexShard.store().directory());
  // calls to readCommit will fail if a valid commit point and all its segments are not in the store.
  SegmentInfos.readCommit(indexShard.store().directory(), lastCommitSegmentsFileName);
@@ -392,7 +497,8 @@ private void waitForReplicaUpdate() throws Exception {
  });
  }
 
- private IndexShard getIndexShard(Index index, String node) {
+ private IndexShard getIndexShard(String node) {
+ final Index index = resolveIndex(INDEX_NAME);
  IndicesService indicesService = internalCluster().getInstance(IndicesService.class, node);
  IndexService indexService = indicesService.indexServiceSafe(index);
  final Optional<Integer> shardId = indexService.shardIds().stream().findFirst();
@@ -409,7 +515,8 @@ private List<ShardSegments[]> getShardSegments(IndicesSegmentResponse indicesSeg
  }
 
  private Map<String, Segment> getLatestSegments(ShardSegments segments) {
- final Long latestPrimaryGen = segments.getSegments().stream().map(Segment::getGeneration).max(Long::compare).get();
+ final Optional<Long> generation = segments.getSegments().stream().map(Segment::getGeneration).max(Long::compare);
+ final Long latestPrimaryGen = generation.get();
  return segments.getSegments()
  .stream()
  .filter(s -> s.getGeneration() == latestPrimaryGen)
@@ -419,4 +526,31 @@ private Map<String, Segment> getLatestSegments(ShardSegments segments) {
  private Map<Boolean, List<ShardSegments>> segmentsByShardType(ShardSegments[] replicationGroupSegments) {
  return Arrays.stream(replicationGroupSegments).collect(Collectors.groupingBy(s -> s.getShardRouting().primary()));
  }
+
+ @Nullable
+ private ShardRouting getShardRoutingForNodeName(String nodeName) {
+ final ClusterState state = client(internalCluster().getClusterManagerName()).admin().cluster().prepareState().get().getState();
+ for (IndexShardRoutingTable shardRoutingTable : state.routingTable().index(INDEX_NAME)) {
+ for (ShardRouting shardRouting : shardRoutingTable.activeShards()) {
+ final String nodeId = shardRouting.currentNodeId();
+ final DiscoveryNode discoveryNode = state.nodes().resolveNode(nodeId);
+ if (discoveryNode.getName().equals(nodeName)) {
+ return shardRouting;
+ }
+ }
+ }
+ return null;
+ }
+
+ private void assertDocCounts(int expectedDocCount, String... nodeNames) {
+ for (String node : nodeNames) {
+ assertHitCount(client(node).prepareSearch(INDEX_NAME).setSize(0).setPreference("_only_local").get(), expectedDocCount);
+ }
+ }
+
+ private DiscoveryNode getNodeContainingPrimaryShard() {
+ final ClusterState state = client(internalCluster().getClusterManagerName()).admin().cluster().prepareState().get().getState();
+ final ShardRouting primaryShard = state.routingTable().index(INDEX_NAME).shard(0).primaryShard();
+ return state.nodes().resolveNode(primaryShard.currentNodeId());
+ }
 }

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

@@ -533,7 +533,7 @@ public synchronized IndexShard createShard(
  () -> globalCheckpointSyncer.accept(shardId),
  retentionLeaseSyncer,
  circuitBreakerService,
- this.indexSettings.isSegRepEnabled() && routing.primary() ? checkpointPublisher : null
+ this.indexSettings.isSegRepEnabled() ? checkpointPublisher : null
  );
  eventListener.indexShardStateChanged(indexShard, null, indexShard.state(), "shard created");
  eventListener.afterIndexShardCreated(indexShard);

server/src/main/java/org/opensearch/index/engine/NRTReplicationEngine.java

@@ -129,6 +129,23 @@ public synchronized void updateSegments(final SegmentInfos infos, long seqNo) th
  localCheckpointTracker.fastForwardProcessedSeqNo(seqNo);
  }
 
+ /**
+ * Persist the latest live SegmentInfos.
+ *
+ * This method creates a commit point from the latest SegmentInfos. It is intended to be used when this shard is about to be promoted as the new primary.
+ *
+ * TODO: If this method is invoked while the engine is currently updating segments on its reader, wait for that update to complete so the updated segments are used.
+ *
+ *
+ * @throws IOException - When there is an IO error committing the SegmentInfos.
+ */
+ public void commitSegmentInfos() throws IOException {
+ // TODO: This method should wait for replication events to finalize.
+ final SegmentInfos latestSegmentInfos = getLatestSegmentInfos();
+ store.commitSegmentInfos(latestSegmentInfos, localCheckpointTracker.getMaxSeqNo(), localCheckpointTracker.getProcessedCheckpoint());
+ translogManager.syncTranslog();
+ }
+
  @Override
  public String getHistoryUUID() {
  return loadHistoryUUID(lastCommittedSegmentInfos.userData);

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

@@ -40,7 +40,7 @@ public void beforeRefresh() throws IOException {
 
  @Override
  public void afterRefresh(boolean didRefresh) throws IOException {
- if (didRefresh && shard.getReplicationTracker().isPrimaryMode()) {
+ if (didRefresh && shard.state() != IndexShardState.CLOSED && shard.getReplicationTracker().isPrimaryMode()) {
  publisher.publish(shard);
  }
  }

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

@@ -164,7 +164,6 @@
 import org.opensearch.indices.recovery.RecoveryTarget;
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
-import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
 import org.opensearch.rest.RestStatus;
@@ -240,6 +239,7 @@ public class IndexShard extends AbstractIndexShardComponent implements IndicesCl
  private final GlobalCheckpointListeners globalCheckpointListeners;
  private final PendingReplicationActions pendingReplicationActions;
  private final ReplicationTracker replicationTracker;
+ private final SegmentReplicationCheckpointPublisher checkpointPublisher;
 
  protected volatile ShardRouting shardRouting;
  protected volatile IndexShardState state;
@@ -304,7 +304,6 @@ Runnable getGlobalCheckpointSyncer() {
  private final AtomicReference<Translog.Location> pendingRefreshLocation = new AtomicReference<>();
  private final RefreshPendingLocationListener refreshPendingLocationListener;
  private volatile boolean useRetentionLeasesInPeerRecovery;
- private final ReferenceManager.RefreshListener checkpointRefreshListener;
 
  public IndexShard(
  final ShardRouting shardRouting,
@@ -410,11 +409,7 @@ public boolean shouldCache(Query query) {
  persistMetadata(path, indexSettings, shardRouting, null, logger);
  this.useRetentionLeasesInPeerRecovery = replicationTracker.hasAllPeerRecoveryRetentionLeases();
  this.refreshPendingLocationListener = new RefreshPendingLocationListener();
- if (checkpointPublisher != null) {
- this.checkpointRefreshListener = new CheckpointRefreshListener(this, checkpointPublisher);
- } else {
- this.checkpointRefreshListener = null;
- }
+ this.checkpointPublisher = checkpointPublisher;
  }
 
  public ThreadPool getThreadPool() {
@@ -614,6 +609,11 @@ public void updateShardState(
  + newRouting;
  assert getOperationPrimaryTerm() == newPrimaryTerm;
  try {
+ if (indexSettings.isSegRepEnabled()) {
+ // this Shard's engine was read only, we need to update its engine before restoring local history from xlog.
+ assert newRouting.primary() && currentRouting.primary() == false;
+ promoteNRTReplicaToPrimary();
+ }
  replicationTracker.activatePrimaryMode(getLocalCheckpoint());
  ensurePeerRecoveryRetentionLeasesExist();
  /*
@@ -3220,11 +3220,11 @@ private EngineConfig newEngineConfig(LongSupplier globalCheckpointSupplier) {
  }
  };
 
- final List<ReferenceManager.RefreshListener> internalRefreshListener;
- if (this.checkpointRefreshListener != null) {
- internalRefreshListener = Arrays.asList(new RefreshMetricUpdater(refreshMetric), checkpointRefreshListener);
- } else {
- internalRefreshListener = Collections.singletonList(new RefreshMetricUpdater(refreshMetric));
+ final List<ReferenceManager.RefreshListener> internalRefreshListener = new ArrayList<>();
+ internalRefreshListener.add(new RefreshMetricUpdater(refreshMetric));
+
+ if (this.checkpointPublisher != null && indexSettings.isSegRepEnabled() && shardRouting.primary()) {
+ internalRefreshListener.add(new CheckpointRefreshListener(this, this.checkpointPublisher));
  }
 
  return this.engineConfigFactory.newEngineConfig(
@@ -4108,4 +4108,26 @@ RetentionLeaseSyncer getRetentionLeaseSyncer() {
  public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
  return getEngine().getSegmentInfosSnapshot();
  }
+
+ /**
+ * With segment replication enabled - prepare the shard's engine to be promoted as the new primary.
+ *
+ * If this shard is currently using a replication engine, this method:
+ * 1. Invokes {@link NRTReplicationEngine#commitSegmentInfos()} to ensure the engine can be reopened as writeable from the latest refresh point.
+ * InternalEngine opens its IndexWriter from an on-disk commit point, but this replica may have recently synced from a primary's refresh point, meaning it has documents searchable in its in-memory SegmentInfos
+ * that are not part of a commit point. This ensures that those documents are made part of a commit and do not need to be reindexed after promotion.
+ * 2. Invokes resetEngineToGlobalCheckpoint - This call performs the engine swap, opening up as a writeable engine and replays any operations in the xlog. The operations indexed from xlog here will be
+ * any ack'd writes that were not copied to this replica before promotion.
+ */
+ private void promoteNRTReplicaToPrimary() {
+ assert shardRouting.primary() && indexSettings.isSegRepEnabled();
+ getReplicationEngine().ifPresentOrElse(engine -> {
+ try {
+ engine.commitSegmentInfos();
+ resetEngineToGlobalCheckpoint();
+ } catch (IOException e) {
+ throw new EngineException(shardId, "Unable to update replica to writeable engine, failing shard", e);
+ }
+ }, () -> { throw new EngineException(shardId, "Expected replica engine to be of type NRTReplicationEngine"); });
+ }
 }