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,12 @@ 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<>();
         /*
          * 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));
+        channel.queue.put(new QueuedMessage(-1, type, null, future));
         return future;
     }
 
@@ -108,7 +105,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);
             }
         }
     }
@@ -227,75 +224,92 @@ 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;
-
+        public void sendStartMessage(CompletableFuture<Void> future)
+                                     throws TransportException {
+            this.setControlFuture(future);
+            try {
+                this.client.startSessionAsync().whenComplete((r, e) -> {
                 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);
                 }
-            });
+                this.completeControlFuture(future, e);
+                });
+            } catch (TransportException e) {
+                this.completeControlFuture(future, e);
+                throw e;
+            }
         }
 
-        public void sendFinishMessage() throws TransportException {
-            this.client.finishSessionAsync().whenComplete((r, e) -> {
-                CompletableFuture<Void> future = this.futureRef.get();
-                assert future != null;
-
+        public void sendFinishMessage(CompletableFuture<Void> future)
+                                      throws TransportException {
+            this.setControlFuture(future);
+            try {
+                this.client.finishSessionAsync().whenComplete((r, e) -> {
                 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);
                 }
-            });
+                this.completeControlFuture(future, e);
+                });
+            } catch (TransportException e) {
+                this.completeControlFuture(future, e);
+                throw e;
+            }
+        }
+
+        public void transportExceptionCaught(TransportException cause) {
+            CompletableFuture<Void> future = this.controlFutureRef.getAndSet(null);
+            if (future != null) {
+                future.completeExceptionally(cause);
+            }
+        }
+
+        private void setControlFuture(CompletableFuture<Void> future) {
+            if (!this.controlFutureRef.compareAndSet(null, future)) {
+                ComputerException e = new ComputerException(
+                                      "The origin future must be null");
+                future.completeExceptionally(e);
+                throw e;
+            }
+        }
+
+        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/sender/QueuedMessageSenderTest.java

@@ -17,8 +17,18 @@
 
 package org.apache.hugegraph.computer.core.sender;
 
+import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.apache.hugegraph.computer.core.common.exception.TransportException;
 import org.apache.hugegraph.computer.core.config.ComputerOptions;
 import org.apache.hugegraph.computer.core.config.Config;
+import org.apache.hugegraph.computer.core.network.message.MessageType;
 import org.apache.hugegraph.computer.core.worker.MockComputation2;
 import org.apache.hugegraph.computer.suite.unit.UnitTestBase;
 import org.apache.hugegraph.testutil.Assert;
@@ -64,4 +74,147 @@ public void testInitAndClose() {
         Assert.assertTrue(ImmutableSet.of(Thread.State.TERMINATED)
                                       .contains(sendExecutor.getState()));
     }
+
+    @Test
+    public void testControlFutureCanQueueNextControlBeforeCompletionDependentFinishes()
+            throws Exception {
+        QueuedMessageSender sender = new QueuedMessageSender(this.config);
+        ControlFutureClient client = new ControlFutureClient();
+        sender.addWorkerClient(1, client);
+        sender.addWorkerClient(2, new MockTransportClient());
+        sender.init();
+
+        CountDownLatch completionStarted = new CountDownLatch(1);
+        CountDownLatch allowCompletion = new CountDownLatch(1);
+        Thread completionThread = null;
+        try {
+            CompletableFuture<Void> startFuture = sender.send(1,
+                                                               MessageType.START);
+            Assert.assertTrue(client.awaitStart());
+            startFuture.whenComplete((r, e) -> {
+                completionStarted.countDown();
+                try {
+                    allowCompletion.await();
+                } catch (InterruptedException exception) {
+                    Thread.currentThread().interrupt();
+                    throw new AssertionError(exception);
+                }
+            });
+
+            completionThread = new Thread(client::completeStart);
+            completionThread.start();
+            Assert.assertTrue(completionStarted.await(1, TimeUnit.SECONDS));
+
+            CompletableFuture<Void> finishFuture = sender.send(1,
+                                                                MessageType.FINISH);
+            allowCompletion.countDown();
+            completionThread.join(TimeUnit.SECONDS.toMillis(1));
+            Assert.assertFalse(completionThread.isAlive());
+            Assert.assertTrue(client.awaitFinish());
+            client.completeFinish();
+            finishFuture.get(1, TimeUnit.SECONDS);
+        } finally {
+            allowCompletion.countDown();
+            if (completionThread != null) {
+                completionThread.join(TimeUnit.SECONDS.toMillis(1));
+            }
+            sender.close();
+        }
+    }
+
+    @Test
+    public void testTransportExceptionCompletesInFlightControlFuture()
+            throws Exception {
+        QueuedMessageSender sender = new QueuedMessageSender(this.config);
+        ControlFutureClient client = new ControlFutureClient();
+        sender.addWorkerClient(1, client);
+        sender.addWorkerClient(2, new MockTransportClient());
+        sender.init();
+
+        try {
+            CompletableFuture<Void> startFuture = sender.send(1,
+                                                               MessageType.START);
+            Assert.assertTrue(client.awaitStart());
+
+            TransportException cause =
+                    new TransportException("connection failed");
+            sender.transportExceptionCaught(cause, client.connectionId());
+            assertFutureFailedWith(startFuture, cause);
+
+            client.completeStart();
+            assertFutureFailedWith(startFuture, cause);
+        } finally {
+            sender.close();
+        }
+    }
+
+    private static void assertFutureFailedWith(CompletableFuture<Void> future,
+                                                Throwable cause)
+            throws InterruptedException, TimeoutException {
+        try {
+            future.get(1, TimeUnit.SECONDS);
+            Assert.fail("Expected control future to fail");
+        } catch (ExecutionException exception) {
+            Assert.assertSame(cause, exception.getCause());
+        }
+    }
+
+    private static class ControlFutureClient extends MockTransportClient {
+
+        private final CountDownLatch startCalled;
+        private final CountDownLatch finishCalled;
+        private final CompletableFuture<Void> startFuture;
+        private final CompletableFuture<Void> finishFuture;
+
+        public ControlFutureClient() {
+            this.startCalled = new CountDownLatch(1);
+            this.finishCalled = new CountDownLatch(1);
+            this.startFuture = new CompletableFuture<>();
+            this.finishFuture = new CompletableFuture<>();
+        }
+
+        @Override
+        public CompletableFuture<Void> startSessionAsync() {
+            this.startCalled.countDown();
+            return this.startFuture;
+        }
+
+        @Override
+        public CompletableFuture<Void> finishSessionAsync() {
+            this.finishCalled.countDown();
+            return this.finishFuture;
+        }
+
+        @Override
+        public boolean send(MessageType messageType, int partition,
+                            ByteBuffer buffer) {
+            return true;
+        }
+
+        @Override
+        public boolean sessionActive() {
+            return false;
+        }
+
+        @Override
+        public InetSocketAddress remoteAddress() {
+            return new InetSocketAddress("127.0.0.1", 8080);
+        }
+
+        public boolean awaitStart() throws InterruptedException {
+            return this.startCalled.await(1, TimeUnit.SECONDS);
+        }
+
+        public boolean awaitFinish() throws InterruptedException {
+            return this.finishCalled.await(1, TimeUnit.SECONDS);
+        }
+
+        public void completeStart() {
+            this.startFuture.complete(null);
+        }
+
+        public void completeFinish() {
+            this.finishFuture.complete(null);
+        }
+    }
 }