fix: prevent sender control future race

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/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;
+    }
 }

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

@@ -84,22 +84,35 @@ public void addWorkerClient(int workerId, TransportClient client) {
     @Override
     public CompletableFuture<Void> send(int workerId, MessageType type)
                                         throws InterruptedException {
+        E.checkArgument(type == MessageType.START ||
+                        type == MessageType.FINISH,
+                        "The control message type must be START or FINISH, " +
+                        "but got '%s'", type);
         WorkerChannel channel = this.channels[channelId(workerId)];
-        CompletableFuture<Void> future = channel.newFuture();
-        future.whenComplete((r, e) -> {
-            channel.resetFuture(future);
-        });
+        CompletableFuture<Void> future = new CompletableFuture<>();
+        if (!channel.setControlFuture(future)) {
+            return future;
+        }
         /*
          * Control message just need message type is enough,
          * partitionId = -1 and buffer = null represents a meaningless value
          */
-        channel.queue.put(new QueuedMessage(-1, type, null));
+        try {
+            channel.queue.put(new QueuedMessage(-1, type, null, future));
+        } catch (InterruptedException e) {
+            channel.completeControlFuture(future, e);
+            throw e;
+        }
         return future;
     }
 
     @Override
     public void send(int workerId, QueuedMessage message)
                      throws InterruptedException {
+        E.checkArgument(message.type() != null &&
+                        message.type().category() == MessageType.Category.DATA,
+                        "The queued message type must be DATA, but got '%s'",
+                        message.type());
         WorkerChannel channel = this.channels[channelId(workerId)];
         channel.queue.put(message);
     }
@@ -108,7 +121,7 @@ public void send(int workerId, QueuedMessage message)
     public void transportExceptionCaught(TransportException cause, ConnectionId connectionId) {
         for (WorkerChannel channel : this.channels) {
             if (channel.client.connectionId().equals(connectionId)) {
-                channel.futureRef.get().completeExceptionally(cause);
+                channel.transportExceptionCaught(cause);
             }
         }
     }
@@ -137,11 +150,19 @@ public void run() {
                             ++emptyQueueCount;
                             continue;
                         }
-                        if (channel.doSend(message)) {
-                            // Only consume the message after it is sent
+                        try {
+                            if (channel.doSend(message)) {
+                                // Only consume the message after it is sent
+                                channel.queue.take();
+                            } else {
+                                ++busyClientCount;
+                            }
+                        } catch (TransportException | RuntimeException e) {
+                            channel.failDataSend(e);
+                            // Discard the failed data message to keep sending
                             channel.queue.take();
-                        } else {
-                            ++busyClientCount;
+                            LOG.warn("Failed to send {} message to {}, " +
+                                     "discard it", message.type(), channel, e);
                         }
                     }
                     int channelCount = channels.length;
@@ -172,9 +193,6 @@ public void run() {
                                   "Interrupted when waiting for message " +
                                   "queue not empty");
                     }
-                } catch (TransportException e) {
-                    // TODO: should handle this in main workflow thread
-                    throw new ComputerException("Failed to send message", e);
                 }
             }
             LOG.info("The send-executor is terminated");
@@ -227,75 +245,139 @@ private static class WorkerChannel {
         private final MessageQueue queue;
         // Each target worker has a TransportClient
         private final TransportClient client;
-        private final AtomicReference<CompletableFuture<Void>> futureRef;
+        private final AtomicReference<CompletableFuture<Void>> controlFutureRef;
+        private final AtomicReference<Throwable> dataFailureRef;
 
         public WorkerChannel(int workerId, MessageQueue queue,
                              TransportClient client) {
             this.workerId = workerId;
             this.queue = queue;
             this.client = client;
-            this.futureRef = new AtomicReference<>();
-        }
-
-        public CompletableFuture<Void> newFuture() {
-            CompletableFuture<Void> future = new CompletableFuture<>();
-            if (!this.futureRef.compareAndSet(null, future)) {
-                throw new ComputerException("The origin future must be null");
-            }
-            return future;
-        }
-
-        public void resetFuture(CompletableFuture<Void> future) {
-            if (!this.futureRef.compareAndSet(future, null)) {
-                throw new ComputerException("Failed to reset futureRef, " +
-                                            "expect future object is %s, " +
-                                            "but some thread modified it",
-                                            future);
-            }
+            this.controlFutureRef = new AtomicReference<>();
+            this.dataFailureRef = new AtomicReference<>();
         }
 
         public boolean doSend(QueuedMessage message)
                               throws TransportException, InterruptedException {
             switch (message.type()) {
                 case START:
-                    this.sendStartMessage();
+                    this.sendStartMessage(this.controlFuture(message));
                     return true;
                 case FINISH:
-                    this.sendFinishMessage();
+                    this.sendFinishMessage(this.controlFuture(message));
                     return true;
                 default:
                     return this.sendDataMessage(message);
             }
         }
 
-        public void sendStartMessage() throws TransportException {
-            this.client.startSessionAsync().whenComplete((r, e) -> {
-                CompletableFuture<Void> future = this.futureRef.get();
-                assert future != null;
-
-                if (e != null) {
-                    LOG.info("Failed to start session connected to {}", this);
-                    future.completeExceptionally(e);
-                } else {
-                    LOG.info("Start session connected to {}", this);
-                    future.complete(null);
+        private CompletableFuture<Void> controlFuture(QueuedMessage message) {
+            CompletableFuture<Void> future = message.controlFuture();
+            E.checkState(future != null,
+                         "The control future can't be null for message '%s'",
+                         message.type());
+            return future;
+        }
+
+        public void sendStartMessage(CompletableFuture<Void> future) {
+            if (!this.controlFutureInFlight(future)) {
+                return;
+            }
+            try {
+                this.client.startSessionAsync().whenComplete((r, e) -> {
+                    if (e != null) {
+                        LOG.info("Failed to start session connected to {}", this);
+                    } else {
+                        LOG.info("Start session connected to {}", this);
+                    }
+                    this.completeControlFuture(future, e);
+                });
+            } catch (TransportException e) {
+                this.completeControlFuture(future, e);
+            } catch (RuntimeException e) {
+                this.completeControlFuture(future, e);
+            }
+        }
+
+        public void sendFinishMessage(CompletableFuture<Void> future) {
+            if (!this.controlFutureInFlight(future)) {
+                return;
+            }
+            try {
+                this.client.finishSessionAsync().whenComplete((r, e) -> {
+                    if (e != null) {
+                        LOG.info("Failed to finish session connected to {}", this);
+                    } else {
+                        LOG.info("Finish session connected to {}", this);
+                    }
+                    this.completeControlFuture(future, e);
+                });
+            } catch (TransportException e) {
+                this.completeControlFuture(future, e);
+            } catch (RuntimeException e) {
+                this.completeControlFuture(future, e);
+            }
+        }
+
+        public void transportExceptionCaught(TransportException cause) {
+            CompletableFuture<Void> future = this.controlFutureRef.get();
+            if (future == null) {
+                this.failDataSend(cause);
+            } else {
+                this.completeControlFuture(future, cause);
+            }
+        }
+
+        public void failControlFuture(Throwable cause) {
+            CompletableFuture<Void> future = this.controlFutureRef.getAndSet(null);
+            if (future != null) {
+                future.completeExceptionally(cause);
+            }
+        }
+
+        public void failDataSend(Throwable cause) {
+            this.dataFailureRef.compareAndSet(null, cause);
+            this.failControlFuture(this.dataFailureRef.get());
+        }
+
+        private boolean setControlFuture(CompletableFuture<Void> future) {
+            Throwable failure = this.dataFailureRef.get();
+            if (failure != null) {
+                future.completeExceptionally(failure);
+                return false;
+            }
+            if (this.controlFutureRef.compareAndSet(null, future)) {
+                failure = this.dataFailureRef.get();
+                if (failure == null) {
+                    return true;
                 }
-            });
+                this.completeControlFuture(future, failure);
+                return false;
+            }
+            ComputerException e = new ComputerException(
+                                  "The origin future must be null");
+            future.completeExceptionally(e);
+            return false;
+        }
+
+        private boolean controlFutureInFlight(CompletableFuture<Void> future) {
+            return future != null && this.controlFutureRef.get() == future;
         }
 
-        public void sendFinishMessage() throws TransportException {
-            this.client.finishSessionAsync().whenComplete((r, e) -> {
-                CompletableFuture<Void> future = this.futureRef.get();
-                assert future != null;
-
-                if (e != null) {
-                    LOG.info("Failed to finish session connected to {}", this);
-                    future.completeExceptionally(e);
-                } else {
-                    LOG.info("Finish session connected to {}", this);
-                    future.complete(null);
+        private void completeControlFuture(CompletableFuture<Void> future,
+                                           Throwable cause) {
+            E.checkState(future != null, "The control future can't be null");
+            if (!this.controlFutureRef.compareAndSet(future, null)) {
+                if (cause != null) {
+                    this.failDataSend(cause);
                 }
-            });
+                return;
+            }
+            if (cause == null) {
+                future.complete(null);
+            } else {
+                future.completeExceptionally(cause);
+            }
         }
 
         public boolean sendDataMessage(QueuedMessage message)

computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/worker/WorkerService.java

@@ -63,6 +63,7 @@ public class WorkerService implements Closeable {
 
     private volatile boolean inited;
     private volatile boolean closed;
+    private volatile boolean registered;
 
     private final ComputerContext context;
     private final Map<Integer, ContainerInfo> workers;
@@ -83,6 +84,7 @@ public WorkerService() {
         this.workers = new HashMap<>();
         this.inited = false;
         this.closed = false;
+        this.registered = false;
         this.shutdownHook = new ShutdownHook();
     }
 
@@ -113,6 +115,7 @@ public synchronized void init(Config config) {
 
             LOG.info("{} register WorkerService", this);
             this.bsp4Worker.workerInitDone();
+            this.registered = true;
             this.connectToWorkers();
 
             this.computeManager = new ComputeManager(this.workerInfo.id(), this.context,
@@ -176,7 +179,6 @@ public synchronized void close() {
                 this.computeManager.close();
             } else {
                 LOG.warn("The computeManager is null");
-                return;
             }
         } catch (Exception e) {
             LOG.error("Error when closing ComputeManager", e);
@@ -194,8 +196,12 @@ public synchronized void close() {
         }
 
         try {
-            this.bsp4Worker.workerCloseDone();
-            this.bsp4Worker.close();
+            if (this.bsp4Worker != null) {
+                if (this.registered) {
+                    this.bsp4Worker.workerCloseDone();
+                }
+                this.bsp4Worker.close();
+            }
         } catch (Exception e) {
             LOG.error("Error while closing bsp4Worker", e);
         }

computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/network/netty/NettyTransportClientTest.java

@@ -127,7 +127,6 @@ public void testSend() throws IOException {
 
     @Test
     public void testDataUniformity() throws IOException {
-        NettyTransportClient client = (NettyTransportClient) this.oneClient();
         byte[] sourceBytes1 = StringEncodeUtil.encode("test data message");
         byte[] sourceBytes2 = StringEncodeUtil.encode("test data edge");
         byte[] sourceBytes3 = StringEncodeUtil.encode("test data vertex");
@@ -165,6 +164,7 @@ public void testDataUniformity() throws IOException {
             return null;
         }).when(serverHandler).handle(Mockito.any(), Mockito.eq(1), Mockito.any());
 
+        NettyTransportClient client = (NettyTransportClient) this.oneClient();
         client.startSession();
         client.send(MessageType.MSG, 1, ByteBuffer.wrap(sourceBytes1));
         client.send(MessageType.EDGE, 1, ByteBuffer.wrap(sourceBytes2));
@@ -274,12 +274,12 @@ public void testFlowControl() throws IOException {
 
     @Test
     public void testHandlerException() throws IOException {
-        NettyTransportClient client = (NettyTransportClient) this.oneClient();
-        client.startSession();
-
         Mockito.doThrow(new RuntimeException("test exception")).when(serverHandler)
                .handle(Mockito.any(), Mockito.anyInt(), Mockito.any());
 
+        NettyTransportClient client = (NettyTransportClient) this.oneClient();
+        client.startSession();
+
         ByteBuffer buffer = ByteBuffer.wrap(StringEncodeUtil.encode("test data"));
         boolean send = client.send(MessageType.MSG, 1, buffer);
         Assert.assertTrue(send);