Fix interrupted KVM snapshot leaving VM on orphan overlay

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

@@ -663,6 +663,35 @@ protected void reconnect(final Link link, String preferredMSHost, boolean forTra
         logger.info("Connected to the host: {}", shell.getConnectedHost());
     }
 
+    protected void waitForCommandsAndReconnect(final Link link) {
+        int inProgress = commandsInProgress.get();
+        if (inProgress > 0) {
+            logger.warn("Agent disconnected with {} command(s) still in progress. Waiting for them to complete before reconnecting.", inProgress);
+
+            long maxWaitMs = 300000;
+            long waitedMs = 0;
+            while (commandsInProgress.get() > 0 && waitedMs < maxWaitMs) {
+                try {
+                    Thread.sleep(1000);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    break;
+                }
+                waitedMs += 1000;
+            }
+
+            if (commandsInProgress.get() > 0) {
+                logger.warn("Timed out waiting for {} command(s) to complete after {} ms. Proceeding with reconnect — "
+                        + "in-progress commands (e.g. snapshot merges) may not have completed.", commandsInProgress.get(), waitedMs);
+            } else {
+                logger.info("All in-progress commands completed. Proceeding with reconnect.");
+            }
+        }
+
+        logger.info("Reconnecting instead of exiting to avoid losing in-flight work (e.g. snapshot overlays).");
+        reconnect(link);
+    }
+
     protected void closeAndTerminateLink(final Link link) {
         if (link == null) {
             return;
@@ -700,7 +729,7 @@ public void processStartupAnswer(final Answer answer, final Response response, f
                 reconnect(link);
                 return;
             }
-            System.exit(1);
+            waitForCommandsAndReconnect(link);
         }
         if (!answerValid) {
             logger.warn("Threw away a startup answer because we're reconnecting.");
@@ -772,7 +801,7 @@ protected void processRequest(final Request request, final Link link) {
                             reconnect(link);
                             return;
                         }
-                        System.exit(1);
+                        waitForCommandsAndReconnect(link);
                         return;
                     } else if (cmd instanceof MaintainCommand) {
                         logger.debug("Received maintainCommand, do not cancel current tasks");

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

@@ -22,10 +22,12 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 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.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -254,4 +256,43 @@ public void testStopAndCleanupConnectionValidConnectionWaitForStopWaitsForStartu
         verify(mockConnection).cleanUp();
         verify(mockBackoff, times(3)).waitBeforeRetry();
     }
+
+    @Test
+    public void testWaitForCommandsAndReconnectNoCommandsInProgress() throws ConfigurationException {
+        Agent spyAgent = spy(new Agent(shell, 1, serverResource));
+        ReflectionTestUtils.setField(spyAgent, "logger", logger);
+        doNothing().when(spyAgent).reconnect(any(Link.class));
+
+        Link mockLink = mock(Link.class);
+        spyAgent.commandsInProgress.set(0);
+        spyAgent.waitForCommandsAndReconnect(mockLink);
+
+        verify(spyAgent).reconnect(mockLink);
+        verify(logger).info(eq("Reconnecting instead of exiting to avoid losing in-flight work (e.g. snapshot overlays)."));
+    }
+
+    @Test
+    public void testWaitForCommandsAndReconnectCommandsCompleteBeforeTimeout() throws ConfigurationException {
+        Agent spyAgent = spy(new Agent(shell, 1, serverResource));
+        ReflectionTestUtils.setField(spyAgent, "logger", logger);
+        doNothing().when(spyAgent).reconnect(any(Link.class));
+
+        Link mockLink = mock(Link.class);
+        spyAgent.commandsInProgress.set(1);
+
+        Thread decrementThread = new Thread(() -> {
+            try {
+                Thread.sleep(500);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            }
+            spyAgent.commandsInProgress.set(0);
+        });
+        decrementThread.start();
+
+        spyAgent.waitForCommandsAndReconnect(mockLink);
+
+        verify(spyAgent).reconnect(mockLink);
+        verify(logger).info(eq("All in-progress commands completed. Proceeding with reconnect."));
+    }
 }