BIgtable: bug fix - Fix Reframer error handling (#2920)

Author: igorbernstein2

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/reframing/ReframingResponseObserver.java

@@ -193,10 +193,23 @@ private void onCancel() {
    */
   @Override
   protected void onResponseImpl(InnerT response) {
+    boolean shoudCancelStream = false;
+
     synchronized (lock) {
       Preconditions.checkState(awaitingInner, "Received unsolicited response from upstream");
       awaitingInner = false;
-      reframer.push(response);
+
+      try {
+        reframer.push(response);
+      } catch (Throwable t) {
+        if (error == null) {
+          shoudCancelStream = true;
+          error = t;
+        }
+      }
+    }
+    if (shoudCancelStream) {
+      innerController.cancel();
     }
     deliver();
   }
@@ -262,6 +275,7 @@ private void deliver() {
 
       if (forceClose) {
         outerResponseObserver.onError(t);
+        innerController.cancel();
       }
     }
   }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/reframing/ReframingResponseObserverTest.java

@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.gaxx.reframing;
 
 import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable.StreamControllerStash;
 import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockResponseObserver;
 import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall;
 import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable;
@@ -34,6 +35,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -311,6 +313,67 @@ public void run() {
     Truth.assertThat(latch.await(1, TimeUnit.MINUTES)).isTrue();
   }
 
+  @Test
+  public void testReframerPushError() throws Exception {
+    MockResponseObserver<String> outerObserver = new MockResponseObserver<>(true);
+    Reframer<String, String> reframer =
+        new DasherizingReframer(1) {
+          @Override
+          public void push(String response) {
+            if (response.equals("boom")) {
+              throw new IllegalStateException("fake error");
+            }
+            super.push(response);
+          }
+        };
+
+    ReframingResponseObserver<String, String> middleware =
+        new ReframingResponseObserver<>(outerObserver, reframer);
+    ServerStreamingStashCallable<String, String> innerCallable =
+        new ServerStreamingStashCallable<>(ImmutableList.of("a", "boom", "c"));
+
+    innerCallable.call("request", middleware);
+
+    Truth.assertThat(outerObserver.getFinalError()).isInstanceOf(IllegalStateException.class);
+    Truth.assertThat(outerObserver.getFinalError()).hasMessage("fake error");
+    Truth.assertThat(outerObserver.popNextResponse()).isEqualTo("a");
+    Truth.assertThat(outerObserver.popNextResponse()).isNull();
+  }
+
+  @Test
+  public void testReframerPopError() {
+    final AtomicInteger popCount = new AtomicInteger();
+
+    MockResponseObserver<String> outerObserver = new MockResponseObserver<>(true);
+    Reframer<String, String> reframer =
+        new DasherizingReframer(1) {
+          @Override
+          public String pop() {
+            if (popCount.incrementAndGet() == 2) {
+              throw new IllegalStateException("fake error");
+            }
+            return super.pop();
+          }
+        };
+
+    ReframingResponseObserver<String, String> middleware =
+        new ReframingResponseObserver<>(outerObserver, reframer);
+    ServerStreamingStashCallable<String, String> innerCallable =
+        new ServerStreamingStashCallable<>(ImmutableList.of("a", "boom", "c"));
+
+    innerCallable.call("request", middleware);
+    StreamControllerStash<String> lastCall = innerCallable.popLastCall();
+
+    Truth.assertThat(outerObserver.getFinalError()).isInstanceOf(IllegalStateException.class);
+    Truth.assertThat(outerObserver.getFinalError()).hasMessage("fake error");
+    Truth.assertThat(outerObserver.popNextResponse()).isEqualTo("a");
+    Truth.assertThat(outerObserver.popNextResponse()).isNull();
+    Truth.assertThat(popCount.get()).isEqualTo(2);
+
+    Truth.assertThat(lastCall.getError()).isInstanceOf(CancellationException.class);
+    Truth.assertThat(lastCall.getNumDelivered()).isEqualTo(2);
+  }
+
   /**
    * A simple implementation of a {@link Reframer}. The input string is split by dash, and the
    * output is concatenated by dashes. The test can verify M:N behavior by adjusting the

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/testing/FakeStreamingApi.java

@@ -30,7 +30,9 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Queue;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.TimeUnit;
 
 @InternalApi("for testing")
 public class FakeStreamingApi {
@@ -118,6 +120,7 @@ public static class ServerStreamingStashCallable<RequestT, ResponseT>
     private ResponseObserver<ResponseT> actualObserver;
     private RequestT actualRequest;
     private List<ResponseT> responseList;
+    private BlockingQueue<StreamControllerStash<ResponseT>> calls = Queues.newLinkedBlockingQueue();
 
     public ServerStreamingStashCallable() {
       responseList = new ArrayList<>();
@@ -138,9 +141,19 @@ public void call(
 
       StreamControllerStash<ResponseT> controller =
           new StreamControllerStash<>(responseList, responseObserver);
+      calls.add(controller);
       controller.start();
     }
 
+    public StreamControllerStash<ResponseT> popLastCall() {
+      try {
+        return calls.poll(1, TimeUnit.SECONDS);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException(e);
+      }
+    }
+
     public ApiCallContext getContext() {
       return context;
     }
@@ -154,11 +167,12 @@ public RequestT getActualRequest() {
     }
 
     // Minimal implementation of back pressure aware stream controller. Not threadsafe
-    private static class StreamControllerStash<ResponseT> implements StreamController {
+    public static class StreamControllerStash<ResponseT> implements StreamController {
       final ResponseObserver<ResponseT> observer;
       final Queue<ResponseT> queue;
       boolean autoFlowControl = true;
       long numPending;
+      long numDelivered;
       Throwable error;
       boolean delivering, closed;
 
@@ -168,6 +182,14 @@ public StreamControllerStash(
         this.queue = Queues.newArrayDeque(responseList);
       }
 
+      public Throwable getError() {
+        return error;
+      }
+
+      public long getNumDelivered() {
+        return numDelivered;
+      }
+
       public void start() {
         observer.onStart(this);
         if (autoFlowControl) {
@@ -200,6 +222,7 @@ private void deliver() {
         try {
           while (error == null && numPending > 0 && !queue.isEmpty()) {
             numPending--;
+            numDelivered++;
             observer.onResponse(queue.poll());
           }