merge: #8853

8853: [Backport stable/1.2] fix: notify new `SnapshotReplicationListener`s about missed replications r=oleschoenburg a=github-actions[bot] # Description Backport of #8834 to `stable/1.2`. relates to #8830 Co-authored-by: Ole Schönburg <ole.schoenburg@gmail.com>
camunda · Feb 28, 2022 · 88024a3 · 88024a3
2 parents 1e040dc + e7a053f
commit 88024a3
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 4 deletions.
diff --git a/atomix/cluster/src/main/java/io/atomix/raft/impl/RaftContext.java b/atomix/cluster/src/main/java/io/atomix/raft/impl/RaftContext.java
@@ -128,6 +128,9 @@ public class RaftContext implements AutoCloseable, HealthMonitorable {
   private volatile HealthStatus health = HealthStatus.HEALTHY;
 
   private boolean ongoingTransition = false;
+  // Keeps track of snapshot replication to notify new listeners about missed events
+  private MissedSnapshotReplicationEvents missedSnapshotReplicationEvents =
+      MissedSnapshotReplicationEvents.NONE;
 
   private long lastHeartbeat;
   private final RaftPartitionConfig partitionConfig;
@@ -466,7 +469,25 @@ public long setCommitIndex(final long commitIndex) {
    */
   public void addSnapshotReplicationListener(
       final SnapshotReplicationListener snapshotReplicationListener) {
-    snapshotReplicationListeners.add(snapshotReplicationListener);
+    threadContext.execute(
+        () -> {
+          snapshotReplicationListeners.add(snapshotReplicationListener);
+          // Notify listener immediately if it registered during an ongoing replication.
+          // This is to prevent missing necessary state transitions.
+          switch (missedSnapshotReplicationEvents) {
+            case STARTED:
+              snapshotReplicationListener.onSnapshotReplicationStarted();
+              break;
+            case COMPLETED:
+              {
+                snapshotReplicationListener.onSnapshotReplicationStarted();
+                snapshotReplicationListener.onSnapshotReplicationCompleted(term);
+                break;
+              }
+            default:
+              break;
+          }
+        });
   }
 
   /**
@@ -476,15 +497,24 @@ public void addSnapshotReplicationListener(
    */
   public void removeSnapshotReplicationListener(
       final SnapshotReplicationListener snapshotReplicationListener) {
-    snapshotReplicationListeners.remove(snapshotReplicationListener);
+    threadContext.execute(() -> snapshotReplicationListeners.remove(snapshotReplicationListener));
   }
 
   public void notifySnapshotReplicationStarted() {
-    snapshotReplicationListeners.forEach(SnapshotReplicationListener::onSnapshotReplicationStarted);
+    threadContext.execute(
+        () -> {
+          missedSnapshotReplicationEvents = MissedSnapshotReplicationEvents.STARTED;
+          snapshotReplicationListeners.forEach(
+              SnapshotReplicationListener::onSnapshotReplicationStarted);
+        });
   }
 
   public void notifySnapshotReplicationCompleted() {
-    snapshotReplicationListeners.forEach(l -> l.onSnapshotReplicationCompleted(term));
+    threadContext.execute(
+        () -> {
+          snapshotReplicationListeners.forEach(l -> l.onSnapshotReplicationCompleted(term));
+          missedSnapshotReplicationEvents = MissedSnapshotReplicationEvents.COMPLETED;
+        });
   }
 
   /**
@@ -1121,4 +1151,14 @@ public enum State {
     ACTIVE,
     READY,
   }
+
+  /**
+   * Keeps track of potentially missed snapshot replication events to properly notify newly
+   * registered listeners.
+   */
+  private enum MissedSnapshotReplicationEvents {
+    NONE,
+    STARTED,
+    COMPLETED
+  }
 }
diff --git a/atomix/cluster/src/test/java/io/atomix/raft/SnapshotReplicationListenerTest.java b/atomix/cluster/src/test/java/io/atomix/raft/SnapshotReplicationListenerTest.java
@@ -18,6 +18,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import java.util.concurrent.CountDownLatch;
@@ -58,4 +59,44 @@ public void shouldNotifySnapshotReplicationListener() throws Throwable {
     verify(snapshotReplicationListener, timeout(1_000).times(1))
         .onSnapshotReplicationCompleted(follower.getTerm());
   }
+
+  @Test
+  public void shouldCallStartedOnRegister() {
+    // given
+    final var snapshotReplicationListener = mock(SnapshotReplicationListener.class);
+    final var follower = raftRule.getFollower().orElseThrow();
+    // when
+    follower.getContext().notifySnapshotReplicationStarted();
+    follower.getContext().addSnapshotReplicationListener(snapshotReplicationListener);
+    // then
+    verify(snapshotReplicationListener, timeout(1_000).times(1)).onSnapshotReplicationStarted();
+  }
+
+  @Test
+  public void shouldCallStartedAndCompletedOnRegister() {
+    // given
+    final var snapshotReplicationListener = mock(SnapshotReplicationListener.class);
+    final var follower = raftRule.getFollower().orElseThrow();
+    // when
+    follower.getContext().notifySnapshotReplicationStarted();
+    follower.getContext().notifySnapshotReplicationCompleted();
+    follower.getContext().addSnapshotReplicationListener(snapshotReplicationListener);
+    // then
+    verify(snapshotReplicationListener, timeout(1_000).times(1)).onSnapshotReplicationStarted();
+    verify(snapshotReplicationListener, timeout(1_000).times(1))
+        .onSnapshotReplicationCompleted(follower.getTerm());
+  }
+
+  @Test
+  public void shouldNotCallListenerOnRegister() {
+    // given
+    final var snapshotReplicationListener = mock(SnapshotReplicationListener.class);
+    final var follower = raftRule.getFollower().orElseThrow();
+    // when
+    follower.getContext().addSnapshotReplicationListener(snapshotReplicationListener);
+    // then
+    verify(snapshotReplicationListener, times(0)).onSnapshotReplicationStarted();
+    verify(snapshotReplicationListener, times(0))
+        .onSnapshotReplicationCompleted(follower.getTerm());
+  }
 }