server/src/main/java/org/opensearch/index/shard/IndexShard.java (1)

138-139: Guard getLastRefreshedCheckpoint() against non‑InternalEngine engines

getLastRefreshedCheckpoint() unconditionally casts Engine to InternalEngine. If a shard ever runs with a different Engine implementation (e.g., a custom EngineFactory, ReadOnlyEngine, or future engine type) and this method is called, it will fail with a ClassCastException (assert only catches it in tests).

Consider adding an explicit runtime check and a clearer failure mode, e.g.:

-    public long getLastRefreshedCheckpoint() {
-        Engine engine = getEngine();
-        assert engine instanceof InternalEngine;
-        return ((InternalEngine) engine).lastRefreshedCheckpoint();
-    }
+    public long getLastRefreshedCheckpoint() {
+        final Engine engine = getEngine();
+        if (engine instanceof InternalEngine internalEngine) {
+            return internalEngine.lastRefreshedCheckpoint();
+        }
+        throw new IllegalStateException(
+            "lastRefreshedCheckpoint is only available for InternalEngine but found [" + engine.getClass().getName() + "]"
+        );
+    }

This keeps the assumption explicit and avoids a silent hard ClassCastException if the invariant is ever violated.

Also applies to: 4028-4033

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationPrimaryPromotionIT.java (2)

34-38: Static test coordination state should be reset for future extensibility

The static coordination state:

public static boolean lockEnable = false;
public static CountDownLatch indexLuceneLatch = new CountDownLatch(1);
public static CountDownLatch flushLatch = new CountDownLatch(1);
public static CountDownLatch refreshLatch = new CountDownLatch(1);

combined with the mocked engines is fine for a single test, but it will behave oddly if additional tests are ever added to this class (latches already at zero, lockEnable left true, etc.).

Consider resetting all of these in a @Before (or @After) method, e.g. reassigning new CountDownLatch(1) instances and lockEnable = false, so each test starts from a known clean state.

Also applies to: 53-100

---

39-42: Ensure the @Before setup method actually runs

setup() is declared private:

@Before
private void setup() {
    internalCluster().startClusterManagerOnlyNode();
}

Depending on the JUnit/OpenSearch test harness configuration, private @Before methods may not be discovered. If that happens, this setup will silently be skipped and the test will rely on default cluster behaviour.

To be safe and consistent with typical JUnit usage, make this method public (and optionally rename to setUp):

-    @Before
-    private void setup() {
+    @Before
+    public void setup() {
         internalCluster().startClusterManagerOnlyNode();
     }

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationPrimaryPromotionIT.java (2)

109-163: Ensure test actually runs and cannot hang indefinitely

A few issues in the main test method impact its effectiveness and robustness:

1. Missing @Test annotation

The method testPrimaryStopped_ReplicaPromoted_no_data_loss is not annotated with @Test. Under JUnit 4 (as used by OpenSearchIntegTestCase), this will typically not be discovered as a test method. Add the annotation and import:

-import org.junit.Before;
+import org.junit.Before;
+import org.junit.Test;
...
-    // Used to test that primary promotion does not result in data loss.
-    public void testPrimaryStopped_ReplicaPromoted_no_data_loss() throws Exception {
+    // Used to test that primary promotion does not result in data loss.
+    @Test
+    public void testPrimaryStopped_ReplicaPromoted_no_data_loss() throws Exception {

This is critical because this IT is the primary guard against the data‑loss scenario.

2. Unbounded waits on latches and join() can hang the suite (previously flagged)

flushLatch.await(), refreshLatch.await(), and writeThread.join() are all unbounded. If the mocked engine or replication flow regresses, the test will hang indefinitely and stall the IT suite.

You can mitigate this with bounded waits and assertions:

-        flushLatch.await();
+        assertTrue("flushLatch timed out", flushLatch.await(30, TimeUnit.SECONDS));
...
-        refreshLatch.await();
-        writeThread.join();
+        assertTrue("refreshLatch timed out", refreshLatch.await(30, TimeUnit.SECONDS));
+        writeThread.join(TimeUnit.SECONDS.toMillis(30));
+        assertFalse("writeThread is still alive", writeThread.isAlive());

(Requires import java.util.concurrent.TimeUnit; and using assertTrue/assertFalse from the base test class.)

3. Mocked publish‑checkpoint exception is permanent; clarify intent (previously flagged in spirit)

The replica’s PublishCheckpointAction handler always throws a RemoteTransportException:

replicaTransportService.addRequestHandlingBehavior(
    PublishCheckpointAction.ACTION_NAME + TransportReplicationAction.REPLICA_ACTION_SUFFIX,
    (handler, request, channel, task) -> {
        throw new RemoteTransportException("mock remote transport exception", new OpenSearchRejectedExecutionException());
    }
);

If the intention is to simulate a permanent failure on publishing checkpoints, adding a short comment here would make that explicit. If the goal is to fail only the first attempt and allow retries to succeed, consider guarding this with a flag (e.g., an AtomicBoolean fail‑once pattern) so later attempts go through.

---

160-161: Use assertEquals for clearer hit‑count assertion (already suggested in earlier review)

The assertion:

assertTrue(response.getHits().getTotalHits().value() == 2);

works but produces a less helpful failure message. Prefer assertEquals with the correct long literal:

-        assertTrue(response.getHits().getTotalHits().value() == 2);
+        assertEquals(2L, response.getHits().getTotalHits().value());

This will show expected vs. actual values directly on failure.

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationPrimaryPromotionIT.java (1)

41-44: Static latch and flag state should be reset for test isolation

lockEnable and the CountDownLatch instances are public static and never reset. This is fine for a single test today, but will make the class fragile if additional tests are added or if the class is reused, since stale latch state (count already at 0, lockEnable left true) will change the control‑flow assumptions.

Consider reinitializing these in a @Before/@After method (or localizing them to the test) and, if they remain shared across threads, making lockEnable volatile to guarantee visibility.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationPrimaryPromotionIT.java (4)

47-50: JUnit lifecycle method should not be private.

With JUnit 4, methods annotated with @Before must be public; declaring private void setup() can cause the runner to treat it as an invalid lifecycle method and fail or ignore it. That would prevent the cluster-manager-only node from being started as intended.

Change the signature to be public:

     @Before
-    private void setup() {
+    public void setup() {
         internalCluster().startClusterManagerOnlyNode();
     }

---

80-92: Guard engine-level latch wait with a timeout and handle interruption explicitly.

MockInternalEngine.generateSeqNoForOperationOnPrimary blocks on indexLuceneLatch.await() without a timeout and wraps all Exceptions in a RuntimeException. If the latch is never released (e.g., test regression, early failure), this will block the engine thread indefinitely and hang the test run.

Add a bounded wait with an assertion:

                 if (lockEnable) {
                     flushLatch.countDown();
-                    indexLuceneLatch.await();
+                    if (!indexLuceneLatch.await(30, TimeUnit.SECONDS)) {
+                        throw new AssertionError("indexLuceneLatch timed out");
+                    }
                 }
-            } catch (Exception e) {
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
                 throw new RuntimeException(e);
             }

---

120-144: Add timeout to thread join and capture exceptions from the write thread.

Two issues with the write thread:

1. Line 144: Unbounded writeThread.join() can hang the test
   If the write thread encounters an unexpected failure or deadlock, this will block indefinitely. Add a timeout:

   -        writeThread.join();
   +        writeThread.join(TimeUnit.SECONDS.toMillis(30));
   +        assertFalse("writeThread still alive after timeout", writeThread.isAlive());

2. Line 120: Exceptions thrown inside the lambda are silently lost
   If client().prepareIndex() throws an exception (e.g., rejection, timeout), the test cannot detect it and may incorrectly pass. Capture and propagate exceptions:

   +        AtomicReference<Exception> writeError = new AtomicReference<>();
           Thread writeThread = new Thread(() -> {
   -            client().prepareIndex(INDEX_NAME).setId("2").setSource("foo2", "bar2").get();
   +            try {
   +                client().prepareIndex(INDEX_NAME).setId("2").setSource("foo2", "bar2").get();
   +            } catch (Exception e) {
   +                writeError.set(e);
   +            }
           });
           writeThread.start();
           ...
           writeThread.join(TimeUnit.SECONDS.toMillis(30));
   +        assertNull("Write thread failed: " + writeError.get(), writeError.get());

---

133-138: Clarify whether checkpoint publish should fail permanently or transiently.

The mock transport exception is thrown unconditionally for all PublishCheckpointAction requests throughout the entire test. If the intent is to simulate a transient network failure where retries eventually succeed, you should use a flag or counter to allow subsequent requests to pass.

If permanent failure is the intended scenario to validate data loss prevention, add a comment clarifying that all checkpoint publishes are expected to fail:

+        // Simulate permanent network partition: all checkpoint publishes to replica will fail
         replicaTransportService.addRequestHandlingBehavior(
             PublishCheckpointAction.ACTION_NAME + TransportReplicationAction.REPLICA_ACTION_SUFFIX,
             (handler, request, channel, task) -> {
                 throw new RemoteTransportException("mock remote transport exception", new OpenSearchRejectedExecutionException());
             }
         );

Otherwise, if retries should succeed after the first failure:

+        AtomicBoolean failOnce = new AtomicBoolean(true);
         replicaTransportService.addRequestHandlingBehavior(
             PublishCheckpointAction.ACTION_NAME + TransportReplicationAction.REPLICA_ACTION_SUFFIX,
             (handler, request, channel, task) -> {
+                if (failOnce.compareAndSet(true, false)) {
                     throw new RemoteTransportException("mock remote transport exception", new OpenSearchRejectedExecutionException());
+                }
+                handler.messageReceived(request, channel, task);
             }
         );

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro