shade.polaris.io.grpc.servlet.jakarta.AsyncServletOutputStreamWriter Maven / Gradle / Ivy
Show all versions of polaris-all Show documentation
/*
* Copyright 2019 The gRPC Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.grpc.servlet.jakarta;
import static com.google.common.base.Preconditions.checkState;
import static io.grpc.servlet.jakarta.ServletServerStream.toHexString;
import static java.util.logging.Level.FINE;
import static java.util.logging.Level.FINEST;
import com.google.common.annotations.VisibleForTesting;
import io.grpc.InternalLogId;
import io.grpc.servlet.jakarta.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.function.BiFunction;
import java.util.function.BooleanSupplier;
import java.util.logging.Logger;
import javax.annotation.CheckReturnValue;
import javax.annotation.Nullable;
import jakarta.servlet.AsyncContext;
import jakarta.servlet.ServletOutputStream;
/** Handles write actions from the container thread and the application thread. */
final class AsyncServletOutputStreamWriter {
/**
* Memory boundary for write actions.
*
*
* 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
* }
*
*
* 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 = new AtomicReference<>(WriteState.DEFAULT);
private final Log log;
private final BiFunction writeAction;
private final ActionItem flushAction;
private final ActionItem completeAction;
private final BooleanSupplier isReady;
/**
* 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 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,
ServletTransportState transportState,
InternalLogId logId) throws IOException {
Logger logger = Logger.getLogger(AsyncServletOutputStreamWriter.class.getName());
this.log = new Log() {
@Override
public void fine(String str, Object... params) {
if (logger.isLoggable(FINE)) {
logger.log(FINE, "[" + logId + "]" + str, params);
}
}
@Override
public void finest(String str, Object... params) {
if (logger.isLoggable(FINEST)) {
logger.log(FINEST, "[" + logId + "] " + str, params);
}
}
};
ServletOutputStream outputStream = asyncContext.getResponse().getOutputStream();
this.writeAction = (byte[] bytes, Integer numBytes) -> () -> {
outputStream.write(bytes, 0, numBytes);
transportState.runOnTransportThread(() -> transportState.onSentBytes(numBytes));
log.finest("outbound data: length={0}, bytes={1}", numBytes, toHexString(bytes, numBytes));
};
this.flushAction = () -> {
log.finest("flushBuffer");
asyncContext.getResponse().flushBuffer();
};
this.completeAction = () -> {
log.fine("call is completing");
transportState.runOnTransportThread(
() -> {
transportState.complete();
asyncContext.complete();
log.fine("call completed");
});
};
this.isReady = () -> outputStream.isReady();
}
/**
* Constructor without java.util.logging and jakarta.servlet.* dependency, so that Lincheck can run.
*
* @param writeAction Provides an {@link ActionItem} to write given bytes with specified length.
* @param isReady Indicates whether the writer can write bytes at the moment (asynchronously).
*/
@VisibleForTesting
AsyncServletOutputStreamWriter(
BiFunction writeAction,
ActionItem flushAction,
ActionItem completeAction,
BooleanSupplier isReady,
Log log) {
this.writeAction = writeAction;
this.flushAction = flushAction;
this.completeAction = completeAction;
this.isReady = isReady;
this.log = log;
}
/** Called from application thread. */
void writeBytes(byte[] bytes, int numBytes) throws IOException {
runOrBuffer(writeAction.apply(bytes, numBytes));
}
/** Called from application thread. */
void flush() throws IOException {
runOrBuffer(flushAction);
}
/** Called from application thread. */
void complete() {
try {
runOrBuffer(completeAction);
} catch (IOException ignore) {
// actually completeAction does not throw IOException
}
}
/** Called from the container thread {@link jakarta.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;
}
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");
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
}
parkingThread = null;
}
/**
* Either execute the write action directly, or buffer the action and let the container thread
* drain it.
*
* 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");
}
} 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
}
}
/** Write actions, e.g. writeBytes, flush, complete. */
@FunctionalInterface
@VisibleForTesting
interface ActionItem {
void run() throws IOException;
}
@VisibleForTesting // Lincheck test can not run with java.util.logging dependency.
interface Log {
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.
*
*
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 jakarta.servlet.ServletOutputStream#isReady()} is true.
*
*
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 jakarta.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);
}
}
}