fix: notify new SnapshotReplicationListeners about ongoing replication

We are using `SnapshotReplicationListener`s to transition to inactive
when a snapshot replication starts. Snapshot replication can start
before any listeners have registered, which means that the listener will
only be notified about snapshot replication finishing, triggering a
transition to follower without first transitioning to inactive.

Here we are keeping track of ongoing snapshot replication to immediately
notify listeners when they are registering.

(cherry picked from commit 3d1f6c5)

atomix/cluster/src/main/java/io/atomix/raft/impl/RaftContext.java

@@ -127,6 +127,7 @@ public class RaftContext implements AutoCloseable, HealthMonitorable {
   private PersistedSnapshot currentSnapshot;
 
   private boolean ongoingTransition = false;
+  private boolean ongoingSnapshotReplication = false;
 
   @SuppressWarnings("java:S3077") // allow volatile here, health is immutable
   private volatile HealthReport health = HealthReport.healthy(this);
@@ -469,6 +470,11 @@ public long setCommitIndex(final long commitIndex) {
   public void addSnapshotReplicationListener(
       final SnapshotReplicationListener snapshotReplicationListener) {
     snapshotReplicationListeners.add(snapshotReplicationListener);
+    if (ongoingSnapshotReplication) {
+      // Notify listener immediately if it registered during an ongoing replication.
+      // This is to prevent missing necessary state transitions.
+      snapshotReplicationListener.onSnapshotReplicationStarted();
+    }
   }
 
   /**
@@ -482,11 +488,13 @@ public void removeSnapshotReplicationListener(
   }
 
   public void notifySnapshotReplicationStarted() {
+    ongoingSnapshotReplication = true;
     snapshotReplicationListeners.forEach(SnapshotReplicationListener::onSnapshotReplicationStarted);
   }
 
   public void notifySnapshotReplicationCompleted() {
     snapshotReplicationListeners.forEach(l -> l.onSnapshotReplicationCompleted(term));
+    ongoingSnapshotReplication = false;
   }
 
   /**

atomix/cluster/src/test/java/io/atomix/raft/SnapshotReplicationListenerTest.java

@@ -17,6 +17,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
 
@@ -59,4 +60,16 @@ public void shouldNotifySnapshotReplicationListener() throws Throwable {
     verify(snapshotReplicationListener, timeout(1_000).times(1))
         .onSnapshotReplicationCompleted(follower.getTerm());
   }
+
+  @Test
+  public void shouldNotifyOnRegisteringListener() {
+    // given
+    final var snapshotReplicationListener = mock(SnapshotReplicationListener.class);
+    final var follower = raftRule.getFollower().orElseThrow();
+    // then
+    follower.getContext().notifySnapshotReplicationStarted();
+    verify(snapshotReplicationListener, never()).onSnapshotReplicationStarted();
+    follower.getContext().addSnapshotReplicationListener(snapshotReplicationListener);
+    verify(snapshotReplicationListener).onSnapshotReplicationStarted();
+  }
 }