modify finishSend

Author: Radeity

computer-core/src/main/java/org/apache/hugegraph/computer/core/compute/ComputeManager.java

@@ -168,7 +168,7 @@ public WorkerStat compute(WorkerContext context, int superstep) {
                                         "partition parallel compute", t);
         }
 
-        this.sendManager.finishSend(MessageType.MSG);
+        this.sendManager.finishSend(MessageType.MSG, true);
 
         // After compute and send finish signal.
         Map<Integer, MessageStat> recvStats = this.recvManager.messageStats();

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

@@ -81,15 +81,15 @@ public void loadGraph() {
             Vertex vertex = iterator.next();
             this.sendManager.sendVertex(vertex);
         }
-        this.sendManager.finishSend(MessageType.VERTEX);
+        this.sendManager.finishSend(MessageType.VERTEX, false);
 
         this.sendManager.startSend(MessageType.EDGE);
         iterator = this.loadService.createIteratorFromEdge();
         while (iterator.hasNext()) {
             Vertex vertex = iterator.next();
             this.sendManager.sendEdge(vertex);
         }
-        this.sendManager.finishSend(MessageType.EDGE);
+        this.sendManager.finishSend(MessageType.EDGE, true);
         this.sendManager.clearBuffer();
     }
 }

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

@@ -118,9 +118,9 @@ public void onChannelInactive(ConnectionId connectionId) {
         @Override
         public void exceptionCaught(TransportException cause,
                                     ConnectionId connectionId) {
-            LOG.error("Channel for connectionId {} occurred exception, shutdown all clients",
+            LOG.error("Channel for connectionId {} occurred exception",
                       connectionId, cause);
-            DataClientManager.this.sender.exceptionCaught(cause, connectionId);
+            DataClientManager.this.sender.transportExceptionCaught(cause, connectionId);
         }
     }
 }

computer-core/src/main/java/org/apache/hugegraph/computer/core/receiver/MessageRecvManager.java

@@ -63,7 +63,6 @@ public class MessageRecvManager implements Manager, MessageHandler {
     private int expectedFinishMessages;
     private CompletableFuture<Void> finishAllMessagesFuture;
     private CompletableFuture<Boolean>[] finishMessageFutures;
-    private AtomicInteger finishMessageCount;
 
     private long waitFinishMessagesTimeout;
     private long superstep;
@@ -92,10 +91,14 @@ public void init(Config config) {
         this.edgePartitions = new EdgeMessageRecvPartitions(
                               this.context, fileGenerator, this.sortManager);
         this.workerCount = config.get(ComputerOptions.JOB_WORKERS_COUNT);
-        // One for vertex and one for edge.
-        this.expectedFinishMessages = this.workerCount * 2;
 
-        initFinishMessageFuture();
+        this.expectedFinishMessages = this.workerCount;
+        this.finishMessageFutures = new CompletableFuture[this.expectedFinishMessages];
+
+        for (int i = 0; i < this.expectedFinishMessages; i++) {
+            this.finishMessageFutures[i] = new CompletableFuture<>();
+        }
+        this.finishAllMessagesFuture = CompletableFuture.allOf(finishMessageFutures);
 
         this.waitFinishMessagesTimeout = config.get(
              ComputerOptions.WORKER_WAIT_FINISH_MESSAGES_TIMEOUT);
@@ -107,9 +110,12 @@ public void beforeSuperstep(Config config, int superstep) {
                                                this.fileManager, superstep);
         this.messagePartitions = new ComputeMessageRecvPartitions(
                                  this.context, fileGenerator, this.sortManager);
-        this.expectedFinishMessages = this.workerCount;
 
-        initFinishMessageFuture();
+        for (int i = 0; i < this.expectedFinishMessages; i++) {
+            this.finishMessageFutures[i] = new CompletableFuture<>();
+        }
+        this.finishAllMessagesFuture = CompletableFuture.allOf(finishMessageFutures);
+
 
         this.superstep = superstep;
 
@@ -151,15 +157,6 @@ public void exceptionCaught(TransportException cause,
         }
     }
 
-    private void initFinishMessageFuture() {
-        this.finishMessageFutures = new CompletableFuture[this.expectedFinishMessages];
-        for (int i = 0; i < this.expectedFinishMessages; i++) {
-            this.finishMessageFutures[i] = new CompletableFuture<>();
-        }
-        this.finishAllMessagesFuture = CompletableFuture.allOf(finishMessageFutures);
-        this.finishMessageCount = new AtomicInteger(0);
-    }
-
     public void waitReceivedAllMessages() {
         if (finishAllMessagesFuture != null) {
             try {
@@ -222,8 +219,7 @@ public void onStarted(ConnectionId connectionId) {
     @Override
     public void onFinished(ConnectionId connectionId) {
         LOG.debug("ConnectionId {} finished", connectionId);
-        int messageIdx = this.finishMessageCount.getAndIncrement();
-        this.finishMessageFutures[messageIdx].complete(true);
+        this.finishMessageFutures[connectionId.clientIndex()].complete(true);
     }
 
     /**

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

@@ -158,14 +158,16 @@ public void startSend(MessageType type) {
      * into queue
      * @param type the message type
      */
-    public void finishSend(MessageType type) {
+    public void finishSend(MessageType type, boolean sendControlMessage) {
         Map<Integer, MessageSendPartition> all = this.buffers.all();
         MessageStat stat = this.sortAndSendLastBuffer(all, type);
 
         Set<Integer> workerIds = all.keySet().stream()
                                     .map(this.partitioner::workerId)
                                     .collect(Collectors.toSet());
-        this.sendControlMessageToWorkers(workerIds, MessageType.FINISH);
+        if (sendControlMessage == true) {
+            this.sendControlMessageToWorkers(workerIds, MessageType.FINISH);
+        }
         LOG.info("Finish sending message(type={},count={},bytes={})",
                  type, stat.messageCount(), stat.messageBytes());
     }

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

@@ -44,5 +44,5 @@ CompletableFuture<Void> send(int workerId, MessageType type)
      * Invoked when the channel associated with the given connectionId has
      * an exception is thrown processing message.
      */
-    void exceptionCaught(TransportException cause, ConnectionId connectionId);
+    void transportExceptionCaught(TransportException cause, ConnectionId connectionId);
 }

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

@@ -105,7 +105,7 @@ public void send(int workerId, QueuedMessage message)
     }
 
     @Override
-    public void exceptionCaught(TransportException cause, ConnectionId connectionId) {
+    public void transportExceptionCaught(TransportException cause, ConnectionId connectionId) {
         for (WorkerChannel channel : this.channels) {
             if (channel.client.connectionId().equals(connectionId)) {
                 channel.futureRef.get().completeExceptionally(cause);

computer-test/src/main/java/org/apache/hugegraph/computer/core/compute/MockMessageSender.java

@@ -40,7 +40,7 @@ public void send(int workerId, QueuedMessage message) {
     }
 
     @Override
-    public void exceptionCaught(TransportException cause, ConnectionId connectionId) {
+    public void transportExceptionCaught(TransportException cause, ConnectionId connectionId) {
         // pass
     }
 }