From a8d47472b419e4eb493fc69829ab01ecd5dc6eb4 Mon Sep 17 00:00:00 2001
From: Wojciech Lukowicz <wojciech.lukowicz@weareadaptive.com>
Date: Fri, 2 Aug 2024 11:31:04 +0100
Subject: [PATCH] fix connectionId returned by follower session after
 disconnect

Previously, it would return the old connectionId until next reconnect,
now it will return NO_CONNECTION_ID.
---
 .../artio/library/LibraryPoller.java          | 11 ++++++++++
 .../artio/session/InternalSession.java        |  5 +++++
 .../artio/session/SessionWriter.java          |  6 +++++
 ...uenceNumberGatewayToGatewaySystemTest.java | 22 +++++++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/artio-core/src/main/java/uk/co/real_logic/artio/library/LibraryPoller.java b/artio-core/src/main/java/uk/co/real_logic/artio/library/LibraryPoller.java
index aa52297339..b568fa8bd6 100644
--- a/artio-core/src/main/java/uk/co/real_logic/artio/library/LibraryPoller.java
+++ b/artio-core/src/main/java/uk/co/real_logic/artio/library/LibraryPoller.java
@@ -1294,6 +1294,17 @@ public Action onDisconnect(
                     }
                 }
 
+                final WeakReference<SessionWriter> sessionWriterRef =
+                    sessionIdToFollowerSessionWriter.get(session.id());
+                if (sessionWriterRef != null)
+                {
+                    final SessionWriter sessionWriter = sessionWriterRef.get();
+                    if (sessionWriter != null)
+                    {
+                        InternalSession.disconnectWriter(sessionWriter);
+                    }
+                }
+
                 connectionIdToSession.remove(connectionId);
 
                 if (isEngineOwned)
diff --git a/artio-core/src/main/java/uk/co/real_logic/artio/session/InternalSession.java b/artio-core/src/main/java/uk/co/real_logic/artio/session/InternalSession.java
index c6cb5f01e1..413dc5793d 100644
--- a/artio-core/src/main/java/uk/co/real_logic/artio/session/InternalSession.java
+++ b/artio-core/src/main/java/uk/co/real_logic/artio/session/InternalSession.java
@@ -254,6 +254,11 @@ public void linkTo(final SessionWriter sessionWriter)
         sessionWriter.linkTo(this);
     }
 
+    public static void disconnectWriter(final SessionWriter writer)
+    {
+        writer.onDisconnect();
+    }
+
     public static void closeWriter(final SessionWriter writer)
     {
         writer.close();
diff --git a/artio-core/src/main/java/uk/co/real_logic/artio/session/SessionWriter.java b/artio-core/src/main/java/uk/co/real_logic/artio/session/SessionWriter.java
index 3fc76c7e3d..e7782d6617 100644
--- a/artio-core/src/main/java/uk/co/real_logic/artio/session/SessionWriter.java
+++ b/artio-core/src/main/java/uk/co/real_logic/artio/session/SessionWriter.java
@@ -23,6 +23,7 @@
 import uk.co.real_logic.artio.protocol.GatewayPublication;
 import uk.co.real_logic.artio.util.MutableAsciiBuffer;
 
+import static uk.co.real_logic.artio.GatewayProcess.NO_CONNECTION_ID;
 import static uk.co.real_logic.artio.messages.MessageStatus.OK;
 
 /**
@@ -207,6 +208,11 @@ void linkTo(final InternalSession session)
         this.connectionId = session.connectionId();
     }
 
+    void onDisconnect()
+    {
+        connectionId = NO_CONNECTION_ID;
+    }
+
     void checkState()
     {
         if (closed)
diff --git a/artio-system-tests/src/test/java/uk/co/real_logic/artio/system_tests/PersistentSequenceNumberGatewayToGatewaySystemTest.java b/artio-system-tests/src/test/java/uk/co/real_logic/artio/system_tests/PersistentSequenceNumberGatewayToGatewaySystemTest.java
index 6e29aa2eec..50758dd97b 100644
--- a/artio-system-tests/src/test/java/uk/co/real_logic/artio/system_tests/PersistentSequenceNumberGatewayToGatewaySystemTest.java
+++ b/artio-system-tests/src/test/java/uk/co/real_logic/artio/system_tests/PersistentSequenceNumberGatewayToGatewaySystemTest.java
@@ -50,6 +50,7 @@
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 import static uk.co.real_logic.artio.Constants.*;
+import static uk.co.real_logic.artio.GatewayProcess.NO_CONNECTION_ID;
 import static uk.co.real_logic.artio.MonitoringAgentFactory.consumeDistinctErrors;
 import static uk.co.real_logic.artio.Reply.State.ERRORED;
 import static uk.co.real_logic.artio.TestFixtures.launchMediaDriver;
@@ -988,6 +989,27 @@ private void resetSomeSequenceNumbersOfOfflineSessions(
         assertAcceptingSessionHasSequenceIndex(retry ? 2 : 1);
     }
 
+    @Test(timeout = TEST_TIMEOUT_IN_MS)
+    public void followerSessionShouldReturnCurrentConnectionId()
+    {
+        launch(this::nothing);
+
+        final SessionWriter followerSession = createFollowerSession(TEST_REPLY_TIMEOUT_IN_MS);
+        assertEquals(NO_CONNECTION_ID, followerSession.connectionId());
+
+        for (int i = 0; i < 2; i++)
+        {
+            connectPersistingSessions();
+
+            assertNotEquals(NO_CONNECTION_ID, followerSession.connectionId());
+            assertEquals(acceptingSession.connectionId(), followerSession.connectionId());
+
+            disconnectSessions();
+
+            assertEquals(NO_CONNECTION_ID, followerSession.connectionId());
+        }
+    }
+
     private void connectPersistingSessions()
     {
         connectPersistingSessions(AUTOMATIC_INITIAL_SEQUENCE_NUMBER, false);