WIP

Author: panchenko

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

@@ -16,70 +16,42 @@
 
 package io.grpc.servlet;
 
-import static com.google.common.base.Preconditions.checkState;
 import static io.grpc.servlet.ServletServerStream.toHexString;
 import static java.util.logging.Level.FINE;
 import static java.util.logging.Level.FINEST;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.errorprone.annotations.CheckReturnValue;
 import io.grpc.InternalLogId;
 import io.grpc.servlet.ServletServerStream.ServletTransportState;
 import java.io.IOException;
-import java.time.Duration;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.concurrent.locks.LockSupport;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
 import java.util.function.BooleanSupplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
 import javax.servlet.AsyncContext;
 import javax.servlet.ServletOutputStream;
+import org.checkerframework.checker.lock.qual.GuardedBy;
 
 /** Handles write actions from the container thread and the application thread. */
 final class AsyncServletOutputStreamWriter {
 
-  /**
-   * Memory boundary for write actions.
-   *
-   * <pre>
-   * WriteState curState = writeState.get();  // mark a boundary
-   * doSomething();  // do something within the boundary
-   * boolean successful = writeState.compareAndSet(curState, newState); // try to mark a boundary
-   * if (successful) {
-   *   // state has not changed since
-   *   return;
-   * } else {
-   *   // state is changed by another thread while doSomething(), need recompute
-   * }
-   * </pre>
-   *
-   * <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.
-   */
-  private final AtomicReference<WriteState> writeState = new AtomicReference<>(WriteState.DEFAULT);
-
   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 if the servlet output stream is not ready or
-   * the queue is not drained.
+   * New write actions will be buffered into this queue if the servlet output stream is not ready
+   * or the queue is not drained.
    */
   // SPSC queue would do
   private final Queue<ActionItem> writeChain = new ConcurrentLinkedQueue<>();
-  // for a theoretical race condition that onWritePossible() is called immediately after isReady()
-  // returns false and before writeState.compareAndSet()
-  @Nullable
-  private volatile Thread parkingThread;
 
   AsyncServletOutputStreamWriter(
       AsyncContext asyncContext,
@@ -173,40 +145,41 @@ 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");
-    assureReadyAndDrainedTurnsFalse();
-    while (isReady.getAsBoolean()) {
-      WriteState curState = writeState.get();
-
-      ActionItem actionItem = writeChain.poll();
-      if (actionItem != null) {
-        actionItem.run();
-        continue;
+    do {
+      writeLock.lock();
+      try {
+        writeFromQueue();
+        if (!outputKnownToBeReady) {
+          log.finest("onWritePossible: EXIT. The servlet output stream becomes not ready");
+          return;
+        }
+      } finally {
+        writeLock.unlock();
       }
+    } while (!writeChain.isEmpty());
+    log.finest("onWritePossible: EXIT. Queue drained");
+  }
 
-      if (writeState.compareAndSet(curState, curState.withReadyAndDrained(true))) {
-        // state has not changed since.
-        log.finest(
-            "onWritePossible: EXIT. All data available now is sent out and the servlet output"
-                + " stream is still ready");
+  @GuardedBy("writeLock")
+  private boolean outputKnownToBeReady = false;
+
+  private void writeFromQueue() throws IOException {
+    for (;;) {
+      ActionItem actionItem;
+      if (outputKnownToBeReady) {
+        actionItem = writeChain.poll();
+      } else if (isReady.getAsBoolean()) {
+        outputKnownToBeReady = true;
+        actionItem = writeChain.poll();
+      } else {
         return;
       }
-      // else, state changed by another thread (runOrBuffer()), need to drain the writeChain
-      // again
-    }
-    log.finest("onWritePossible: EXIT. The servlet output stream becomes not ready");
-  }
-
-  private void assureReadyAndDrainedTurnsFalse() {
-    // readyAndDrained should have been set to false already.
-    // Just in case due to a race condition readyAndDrained is still true at this moment and is
-    // being set to false by runOrBuffer() concurrently.
-    while (writeState.get().readyAndDrained) {
-      parkingThread = Thread.currentThread();
-      // Try to sleep for an extremely long time to avoid writeState being changed at exactly
-      // the time when sleep time expires (in extreme scenario, such as #9917).
-      LockSupport.parkNanos(Duration.ofHours(1).toNanos()); // should return immediately
+      if (actionItem == null) {
+        return;
+      }
+      actionItem.run();
+      outputKnownToBeReady = false;
     }
-    parkingThread = null;
   }
 
   /**
@@ -216,31 +189,13 @@ private void assureReadyAndDrainedTurnsFalse() {
    * <p>Called from application thread.
    */
   private void runOrBuffer(ActionItem actionItem) throws IOException {
-    WriteState curState = writeState.get();
-    if (curState.readyAndDrained) { // write to the outputStream directly
-      actionItem.run();
-      if (actionItem == completeAction) {
-        return;
-      }
-      if (!isReady.getAsBoolean()) {
-        boolean successful =
-            writeState.compareAndSet(curState, curState.withReadyAndDrained(false));
-        LockSupport.unpark(parkingThread);
-        checkState(successful, "Bug: curState is unexpectedly changed by another thread");
-        log.finest("the servlet output stream becomes not ready");
+    writeChain.offer(actionItem);
+    if (writeLock.tryLock()) { // write to the outputStream directly
+      try {
+        writeFromQueue();
+      } finally {
+        writeLock.unlock();
       }
-    } else { // buffer to the writeChain
-      writeChain.offer(actionItem);
-      if (!writeState.compareAndSet(curState, curState.withReadyAndDrained(false))) {
-        checkState(
-            writeState.get().readyAndDrained,
-            "Bug: onWritePossible() should have changed readyAndDrained to true, but not");
-        ActionItem lastItem = writeChain.poll();
-        if (lastItem != null) {
-          checkState(lastItem == actionItem, "Bug: lastItem != actionItem");
-          runOrBuffer(lastItem);
-        }
-      } // state has not changed since
     }
   }
 
@@ -254,43 +209,11 @@ interface ActionItem {
   @VisibleForTesting // Lincheck test can not run with java.util.logging dependency.
   interface Log {
     default boolean isLoggable(Level level) {
-      return false; 
+      return false;
     }
 
     default void fine(String str, Object...params) {}
 
     default void finest(String str, Object...params) {}
   }
-
-  private static final class WriteState {
-
-    static final WriteState DEFAULT = new WriteState(false);
-
-    /**
-     * The servlet output stream is ready and the writeChain is empty.
-     *
-     * <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 writeChain.
-     */
-    final boolean readyAndDrained;
-
-    WriteState(boolean readyAndDrained) {
-      this.readyAndDrained = readyAndDrained;
-    }
-
-    /**
-     * Only {@code onWritePossible()} can set readyAndDrained to true, and only {@code
-     * runOrBuffer()} can set it to false.
-     */
-    @CheckReturnValue
-    WriteState withReadyAndDrained(boolean readyAndDrained) {
-      return new WriteState(readyAndDrained);
-    }
-  }
 }

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

@@ -63,9 +63,9 @@ public class AsyncServletOutputStreamWriterConcurrencyTest extends VerifierState
   private volatile boolean isReady;
   // when isReadyReturnedFalse, writer.onWritePossible() will be called.
   private volatile boolean isReadyReturnedFalse;
-  private int producerIndex;
-  private int consumerIndex;
-  private int bytesWritten;
+  private volatile int producerIndex;
+  private volatile int consumerIndex;
+  private volatile int bytesWritten;
 
   /** Public no-args constructor. */
   public AsyncServletOutputStreamWriterConcurrencyTest() {