Dedup agent connection requests

Dedupes agent connection requests. Observed scenarios where on agent with old agent code send a storm of connection requests. In such a case, to ensure other agents are not impacted, deduping the connection requests.

Author: mprokopchuk

agent/src/main/java/com/cloud/agent/Agent.java

@@ -159,7 +159,7 @@ public int value() {
     CopyOnWriteArrayList<IAgentControlListener> controlListeners = new CopyOnWriteArrayList<>();
 
     IAgentShell shell;
-    NioConnection connection;
+    NioClient connection;
     ServerResource serverResource;
     Link link;
     Long id;
@@ -857,8 +857,20 @@ public void processStartupAnswer(final StartupAnswer startup, final Response res
             if (serverResource != null && !serverResource.isExitOnFailures()) {
                 logger.trace("{} does not allow exit on failure, reconnecting",
                         serverResource.getClass().getSimpleName());
+                // If the MS flagged this as a duplicate connection, retry the same MS using
+                // the host we just connected to — the agent already knows the correct hostname.
+                final String preferredHost = startup.isRetryCurrentMs() && connection != null
+                        ? connection.getHost() : null;
+                if (preferredHost != null) {
+                    logger.info("Duplicate connection rejected, will retry same MS: {}", preferredHost);
+                }
                 logger.info("Reconnecting for {}", link);
-                requestHandler.submit(() -> reconnect(link, null, false));
+                requestHandler.submit(() -> {
+                    if (preferredHost != null) {
+                        shell.getBackoffAlgorithm().waitBeforeRetry();
+                    }
+                    reconnect(link, preferredHost, false);
+                });
                 return;
             }
             logger.fatal("Got unsuccessful result {} from the answer {}, details: {}",

agent/src/test/java/com/cloud/agent/AgentTest.java

@@ -22,32 +22,41 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
+import java.util.concurrent.ExecutorService;
 
 import javax.naming.ConfigurationException;
 
+import com.cloud.agent.api.StartupAnswer;
+import com.cloud.agent.api.StartupRoutingCommand;
+import com.cloud.utils.backoff.BackoffAlgorithm;
+import com.cloud.utils.nio.NioClient;
 import org.apache.logging.log4j.Logger;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 import org.springframework.test.util.ReflectionTestUtils;
 
 import com.cloud.resource.ServerResource;
 import com.cloud.utils.backoff.impl.ConstantTimeBackoff;
 import com.cloud.utils.nio.Link;
-import com.cloud.utils.nio.NioConnection;
 
 @RunWith(MockitoJUnitRunner.class)
 public class AgentTest {
@@ -222,7 +231,7 @@ public void testStopAndCleanupConnectionConnectionIsNullDoesNothing() {
 
     @Test
     public void testStopAndCleanupConnectionValidConnectionNoWaitStopsAndCleansUp() throws IOException {
-        NioConnection mockConnection = mock(NioConnection.class);
+        NioClient mockConnection = mock(NioClient.class);
         agent.connection = mockConnection;
         agent.stopAndCleanupConnection();
         verify(mockConnection).stop();
@@ -231,7 +240,7 @@ public void testStopAndCleanupConnectionValidConnectionNoWaitStopsAndCleansUp()
 
     @Test
     public void testStopAndCleanupConnectionCleanupThrowsIOExceptionLogsWarning() throws IOException {
-        NioConnection mockConnection = mock(NioConnection.class);
+        NioClient mockConnection = mock(NioClient.class);
         agent.connection = mockConnection;
         doThrow(new IOException("Cleanup failed")).when(mockConnection).cleanUp();
         agent.stopAndCleanupConnection();
@@ -241,7 +250,7 @@ public void testStopAndCleanupConnectionCleanupThrowsIOExceptionLogsWarning() th
 
     @Test
     public void testStopAndCleanupConnectionValidConnectionWaitForStopWaitsForStartupToStop() throws IOException {
-        NioConnection mockConnection = mock(NioConnection.class);
+        NioClient mockConnection = mock(NioClient.class);
         ConstantTimeBackoff mockBackoff = mock(ConstantTimeBackoff.class);
         mockBackoff.setTimeToWait(0);
         agent.connection = mockConnection;
@@ -290,4 +299,77 @@ public void testSelectReconnectionHostWithNullSocketAddressUsesShellNextHost() {
         assertEquals("fallback.host.com", result);
         verify(shell, times(1)).getNextHost();
     }
+
+    private Agent setupAgentSpyForStartupAnswerTests(String currentHost) throws ConfigurationException {
+        Agent spyAgent = Mockito.spy(agent);
+        doNothing().when(spyAgent).reconnect(any(), any(), eq(false));
+
+        NioClient mockConnection = mock(NioClient.class);
+        when(mockConnection.getHost()).thenReturn(currentHost);
+        spyAgent.connection = mockConnection;
+
+        spyAgent.requestHandler = mock(ExecutorService.class);
+        when(serverResource.isExitOnFailures()).thenReturn(false);
+        when(shell.getBackoffAlgorithm()).thenReturn(mock(BackoffAlgorithm.class));
+        return spyAgent;
+    }
+
+    @Test
+    public void testProcessStartupAnswerRetryCurrentMsReconnectsToCurrentHost() throws Exception {
+        String currentHost = "ms-host-0.example.com";
+        Agent spyAgent = setupAgentSpyForStartupAnswerTests(currentHost);
+        Link link = mock(Link.class);
+
+        StartupAnswer answer = new StartupAnswer(new StartupRoutingCommand(), "duplicate");
+        answer.setRetryCurrentMs(true);
+
+        spyAgent.processStartupAnswer(answer, null, link);
+
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        verify(spyAgent.requestHandler).submit(captor.capture());
+        captor.getValue().run();
+
+        verify(spyAgent).reconnect(eq(link), eq(currentHost), eq(false));
+    }
+
+    @Test
+    public void testProcessStartupAnswerRetryCurrentMsAppliesBackoffBeforeReconnect() throws Exception {
+        Agent spyAgent = setupAgentSpyForStartupAnswerTests("ms-host-0.example.com");
+        Link link = mock(Link.class);
+
+        BackoffAlgorithm mockBackoff = mock(BackoffAlgorithm.class);
+        when(shell.getBackoffAlgorithm()).thenReturn(mockBackoff);
+
+        StartupAnswer answer = new StartupAnswer(new StartupRoutingCommand(), "duplicate");
+        answer.setRetryCurrentMs(true);
+
+        spyAgent.processStartupAnswer(answer, null, link);
+
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        verify(spyAgent.requestHandler).submit(captor.capture());
+        captor.getValue().run();
+
+        verify(mockBackoff).waitBeforeRetry();
+    }
+
+    @Test
+    public void testProcessStartupAnswerNoRetryCurrentMsReconnectsWithNullAndNoBackoff() throws Exception {
+        Agent spyAgent = setupAgentSpyForStartupAnswerTests("ms-host-0.example.com");
+        Link link = mock(Link.class);
+
+        BackoffAlgorithm mockBackoff = mock(BackoffAlgorithm.class);
+        when(shell.getBackoffAlgorithm()).thenReturn(mockBackoff);
+
+        StartupAnswer answer = new StartupAnswer(new StartupRoutingCommand(), "not a duplicate");
+        // retryCurrentMs defaults to false
+
+        spyAgent.processStartupAnswer(answer, null, link);
+
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        verify(spyAgent.requestHandler).submit(captor.capture());
+        captor.getValue().run();
+
+        verify(spyAgent).reconnect(eq(link), isNull(), eq(false));
+        verify(mockBackoff, never()).waitBeforeRetry();
+    }
 }

api/src/main/java/com/cloud/agent/api/Answer.java

@@ -43,10 +43,18 @@ public boolean getResult() {
         return result;
     }
 
+    public void setResult(boolean result) {
+        this.result = result;
+    }
+
     public String getDetails() {
         return details;
     }
 
+    public void setDetails(String details) {
+        this.details = details;
+    }
+
     @Override
     public boolean executeInSequence() {
         return false;

core/src/main/java/com/cloud/agent/api/StartupAnswer.java

@@ -30,6 +30,7 @@ public class StartupAnswer extends Answer {
 
     Integer agentHostStatusCheckDelaySec;
     private Map<String, String> params;
+    private boolean retryCurrentMs;
 
     protected StartupAnswer() {
         params = new HashMap<>();
@@ -80,4 +81,12 @@ public Integer getAgentHostStatusCheckDelaySec() {
     public void setAgentHostStatusCheckDelaySec(Integer agentHostStatusCheckDelaySec) {
         this.agentHostStatusCheckDelaySec = agentHostStatusCheckDelaySec;
     }
+
+    public boolean isRetryCurrentMs() {
+        return retryCurrentMs;
+    }
+
+    public void setRetryCurrentMs(boolean retryCurrentMs) {
+        this.retryCurrentMs = retryCurrentMs;
+    }
 }

engine/orchestration/src/main/java/com/cloud/agent/manager/AgentManagerImpl.java

@@ -37,6 +37,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -251,6 +252,7 @@ public class AgentManagerImpl extends ManagerBase implements AgentManager, Handl
     private int maxConcurrentNewAgentConnections;
     private final ConcurrentHashMap<String, Long> newAgentConnections = new ConcurrentHashMap<>();
     protected ScheduledExecutorService newAgentConnectionsMonitor;
+    private final ConcurrentHashMap<String, Boolean> _processingAgentGuids = new ConcurrentHashMap<>();
 
     private boolean _reconcileCommandsEnabled = false;
     private Integer _reconcileCommandInterval;
@@ -2005,37 +2007,93 @@ protected void runInContext() {
                 startups[i] = (StartupCommand)_cmds[i];
             }
 
-            AgentAttache attache = handleConnectedAgent(_link, startups, _request);
-            if (attache == null) {
-                logger.warn("Unable to create attache for agent: {}", _request);
+            if (_link.isTerminated()) {
+                logger.warn("Link is already terminated for agent: {}, skipping connection processing", _request);
+                if (startups[0] != null && startups[0].getGuid() != null) {
+                    _processingAgentGuids.remove(startups[0].getGuid());
+                }
+                unregisterNewConnection(_link.getSocketAddress());
+                return;
+            }
+
+            try {
+                AgentAttache attache = handleConnectedAgent(_link, startups, _request);
+                if (attache == null) {
+                    logger.warn("Unable to create attache for agent: {}", _request);
+                }
+            } finally {
+                if (startups[0] != null && startups[0].getGuid() != null) {
+                    _processingAgentGuids.remove(startups[0].getGuid());
+                }
+                unregisterNewConnection(_link.getSocketAddress());
             }
-            unregisterNewConnection(_link.getSocketAddress());
         }
     }
 
     protected void connectAgent(final Link link, final Command[] cmds, final Request request) {
         // send startup answer to agent in the very beginning, so agent can move on without waiting for the answer for an undetermined time,
         // if we put this logic into another thread pool.
-        Map<String, String> backoffConfiguration = ConfigKeyUtil.toMap(BackoffConfiguration.value());
+        // Build the startup answer — populated for all cases (success and reject)
         StartupAnswer[] answers = new StartupAnswer[cmds.length];
-        Command cmd;
         for (int i = 0; i < cmds.length; i++) {
-            cmd = cmds[i];
-            if (cmd instanceof StartupRoutingCommand || cmd instanceof StartupProxyCommand || cmd instanceof StartupSecondaryStorageCommand
-                    || cmd instanceof StartupStorageCommand) {
+            Command cmd = cmds[i];
+            if (cmd instanceof StartupRoutingCommand || cmd instanceof StartupProxyCommand
+                    || cmd instanceof StartupSecondaryStorageCommand || cmd instanceof StartupStorageCommand) {
                 StartupAnswer answer = new StartupAnswer((StartupCommand) cmds[i], 0, "", "", mgmtServiceConf.getPingInterval());
-                answer.setParams(backoffConfiguration);
+                answer.setParams(ConfigKeyUtil.toMap(BackoffConfiguration.value()));
                 answer.setAgentHostStatusCheckDelaySec(AgentHostStatusCheckDelay.value());
                 answers[i] = answer;
             }
         }
-        Response response = new Response(request, answers[0], _nodeId, -1);
+
+        if (answers[0] == null) {
+            // The leading command has no startup answer (an unhandled StartupCommand subtype, or a
+            // non-startup command as cmds[0]). Avoid the NullPointerException in the send/getResult
+            // path below and release the connection slot before returning.
+            logger.warn("No startup answer built for leading command {}; skipping connect processing", cmds[0] != null ? cmds[0].getClass().getSimpleName() : "null");
+            unregisterNewConnection(link.getSocketAddress());
+            return;
+        }
+
+        // Dedup: if this agent GUID is already being processed, reject and redirect back to this MS
+        String agentGuid = null;
+        for (Command cmd : cmds) {
+            if (cmd instanceof StartupCommand) {
+                agentGuid = ((StartupCommand) cmd).getGuid();
+                break;
+            }
+        }
+        if (agentGuid != null && _processingAgentGuids.putIfAbsent(agentGuid, Boolean.TRUE) != null) {
+            logger.info("Duplicate connection from agent GUID {}, rejecting and redirecting to same MS", agentGuid);
+            answers[0].setResult(false);
+            answers[0].setDetails("Duplicate connection, retry this MS");
+            answers[0].setRetryCurrentMs(true);
+        }
+
         try {
-            link.send(response.toBytes());
+            link.send(new Response(request, answers[0], _nodeId, -1).toBytes());
         } catch (ClosedChannelException e) {
             logger.debug("Failed to send startup answer: {}", e.getMessage(), e);
         }
-        _connectExecutor.execute(new HandleAgentConnectTask(link, cmds, request));
+
+        if (answers[0].getResult()) {
+            try {
+                _connectExecutor.execute(new HandleAgentConnectTask(link, cmds, request));
+            } catch (RejectedExecutionException e) {
+                // The task that would release the dedup GUID and the new-connection slot will not run,
+                // so clean them up here to avoid leaking the connection and permanently stranding the GUID.
+                logger.warn("Failed to schedule agent connect task for GUID {}; releasing dedup state", agentGuid, e);
+                if (agentGuid != null) {
+                    _processingAgentGuids.remove(agentGuid);
+                }
+                unregisterNewConnection(link.getSocketAddress());
+            }
+        } else {
+            // Duplicate/rejected connection: HandleAgentConnectTask (which unregisters the connection) is
+            // not scheduled, so release the new-connection slot here. The GUID is owned by the in-flight
+            // task that claimed it first, so it is intentionally not removed here.
+            unregisterNewConnection(link.getSocketAddress());
+        }
     }
 
     public class AgentHandler extends Task {