lokidundun commented on code in PR #357:
URL: 
https://github.com/apache/hugegraph-computer/pull/357#discussion_r3674430026


##########
computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSender.java:
##########
@@ -227,75 +237,127 @@ 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(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(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);

Review Comment:
   Done



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to