CR

Author: panchenko

servlet/src/main/java/io/grpc/servlet/AsyncServletOutputStreamWriter.java

@@ -26,8 +26,7 @@
 import java.io.IOException;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.locks.StampedLock;
 import java.util.function.BiFunction;
 import java.util.function.BooleanSupplier;
 import java.util.logging.Level;
@@ -38,13 +37,34 @@
 /** Handles write actions from the container thread and the application thread. */
 final class AsyncServletOutputStreamWriter {
 
+  private final StampedLock writeLock = new StampedLock();
+
+  /**
+   * The servlet output stream is ready, and the writeQueue is empty.
+   *
+   * <p>There are two threads, the container thread (calling {@code onWritePossible()}) and the
+   * application thread (calling {@code runOrBuffer()}) that read and update the
+   * writeState. Only onWritePossible() may turn {@code readyAndDrained} from false to true, and
+   * only runOrBuffer() may turn it from true to false.
+   *
+   * <p>readyAndDrained turns from false to true when:
+   * {@code onWritePossible()} exits while currently there is no more data to write, but the last
+   * check of {@link javax.servlet.ServletOutputStream#isReady()} is true.
+   *
+   * <p>readyAndDrained turns from true to false when:
+   * {@code runOrBuffer()} exits while either the action item is written directly to the
+   * servlet output stream and the check of {@link javax.servlet.ServletOutputStream#isReady()}
+   * right after that returns false, or the action item is buffered into the writeQueue.
+   */
+  // @GuardedBy("writeLock")
+  private boolean readyAndDrained;
+
   private final Log log;
   private final BiFunction<byte[], Integer, ActionItem> writeAction;
   private final ActionItem flushAction;
   private final ActionItem completeAction;
   private final BooleanSupplier isReady;
 
-  private final Lock writeLock = new ReentrantLock();
   /**
    * New write actions will be buffered into this queue.
    */
@@ -143,35 +163,21 @@ void complete() {
   /** Called from the container thread {@link javax.servlet.WriteListener#onWritePossible()}. */
   void onWritePossible() throws IOException {
     log.finest("onWritePossible: ENTRY. The servlet output stream becomes ready");
-    do {
-      writeLock.lock();
-      try {
-        if (writeFromQueue() == WriteResult.OUTPUT_NOT_READY) {
-          log.finest("onWritePossible: EXIT. The servlet output stream becomes not ready");
+    long stamp = writeLock.writeLock();
+    try {
+      while (isReady.getAsBoolean()) {
+        ActionItem actionItem = writeQueue.poll();
+        if (actionItem == null) {
+          readyAndDrained = true;
+          log.finest("onWritePossible: EXIT. Queue drained");
           return;
         }
-      } finally {
-        writeLock.unlock();
+        actionItem.run();
       }
-      // retry if runOrBuffer() added tasks, but have not acquired the lock
-    } while (!writeQueue.isEmpty());
-    log.finest("onWritePossible: EXIT. Queue drained");
-  }
-
-  private enum WriteResult {
-    OUTPUT_NOT_READY,
-    QUEUE_DRAINED
-  }
-
-  private WriteResult writeFromQueue() throws IOException {
-    while (isReady.getAsBoolean()) {
-      ActionItem actionItem = writeQueue.poll();
-      if (actionItem == null) {
-        return WriteResult.QUEUE_DRAINED;
-      }
-      actionItem.run();
+      log.finest("onWritePossible: EXIT. The servlet output stream becomes not ready");
+    } finally {
+      writeLock.unlockWrite(stamp);
     }
-    return WriteResult.OUTPUT_NOT_READY;
   }
 
   /**
@@ -182,12 +188,25 @@ private WriteResult writeFromQueue() throws IOException {
    */
   private void runOrBuffer(ActionItem actionItem) throws IOException {
     writeQueue.offer(actionItem);
-    if (writeLock.tryLock()) { // write to the outputStream directly
-      try {
-        writeFromQueue();
-      } finally {
-        writeLock.unlock();
+    long stamp = writeLock.tryWriteLock();
+    if (stamp == 0L) {
+      return;
+    }
+    try {
+      if (readyAndDrained) { // write to the outputStream directly
+        ActionItem toWrite = writeQueue.poll();
+        if (toWrite != null) {
+          toWrite.run();
+          if (toWrite == completeAction) {
+            return;
+          }
+          if (!isReady.getAsBoolean()) {
+            readyAndDrained = false;
+          }
+        }
       }
+    } finally {
+      writeLock.unlockWrite(stamp);
     }
   }
 

servlet/src/threadingTest/java/io/grpc/servlet/AsyncServletOutputStreamWriterConcurrencyTest.java

@@ -103,8 +103,8 @@ private void writeOrFlush() {
 
   private boolean isReady() {
     if (!isReady) {
-      //assertWithMessage("isReady() already returned false, onWritePossible() will be invoked")
-      //.that(isReadyReturnedFalse).isFalse();
+      assertWithMessage("isReady() already returned false, onWritePossible() will be invoked")
+          .that(isReadyReturnedFalse).isFalse();
       isReadyReturnedFalse = true;
     }
     return isReady;