From 5613baf066dc56ee85709c10d757ce8799c4f425 Mon Sep 17 00:00:00 2001
From: Volodymyr Babak <volodymyr.babak@gmail.com>
Date: Fri, 25 Sep 2020 16:52:56 +0300
Subject: [PATCH] Added correct onComplete in case onError

---
 .../java/org/thingsboard/edge/rpc/EdgeGrpcClient.java  | 10 ++++++++--
 .../java/org/thingsboard/edge/rpc/EdgeRpcClient.java   |  2 +-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeGrpcClient.java b/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeGrpcClient.java
index 4c20e97b6a..88fe402668 100644
--- a/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeGrpcClient.java
+++ b/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeGrpcClient.java
@@ -40,6 +40,7 @@ import org.thingsboard.server.gen.edge.UplinkResponseMsg;
 import javax.net.ssl.SSLException;
 import java.io.File;
 import java.net.URISyntaxException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
@@ -107,7 +108,7 @@ public class EdgeGrpcClient implements EdgeRpcClient {
                     } else {
                         log.error("[{}] Failed to establish the connection! Code: {}. Error message: {}.", edgeKey, connectResponseMsg.getResponseCode(), connectResponseMsg.getErrorMsg());
                         try {
-                            EdgeGrpcClient.this.disconnect();
+                            EdgeGrpcClient.this.disconnect(true);
                         } catch (InterruptedException e) {
                             log.error("[{}] Got interruption during disconnect!", edgeKey, e);
                         }
@@ -136,7 +137,12 @@ public class EdgeGrpcClient implements EdgeRpcClient {
     }
 
     @Override
-    public void disconnect() throws InterruptedException {
+    public void disconnect(boolean onError) throws InterruptedException {
+        if (!onError) {
+            try {
+                inputStream.onCompleted();
+            } catch (Exception ignored) {}
+        }
         if (channel != null) {
             channel.shutdown().awaitTermination(timeoutSecs, TimeUnit.SECONDS);
         }
diff --git a/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeRpcClient.java b/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeRpcClient.java
index a66f941922..7aa017c3db 100644
--- a/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeRpcClient.java
+++ b/common/edge-api/src/main/java/org/thingsboard/edge/rpc/EdgeRpcClient.java
@@ -32,7 +32,7 @@ public interface EdgeRpcClient {
                  Consumer<DownlinkMsg> onDownlink,
                  Consumer<Exception> onError);
 
-    void disconnect() throws InterruptedException;
+    void disconnect(boolean onError) throws InterruptedException;
 
     void sendSyncRequestMsg();