Do not hold lock while sending requests, and some minor refactoring.

yayi-google · yayi-google · commit b98d6be0a648 · 2021-02-22T12:24:15.000-08:00
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterV2.java
@@ -23,7 +23,6 @@
 import com.google.common.util.concurrent.Uninterruptibles;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
-import java.time.Duration;
 import java.util.Deque;
 import java.util.LinkedList;
 import java.util.concurrent.TimeUnit;
@@ -53,8 +52,6 @@
 public class StreamWriterV2 implements AutoCloseable {
   private static final Logger log = Logger.getLogger(StreamWriterV2.class.getName());
 
-  private static final Duration DONE_CALLBACK_WAIT_TIMEOUT = Duration.ofMinutes(10);
-
   private Lock lock;
   private Condition hasMessageInWaitingQueue;
 
@@ -104,23 +101,28 @@ private StreamWriterV2(Builder builder) {
     this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.streamConnection =
-        new StreamConnection(builder.client, new RequestCallback() {
-          @Override
-          public void run(AppendRowsResponse response) {
-            requestCallback(response);
-          }
-        }, new DoneCallback() {
-          @Override
-          public void run(Throwable finalStatus) {
-            doneCallback(finalStatus);
-          }
-        });
-    this.appendThread = new Thread(new Runnable() {
-      @Override
-      public void run() {
-        appendLoop();
-      }
-    });
+        new StreamConnection(
+            builder.client,
+            new RequestCallback() {
+              @Override
+              public void run(AppendRowsResponse response) {
+                requestCallback(response);
+              }
+            },
+            new DoneCallback() {
+              @Override
+              public void run(Throwable finalStatus) {
+                doneCallback(finalStatus);
+              }
+            });
+    this.appendThread =
+        new Thread(
+            new Runnable() {
+              @Override
+              public void run() {
+                appendLoop();
+              }
+            });
     this.appendThread.start();
   }
 
@@ -210,10 +212,16 @@ private void appendLoop() {
       try {
         hasMessageInWaitingQueue.await(100, TimeUnit.MILLISECONDS);
         while (!this.waitingRequestQueue.isEmpty()) {
-          localQueue.addLast(this.waitingRequestQueue.pollFirst());
+          AppendRequestAndResponse requestWrapper = this.waitingRequestQueue.pollFirst();
+          this.inflightRequestQueue.addLast(requestWrapper);
+          localQueue.addLast(requestWrapper);
         }
       } catch (InterruptedException e) {
-        log.warning("Interrupted while waiting for message. Error: " + e.toString());
+        log.warning(
+            "Interrupted while waiting for message. Stream: "
+                + streamName
+                + " Error: "
+                + e.toString());
       } finally {
         this.lock.unlock();
       }
@@ -223,43 +231,16 @@ private void appendLoop() {
       }
 
       // TODO: Add reconnection here.
-
-      this.lock.lock();
-      try {
-        while (!localQueue.isEmpty()) {
-          AppendRequestAndResponse requestWrapper = localQueue.pollFirst();
-          this.inflightRequestQueue.addLast(requestWrapper);
-          this.streamConnection.send(requestWrapper.message);
-        }
-      } finally {
-        this.lock.unlock();
+      while (!localQueue.isEmpty()) {
+        this.streamConnection.send(localQueue.pollFirst().message);
       }
     }
 
     log.info("Cleanup starts. Stream: " + streamName);
     // At this point, the waiting queue is drained, so no more requests.
     // We can close the stream connection and handle the remaining inflight requests.
     this.streamConnection.close();
-
-    log.info("Waiting for done callback from stream connection. Stream: " + streamName);
-    long waitDeadlineMs = System.currentTimeMillis() + DONE_CALLBACK_WAIT_TIMEOUT.toMillis();
-    while (true) {
-      if (System.currentTimeMillis() > waitDeadlineMs) {
-        log.warning(
-            "Timeout waiting for done wallback. Skip inflight cleanup. Stream: " + streamName);
-        return;
-      }
-      this.lock.lock();
-      try {
-        if (connectionFinalStatus != null) {
-          // Done callback is received, break.
-          break;
-        }
-      } finally {
-        this.lock.unlock();
-      }
-      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
-    }
+    waitForDoneCallback();
 
     // At this point, there cannot be more callback. It is safe to clean up all inflight requests.
     log.info(
@@ -284,6 +265,22 @@ private boolean waitingQueueDrained() {
     }
   }
 
+  private void waitForDoneCallback() {
+    log.info("Waiting for done callback from stream connection. Stream: " + streamName);
+    while (true) {
+      this.lock.lock();
+      try {
+        if (connectionFinalStatus != null) {
+          // Done callback is received, return.
+          return;
+        }
+      } finally {
+        this.lock.unlock();
+      }
+      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
+    }
+  }
+
   private void cleanupInflightRequests() {
     Throwable finalStatus;
     Deque<AppendRequestAndResponse> localQueue = new LinkedList<AppendRequestAndResponse>();
