IGNITE-22319 Fixed node crashing if a snapshot restore cancelled due to network issues

modules/core/src/main/java/org/apache/ignite/internal/managers/communication/GridIoManager.java

@@ -1958,6 +1958,14 @@ public void removeTransmissionHandler(Object topic) {
                 "on local node [nodeId=" + ctx.localNodeId() + ']'));
     }
 
+    /**
+     * @param topic The topic to interrupt receiver from.
+     * @param ex Exception.
+     */
+    public void interruptTransmissionReceiver(Object topic, Exception ex) {
+        interruptReceiver(rcvCtxs.remove(topic), ex);
+    }
+
     /**
      * This method must be used prior to opening a {@link TransmissionSender} by calling
      * {@link #openTransmissionSender(UUID, Object)} to ensure that remote and local nodes

modules/core/src/main/java/org/apache/ignite/internal/processors/cache/persistence/snapshot/IgniteSnapshotManager.java

@@ -164,7 +164,6 @@
 import org.apache.ignite.internal.util.GridCloseableIteratorAdapter;
 import org.apache.ignite.internal.util.distributed.DistributedProcess;
 import org.apache.ignite.internal.util.distributed.InitMessage;
-import org.apache.ignite.internal.util.future.GridCompoundFuture;
 import org.apache.ignite.internal.util.future.GridCompoundIdentityFuture;
 import org.apache.ignite.internal.util.future.GridFinishedFuture;
 import org.apache.ignite.internal.util.future.GridFutureAdapter;
@@ -3785,47 +3784,32 @@ public synchronized void stop() {
 
             while ((r = queue.poll()) != null)
                 r.acceptException(new IgniteException(SNP_NODE_STOPPING_ERR_MSG));
-
-            Set<RemoteSnapshotFilesRecevier> futs = activeTasks();
-            GridCompoundFuture<Void, Void> stopFut = new GridCompoundFuture<>();
-
-            try {
-                for (IgniteInternalFuture<Void> fut : futs)
-                    stopFut.add(fut);
-
-                stopFut.markInitialized().get();
-            }
-            catch (IgniteCheckedException e) {
-                throw new IgniteException(e);
-            }
         }
 
         /**
          * @param nodeId A node left the cluster.
          */
-        public void onNodeLeft(UUID nodeId) {
-            Set<RemoteSnapshotFilesRecevier> futs = activeTasks();
+        public synchronized void onNodeLeft(UUID nodeId) {
+            if (active == null || active.isDone())
+                return;
+
             ClusterTopologyCheckedException ex = new ClusterTopologyCheckedException("The node from which a snapshot has been " +
                 "requested left the grid");
 
-            futs.forEach(t -> {
+            if (active.rmtNodeId.equals(nodeId))
+                interruptActive(ex);
+
+            queue.forEach(t -> {
                 if (t.rmtNodeId.equals(nodeId))
                     t.acceptException(ex);
             });
         }
 
-        /**
-         * @return The set of currently scheduled tasks, some of them may be already completed.
-         */
-        private Set<RemoteSnapshotFilesRecevier> activeTasks() {
-            Set<RemoteSnapshotFilesRecevier> futs = new HashSet<>(queue);
-
-            RemoteSnapshotFilesRecevier active0 = active;
+        /** Interrupts current active task (if present) to safely shedule next. */
+        private void interruptActive(Exception e) {
+            cctx.kernalContext().io().interruptTransmissionReceiver(DFLT_INITIAL_SNAPSHOT_TOPIC, e);
 
-            if (active0 != null)
-                futs.add(active0);
-
-            return futs;
+            active.acceptException(e);
         }
 
         /** {@inheritDoc} */
@@ -3895,20 +3879,22 @@ private Set<RemoteSnapshotFilesRecevier> activeTasks() {
                 else if (msg instanceof SnapshotFilesFailureMessage) {
                     SnapshotFilesFailureMessage respMsg0 = (SnapshotFilesFailureMessage)msg;
 
-                    RemoteSnapshotFilesRecevier task = active;
+                    synchronized (this) {
+                        RemoteSnapshotFilesRecevier task = active;
 
-                    if (task == null || !task.reqId.equals(respMsg0.id())) {
-                        if (log.isInfoEnabled()) {
-                            log.info("A stale snapshot response message has been received. Will be ignored " +
-                                "[fromNodeId=" + nodeId + ", response=" + respMsg0 + ']');
-                        }
+                        if (task == null || !task.reqId.equals(respMsg0.id())) {
+                            if (log.isInfoEnabled()) {
+                                log.info("A stale snapshot response message has been received. Will be ignored " +
+                                    "[fromNodeId=" + nodeId + ", response=" + respMsg0 + ']');
+                            }
 
-                        return;
-                    }
+                            return;
+                        }
 
-                    if (respMsg0.errorMessage() != null) {
-                        task.acceptException(new IgniteCheckedException("Request cancelled. The snapshot operation stopped " +
-                            "on the remote node with an error: " + respMsg0.errorMessage()));
+                        if (respMsg0.errorMessage() != null) {
+                            interruptActive(new IgniteCheckedException("Request cancelled. The snapshot operation " +
+                                "stopped on the remote node with an error: " + respMsg0.errorMessage()));
+                        }
                     }
                 }
             }
@@ -3941,7 +3927,7 @@ else if (msg instanceof SnapshotFilesFailureMessage) {
         }
 
         /** {@inheritDoc} */
-        @Override public void onException(UUID nodeId, Throwable ex) {
+        @Override public synchronized void onException(UUID nodeId, Throwable ex) {
             RemoteSnapshotFilesRecevier task = active;
 
             if (task == null)

modules/core/src/main/java/org/apache/ignite/internal/processors/cache/persistence/snapshot/SnapshotRestoreProcess.java

@@ -1200,8 +1200,7 @@ private IgniteInternalFuture<Boolean> preload(UUID reqId) {
         }
         catch (Exception ex) {
             opCtx0.errHnd.accept(ex);
-
-            return new GridFinishedFuture<>(ex);
+            retFut.onDone(ex);
         }
 
         return retFut;
@@ -1887,7 +1886,7 @@ private static class SnapshotRestoreContext {
         private volatile Map<Integer, StoredCacheData> cfgs = Collections.emptyMap();
 
         /** Graceful shutdown future. */
-        private volatile IgniteFuture<?> stopFut;
+        private volatile IgniteFuture<Void> stopFut;
 
         /** Operation start time. */
         private final long startTime;

modules/core/src/test/java/org/apache/ignite/internal/processors/cache/persistence/snapshot/IgniteSnapshotRestoreFromRemoteTest.java

@@ -30,6 +30,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -64,6 +66,7 @@
 import static org.apache.ignite.internal.processors.cache.persistence.file.FilePageStoreManager.partId;
 import static org.apache.ignite.internal.processors.cache.persistence.snapshot.IgniteSnapshotManager.resolveSnapshotWorkDirectory;
 import static org.apache.ignite.testframework.GridTestUtils.assertContains;
+import static org.apache.ignite.testframework.GridTestUtils.assertThrowsWithCause;
 
 /** */
 public class IgniteSnapshotRestoreFromRemoteTest extends IgniteClusterSnapshotRestoreBaseTest {
@@ -323,6 +326,46 @@ public void testSnapshotCachesStoppedIfLoadingFailOnRemote() throws Exception {
         ensureCacheAbsent(dfltCacheCfg);
     }
 
+    /** @throws Exception If failed. */
+    @Test
+    public void testRestoreConnectionLost() throws Exception {
+        IgniteEx coord = startDedicatedGrids(SECOND_CLUSTER_PREFIX, 2);
+
+        copyAndShuffle(snpParts, G.allGrids());
+
+        // Start a new node without snapshot working directory.
+        IgniteEx emptyNode = startDedicatedGrid(SECOND_CLUSTER_PREFIX, 2);
+
+        emptyNode.cluster().state(ClusterState.ACTIVE);
+
+        emptyNode.cache(DEFAULT_CACHE_NAME).destroy();
+        awaitPartitionMapExchange();
+
+        CountDownLatch latch = new CountDownLatch(1);
+
+        IgniteSnapshotManager mgr = snp(coord);
+        mgr.remoteSnapshotSenderFactory(new BiFunction<String, UUID, SnapshotSender>() {
+            @Override public SnapshotSender apply(String s, UUID uuid) {
+                return new DelegateSnapshotSender(log, mgr.snapshotExecutorService(), mgr.remoteSnapshotSenderFactory(s, uuid)) {
+                    @Override public void sendPart0(File part, String cacheDirName, GroupPartitionId pair, Long length) {
+                        delegate.sendPart0(part, cacheDirName, pair, length);
+
+                        latch.countDown();
+                    }
+                };
+            }
+        });
+
+        // Restore all cache groups.
+        IgniteFuture<Void> fut = emptyNode.snapshot().restoreSnapshot(SNAPSHOT_NAME, null);
+
+        latch.await(TIMEOUT, TimeUnit.MILLISECONDS);
+
+        coord.close();
+
+        assertThrowsWithCause(() -> fut.get(TIMEOUT), IgniteException.class);
+    }
+
     /**
      * @param snpParts Snapshot parts.
      * @param toNodes List of toNodes to copy parts to.