fix: prevent sender control future race

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);
+    }
+
+    public 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;
     }
+
+    public CompletableFuture<Void> controlFuture() {
+        return this.controlFuture;
+    }
 }

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

@@ -85,15 +85,20 @@ public void addWorkerClient(int workerId, TransportClient client) {
     public CompletableFuture<Void> send(int workerId, MessageType type)
                                         throws InterruptedException {
         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;
     }
 
@@ -108,7 +113,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 +142,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.failControlFuture(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 +185,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 +237,105 @@ 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;
 
         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<>();
         }
 
         public boolean doSend(QueuedMessage message)
                               throws TransportException, InterruptedException {
             switch (message.type()) {
                 case START:
-                    this.sendStartMessage();
+                    this.sendStartMessage(message.controlFuture());
                     return true;
                 case FINISH:
-                    this.sendFinishMessage();
+                    this.sendFinishMessage(message.controlFuture());
                     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);
-                }
-            });
+        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() 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);
-                }
-            });
+        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) {
+            this.failControlFuture(cause);
+        }
+
+        public void failControlFuture(Throwable cause) {
+            CompletableFuture<Void> future = this.controlFutureRef.getAndSet(null);
+            if (future != null) {
+                future.completeExceptionally(cause);
+            }
+        }
+
+        private boolean setControlFuture(CompletableFuture<Void> future) {
+            if (this.controlFutureRef.compareAndSet(null, future)) {
+                return true;
+            }
+            ComputerException e = new ComputerException(
+                                  "The origin future must be null");
+            future.completeExceptionally(e);
+            return false;
+        }
+
+        private boolean controlFutureInFlight(CompletableFuture<Void> future) {
+            return this.controlFutureRef.get() == future;
+        }
+
+        private void completeControlFuture(CompletableFuture<Void> future,
+                                           Throwable cause) {
+            if (!this.controlFutureRef.compareAndSet(future, null)) {
+                return;
+            }
+            if (cause == null) {
+                future.complete(null);
+            } else {
+                future.completeExceptionally(cause);
+            }
         }
 
         public boolean sendDataMessage(QueuedMessage message)

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);