Fix request stream recreation on reconnect

iamborisov · May 8, 2020 · ed81b5a · ed81b5a
1 parent d389406
commit ed81b5a
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 22 deletions.
diff --git a/src/main/java/org/camunda/bpm/engine/grpc/client/request/RequestObserver.java b/src/main/java/org/camunda/bpm/engine/grpc/client/request/RequestObserver.java
@@ -6,4 +6,6 @@
 public interface RequestObserver {
 
     StreamObserver<FetchAndLockRequest> getStreamObserver();
+
+    void reset();
 }
diff --git a/src/main/java/org/camunda/bpm/engine/grpc/client/request/impl/RequestObserverImpl.java b/src/main/java/org/camunda/bpm/engine/grpc/client/request/impl/RequestObserverImpl.java
@@ -32,4 +32,9 @@ public StreamObserver<FetchAndLockRequest> getStreamObserver() {
     private StreamObserver<FetchAndLockRequest> buildStreamObserver() {
         return stub.getStub().fetchAndLock(responseObserver);
     }
+
+    @Override
+    public void reset() {
+        streamObserver = null;
+    }
 }
diff --git a/src/main/java/org/camunda/bpm/engine/grpc/client/worker/impl/HandlerImpl.java b/src/main/java/org/camunda/bpm/engine/grpc/client/worker/impl/HandlerImpl.java
@@ -7,6 +7,7 @@
 import org.camunda.bpm.engine.grpc.client.subscription.SubscriptionRepository;
 import org.camunda.bpm.engine.grpc.client.worker.Handler;
 import org.camunda.bpm.engine.grpc.client.worker.Locker;
+import org.camunda.bpm.engine.grpc.client.worker.Watchdog;
 import org.springframework.stereotype.Component;
 
 @Slf4j
@@ -16,6 +17,8 @@ public class HandlerImpl implements Handler {
 
     private final Locker locker;
 
+    private final Watchdog watchdog;
+
     private final SubscriptionRepository subscriptionRepository;
 
     private final RequestObserver requestObserver;
@@ -24,25 +27,47 @@ public class HandlerImpl implements Handler {
 
     @Override
     public void handle() {
-        if (!subscriptionRepository.isEmpty()) {
-            requestObserver.getStreamObserver().onNext(
-                requestFactory.create()
-            );
-
-            locker.lock();
-        } else {
-            log.info("There is no registered external task subscriptions");
-
-            try {
-                Thread.sleep(500);
-            } catch (InterruptedException e) {
-                log.error("Thread interrupted", e);
-            }
+        if (!handleConnection() || !handleSubscriptions()) {
+            sleep();
         }
     }
 
     @Override
     public void complete() {
         requestObserver.getStreamObserver().onCompleted();
     }
+
+    private boolean handleConnection() {
+        if (!watchdog.watch()) {
+            requestObserver.reset();
+
+            return false;
+        }
+
+        return true;
+    }
+
+    private boolean handleSubscriptions() {
+        if (subscriptionRepository.isEmpty()) {
+            log.info("There is no registered external task subscriptions");
+
+            return false;
+        }
+
+        requestObserver.getStreamObserver().onNext(
+            requestFactory.create()
+        );
+
+        locker.lock();
+
+        return true;
+    }
+
+    private void sleep() {
+        try {
+            Thread.sleep(1000);
+        } catch (InterruptedException e) {
+            log.error("Thread sleep interrupted", e);
+        }
+    }
 }
diff --git a/src/main/java/org/camunda/bpm/engine/grpc/client/worker/impl/WorkerImpl.java b/src/main/java/org/camunda/bpm/engine/grpc/client/worker/impl/WorkerImpl.java
@@ -3,7 +3,6 @@
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.camunda.bpm.engine.grpc.client.worker.Handler;
-import org.camunda.bpm.engine.grpc.client.worker.Watchdog;
 import org.camunda.bpm.engine.grpc.client.worker.Worker;
 import org.springframework.stereotype.Component;
 
@@ -16,8 +15,6 @@ public class WorkerImpl implements Worker {
 
     private final Handler handler;
 
-    private final Watchdog watchdog;
-
     private AtomicBoolean isRunning = new AtomicBoolean(false);
 
     private Thread thread = new Thread(this, getClass().getSimpleName());
@@ -54,11 +51,7 @@ public void stop() throws AlreadyStoppedException {
     public void run() {
         while (isRunning.get()) {
             try {
-                if (watchdog.watch()) {
-                    handler.handle();
-                } else {
-                    Thread.sleep(1000);
-                }
+                handler.handle();
             } catch (Throwable e) {
                 log.error("Exception while handling worker thread", e);
             }