From 86ccdd5a559b7c62f613940c347a184c8a09c28a Mon Sep 17 00:00:00 2001
From: Vaci Koblizek <vaci@vaci.org>
Date: Sat, 24 Oct 2020 16:51:16 +0100
Subject: [PATCH] move cleanup to end of message loop

---
 .../src/main/java/org/capnproto/RpcState.java   | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/runtime/src/main/java/org/capnproto/RpcState.java b/runtime/src/main/java/org/capnproto/RpcState.java
index 1572d60..becaed4 100644
--- a/runtime/src/main/java/org/capnproto/RpcState.java
+++ b/runtime/src/main/java/org/capnproto/RpcState.java
@@ -282,9 +282,9 @@ final class RpcState {
     private final ReferenceQueue<Question> questionRefs = new ReferenceQueue<>();
     private final ReferenceQueue<ImportClient> importRefs = new ReferenceQueue<>();
 
-    RpcState( Capability.Client bootstrapInterface,
-              VatNetwork.Connection connection,
-              CompletableFuture<java.lang.Void> onDisconnect) {
+    RpcState(Capability.Client bootstrapInterface,
+             VatNetwork.Connection connection,
+             CompletableFuture<java.lang.Void> onDisconnect) {
         this.bootstrapInterface = bootstrapInterface;
         this.connection = connection;
         this.onDisconnect = onDisconnect;
@@ -295,6 +295,10 @@ final class RpcState {
         return this.messageLoop;
     }
 
+    public CompletableFuture<java.lang.Void> onDisconnect() {
+        return this.messageLoop;
+    }
+
     CompletableFuture<java.lang.Void> disconnect(Throwable exc) {
         if (isDisconnected()) {
             return CompletableFuture.failedFuture(this.disconnected);
@@ -420,21 +424,20 @@ final class RpcState {
     }
 
     private CompletableFuture<java.lang.Void> doMessageLoop() {
-        this.cleanupImports();
-        this.cleanupQuestions();
-
         if (isDisconnected()) {
             return CompletableFuture.failedFuture(this.disconnected);
         }
 
         return connection.receiveIncomingMessage().thenCompose(message -> {
             try {
-                handleMessage(message);
+                this.handleMessage(message);
             } catch (Exception rpcExc) {
                 // either we received an Abort message from peer
                 // or internal RpcState is bad.
                 return this.disconnect(rpcExc);
             }
+            this.cleanupImports();
+            this.cleanupQuestions();
             return this.doMessageLoop();
 
         }).exceptionallyCompose(exc -> this.disconnect(exc));