Merge pull request #2643 from opencb/TASK-8158

TASK-8158 - GRPC doesn't stop streaming on cancelled requests.

Author: j-coll

opencga-server/src/main/java/org/opencb/opencga/server/grpc/GenericGrpcService.java

@@ -20,6 +20,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.ObjectWriter;
 import io.grpc.MethodDescriptor;
+import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 import org.apache.commons.lang3.time.StopWatch;
 import org.opencb.commons.datastore.core.Query;
@@ -183,39 +184,50 @@ public void run(MethodDescriptor<Request, ?> method, Request request, StreamObse
 
         Exception e = null;
         int numResults = -1;
+        boolean cancelled = false;
         try {
             numResults = requestRunner.run(query, queryOptions);
+            cancelled = isCancelled(streamObserver);
+            if (cancelled) {
+                logger.warn("Request cancelled by client");
+            }
             streamObserver.onCompleted();
         } catch (Exception ex) {
             e = ex;
             logger.error("Catch error: " + e.getMessage(), e);
             streamObserver.onError(ex);
         } finally {
             stopWatch.stop();
-            logEnd(e, stopWatch, numResults, requestDescription);
+            logEnd(e, cancelled, stopWatch, numResults, requestDescription);
         }
 
     }
 
-    private void logEnd(Exception e, StopWatch stopWatch, int numResults, String requestDescription) {
+    protected boolean isCancelled(StreamObserver<?> streamObserver) {
+        return streamObserver instanceof ServerCallStreamObserver
+                && ((ServerCallStreamObserver<?>) streamObserver).isCancelled();
+    }
+
+    private void logEnd(Exception e, boolean cancelled, StopWatch stopWatch, int numResults, String requestDescription) {
         StringBuilder sb = new StringBuilder();
-        boolean ok;
-        if (e == null) {
-            sb.append("OK");
-            ok = true;
-        } else {
+        if (e != null) {
             sb.append("ERROR");
-            ok = false;
+        } else if (cancelled) {
+            sb.append("CANCELLED");
+        } else {
+            sb.append("OK");
         }
         sb.append(", ").append(stopWatch.getTime(TimeUnit.MILLISECONDS)).append("ms");
         if (numResults >= 0) {
             sb.append(", ").append(numResults).append(" results");
         }
         sb.append(", ").append(requestDescription);
-        if (ok) {
-            logger.info(sb.toString());
-        } else {
+        if (e != null) {
             logger.error(sb.toString());
+        } else if (cancelled) {
+            logger.warn(sb.toString());
+        } else {
+            logger.info(sb.toString());
         }
     }
 }

opencga-server/src/main/java/org/opencb/opencga/server/grpc/VariantGrpcService.java

@@ -91,7 +91,7 @@ public void query(Request request, StreamObserver<VariantResponse> responseObser
                         events.add(eventB.build());
                     }
                 }
-                while (iterator.hasNext()) {
+                while (iterator.hasNext() && !genericGrpcService.isCancelled(responseObserver)) {
                     Variant variant = iterator.next();
                     VariantProto.Variant variantProto = converter.convert(variant);
                     VariantResponse.Builder responseBuilder = VariantResponse.newBuilder();