fix: prevent sender control future race

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/input/WorkerInputManager.java

@@ -139,6 +139,7 @@ public void loadGraph() {
                                         "sending edges", e);
         }).join();
         this.sendManager.finishSend(MessageType.EDGE);
+        this.sendManager.checkFatal();
         this.sendManager.clearBuffer();
     }
 

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/manager/Managers.java

@@ -21,6 +21,7 @@
 import java.util.Iterator;
 import java.util.Map;
 
+import org.apache.hugegraph.computer.core.common.exception.ComputerException;
 import org.apache.hugegraph.computer.core.config.Config;
 import org.apache.hugegraph.util.E;
 import org.apache.hugegraph.util.InsertionOrderUtil;
@@ -67,8 +68,22 @@ public void initedAll(Config config) {
     }
 
     public void closeAll(Config config) {
+        ComputerException exception = null;
         for (Manager manager : this.all()) {
-            manager.close(config);
+            try {
+                manager.close(config);
+            } catch (Throwable e) {
+                // Close all managers even if some of them failed
+                if (exception == null) {
+                    exception = new ComputerException(
+                                "Failed to close manager '%s'", e, manager.name());
+                } else {
+                    exception.addSuppressed(e);
+                }
+            }
+        }
+        if (exception != null) {
+            throw exception;
         }
     }
 

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/master/MasterService.java

@@ -159,7 +159,7 @@ public synchronized void close() {
             LOG.error("Error occurred while closing master service", e);
         }
 
-        if (!failed && this.bsp4Master != null) {
+        if (this.inited && !failed && this.bsp4Master != null) {
             this.bsp4Master.waitWorkersCloseDone();
         }
 

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/network/DataServerManager.java

@@ -54,9 +54,8 @@ public void init(Config config) {
 
     @Override
     public void close(Config config) {
-        InetSocketAddress address = this.address();
         this.connectionManager.shutdownServer();
-        LOG.info("DataServerManager closed with address '{}'", address);
+        LOG.info("DataServerManager closed");
     }
 
     public InetSocketAddress address() {

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/MessageSendManager.java

@@ -150,6 +150,7 @@ public void startSend(MessageType type) {
                                     .map(this.partitioner::workerId)
                                     .collect(Collectors.toSet());
         this.sendControlMessageToWorkers(workerIds, MessageType.START);
+        this.sender.checkFatal();
         LOG.info("Start sending message(type={})", type);
     }
 
@@ -166,6 +167,7 @@ public void finishSend(MessageType type) {
                                     .map(this.partitioner::workerId)
                                     .collect(Collectors.toSet());
         this.sendControlMessageToWorkers(workerIds, MessageType.FINISH);
+        this.sender.checkFatal();
         LOG.info("Finish sending message(type={},count={},bytes={})",
                  type, stat.messageCount(), stat.messageBytes());
     }
@@ -178,6 +180,10 @@ public void clearBuffer() {
         this.buffers.clear();
     }
 
+    public void checkFatal() {
+        this.checkException();
+    }
+
     private void sortIfTargetBufferIsFull(WriteBuffers buffer,
                                           int partitionId,
                                           MessageType type) {
@@ -286,6 +292,7 @@ private void sendControlMessageToWorkers(Set<Integer> workerIds,
     }
 
     private void checkException() {
+        this.sender.checkFatal();
         if (this.exception.get() != null) {
             throw new ComputerException("Failed to send message",
                                         this.exception.get());

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/MessageSender.java

@@ -45,4 +45,14 @@ CompletableFuture<Void> send(int workerId, MessageType type)
      * an exception is thrown processing message.
      */
     void transportExceptionCaught(TransportException cause, ConnectionId connectionId);
+
+    /**
+     * Check whether the sender has encountered a fatal error. Implementations
+     * that run background threads should propagate the first fatal error to
+     * callers so that the caller can fail fast instead of hanging on a future
+     * or barrier.
+     */
+    default void checkFatal() {
+        // no-op by default
+    }
 }

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessage.java

@@ -18,6 +18,7 @@
 package org.apache.hugegraph.computer.core.sender;
 
 import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
 
 import org.apache.hugegraph.computer.core.network.message.MessageType;
 
@@ -26,11 +27,18 @@ public class QueuedMessage {
     private final int partitionId;
     private final MessageType type;
     private final ByteBuffer buffer;
+    private final CompletableFuture<Void> controlFuture;
 
     public QueuedMessage(int partitionId, MessageType type, ByteBuffer buffer) {
+        this(partitionId, type, buffer, null);
+    }
+
+    QueuedMessage(int partitionId, MessageType type, ByteBuffer buffer,
+                  CompletableFuture<Void> controlFuture) {
         this.partitionId = partitionId;
         this.type = type;
         this.buffer = buffer;
+        this.controlFuture = controlFuture;
     }
 
     public int partitionId() {
@@ -44,4 +52,8 @@ public MessageType type() {
     public ByteBuffer buffer() {
         return this.buffer;
     }
+
+    CompletableFuture<Void> controlFuture() {
+        return this.controlFuture;
+    }
 }