com.github.netty.protocol.servlet.ServletOutputStream Maven / Gradle / Ivy
The newest version!
package com.github.netty.protocol.servlet;
import com.github.netty.core.util.*;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.handler.codec.http.DefaultLastHttpContent;
import io.netty.handler.codec.http.LastHttpContent;
import io.netty.handler.stream.ChunkedFile;
import io.netty.handler.stream.ChunkedInput;
import io.netty.util.internal.PlatformDependent;
import javax.servlet.WriteListener;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
/**
* Servlet OutputStream
*
* @author wangzihao
*/
public class ServletOutputStream extends javax.servlet.ServletOutputStream implements Recyclable, NettyOutputStream {
public static final ServletResetBufferIOException RESET_BUFFER_EXCEPTION = new ServletResetBufferIOException();
private static final Recycler RECYCLER = new Recycler<>(ServletOutputStream::new);
public static int h2ChunkSize = 81920;
protected final AtomicLong writeBytes = new AtomicLong();
protected final AtomicBoolean isClosed = new AtomicBoolean(false);
protected final AtomicBoolean isSendResponse = new AtomicBoolean(false);
private final CloseListener closeListenerWrapper = new CloseListener();
protected ServletHttpExchange servletHttpExchange;
protected WriteListener writeListener;
protected ChannelProgressivePromise lastContentPromise;
private ChannelProgressivePromise blockPromise;
protected ServletOutputStream() {
}
public static ServletOutputStream newInstance(ServletHttpExchange servletHttpExchange) {
ServletOutputStream instance = RECYCLER.getInstance();
instance.blockPromise = null;
instance.setServletHttpExchange(servletHttpExchange);
instance.writeBytes.set(0);
instance.isSendResponse.set(false);
instance.isClosed.set(false);
return instance;
}
public long getWriteBytes() {
return writeBytes.get();
}
@Override
public ChannelProgressivePromise write(ByteBuffer httpBody) throws IOException {
ByteBuf byteBuf = Unpooled.wrappedBuffer(httpBody);
return writeHttpBody(byteBuf, byteBuf.readableBytes());
}
@Override
public ChannelProgressivePromise write(ByteBuf httpBody) throws IOException {
IOUtil.writerModeToReadMode(httpBody);
return writeHttpBody(httpBody, httpBody.readableBytes());
}
@Override
public ChannelProgressivePromise write(ChunkedInput input) throws IOException {
return writeHttpBody(input, input.length());
}
@Override
public ChannelProgressivePromise write(File file, long position, long count) throws IOException {
if (isHttp2()) {
return writeHttpBody(new ChunkedFile(new RandomAccessFile(file, "r"), position, count, h2ChunkSize), count);
} else {
return writeHttpBody(new DefaultFileRegion(file, position, count), count);
}
}
@Override
public ChannelProgressivePromise write(File httpBody) throws IOException {
long length = httpBody.length();
if (isHttp2()) {
return writeHttpBody(new ChunkedFile(new RandomAccessFile(httpBody, "r"), 0, length, h2ChunkSize), length);
} else {
return writeHttpBody(new DefaultFileRegion(httpBody, 0, length), length);
}
}
public boolean isHttp2() {
return servletHttpExchange.getProtocol().isHttp2();
}
protected ChannelProgressivePromise writeHttpBody(Object httpBody, long length) throws IOException {
if (httpBody instanceof FileRegion) {
servletHttpExchange.getResponse().getNettyResponse().setWriteSendFile(true);
}
try {
checkClosed();
writeResponseHeaderIfNeed();
ServletHttpExchange servletHttpExchange = this.servletHttpExchange;
ChannelHandlerContext context = servletHttpExchange.getChannelHandlerContext();
ChannelProgressivePromise promise = context.newProgressivePromise();
if (length > 0) {
writeBytes.addAndGet(length);
}
long contentLength = servletHttpExchange.getResponse().getContentLength();
// response finish
if (contentLength >= 0 && writeBytes.get() >= contentLength) {
boolean autoFlush = servletHttpExchange.getServletContext().isAutoFlush();
if (httpBody instanceof ByteBuf) {
DefaultLastHttpContent httpContent = new DefaultLastHttpContent((ByteBuf) httpBody, false);
if (autoFlush) {
context.write(httpContent, promise);
} else {
context.writeAndFlush(httpContent, promise);
}
} else {
context.write(httpBody);
if (autoFlush) {
context.write(LastHttpContent.EMPTY_LAST_CONTENT, promise);
} else {
context.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT, promise);
}
}
lastContentPromise = promise;
} else {
// Response continues
context.write(httpBody, promise);
}
// limiting control. avoid buffer memory overflow.
blockIfNeed(promise);
return promise;
} catch (Exception e) {
RecyclableUtil.release(httpBody);
throw e;
}
}
/**
* limiting control. avoid buffer memory overflow.
*
* @param promise last write promise
* @throws IOException wrap ClosedChannelException or other channel exception
*/
private void blockIfNeed(ChannelProgressivePromise promise) throws IOException {
ServletHttpExchange exchange = this.servletHttpExchange;
ChannelHandlerContext context = exchange.getChannelHandlerContext();
long pendingWriteBytes = exchange.getPendingWriteBytes();
if (pendingWriteBytes <= 0) {
return;
}
if (context.executor().inEventLoop()) {
// 1 time slices
Thread.yield();
context.flush();
Thread.yield();
ChannelUtils.forceFlush(context.channel());
} else {
int bufferSize = exchange.getResponse().getBufferSize();
ChannelProgressivePromise blockPromise = this.blockPromise;
boolean requiresFlush = true;
if (pendingWriteBytes >= bufferSize && blockPromise == null) {
context.flush();
requiresFlush = false;
// record promise
this.blockPromise = blockPromise = promise;
}
// Over the double buffer size, block the processing
int doubleBufferSize = bufferSize << 1;
if (pendingWriteBytes >= doubleBufferSize) {
try {
if (blockPromise == null) {
blockPromise = promise;
}
if (requiresFlush) {
context.flush();
}
blockPromise.sync();
} catch (InterruptedException ignored) {
} catch (Exception e) {
throw new IOException("flush fail = " + e, e);
} finally {
this.blockPromise = null;
}
}
}
}
private void writeResponseHeaderIfNeed() {
if (isSendResponse.compareAndSet(false, true)) {
ServletHttpServletResponse servletResponse = servletHttpExchange.getResponse();
NettyHttpResponse nettyResponse = servletResponse.getNettyResponse();
ChannelHandlerContext context = servletHttpExchange.getChannelHandlerContext();
context.write(nettyResponse);
}
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
checkClosed();
if (len == 0) {
return;
}
ChannelHandlerContext context = servletHttpExchange.getChannelHandlerContext();
ByteBuf ioByteBuf = context.alloc().ioBuffer(len);
ioByteBuf.writeBytes(b, off, len);
IOUtil.writerModeToReadMode(ioByteBuf);
writeHttpBody(ioByteBuf, ioByteBuf.readableBytes());
}
@Override
public boolean isReady() {
ServletHttpExchange exchange = this.servletHttpExchange;
if (exchange == null) {
return true;
}
long pendingWriteBytes = exchange.getPendingWriteBytes();
if (pendingWriteBytes <= 0) {
return true;
}
boolean ready = true;
if (!exchange.getChannelHandlerContext().executor().inEventLoop()) {
// limiting control
// pendingWriteBytes < exchange.getResponse().getBufferSize() * 2
ready = pendingWriteBytes < exchange.getResponse().getBufferSize() << 1;
}
return ready;
}
@Override
public void setWriteListener(WriteListener writeListener) {
this.writeListener = writeListener;
}
public void setCloseListener(ChannelFutureListener closeListener) {
this.closeListenerWrapper.setCloseListener(closeListener);
}
@Override
public void write(byte[] b) throws IOException {
write(b, 0, b.length);
}
/**
* Int. Third-party frameworks are all 1 byte, not 4 bytes
*
* @param b byte
* @throws IOException IOException
*/
@Override
public void write(int b) throws IOException {
checkClosed();
byte[] bytes = new byte[]{(byte) b};
write(bytes, 0, bytes.length);
}
@Override
public void flush() throws IOException {
checkClosed();
writeResponseHeaderIfNeed();
ServletHttpExchange exchange = this.servletHttpExchange;
if (exchange != null && !exchange.getServletContext().isAutoFlush()) {
exchange.getChannelHandlerContext().flush();
}
}
/**
* End response object
* The following events indicate that the servlet has satisfied the request and the response object is about to close
* The following events indicate that the servlet has satisfied the request and the response object is about to close：
* ■The service method of the servlet terminates.
* ■The setContentLength or setContentLengthLong method of the response specifies an internal capacity greater than zero, and has already been written to the response.
* ■sendError Method called。
* ■sendRedirect Method called。
* ■AsyncContext.complete Method called
*/
@Override
public void close() {
if (isClosed.compareAndSet(false, true)) {
ChannelFuture closeFuture = lastContentPromise;
if (closeFuture == null) {
ServletHttpExchange exchange = getServletHttpExchange();
ChannelHandlerContext context = exchange.getChannelHandlerContext();
writeResponseHeaderIfNeed();
if (exchange.getServletContext().isAutoFlush()) {
closeFuture = context.write(LastHttpContent.EMPTY_LAST_CONTENT);
} else {
closeFuture = context.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT);
}
closeFuture.addListener(closeListenerWrapper);
} else if (closeFuture.isDone()) {
try {
closeListenerWrapper.operationComplete(closeFuture);
} catch (Exception e) {
PlatformDependent.throwException(e);
}
} else {
closeFuture.addListener(closeListenerWrapper);
}
}
}
/**
* Check if it's closed
*
* @throws ClosedChannelException
*/
protected void checkClosed() throws IOException {
if (isClosed()) {
throw new IOException("Stream closed");
}
ServletHttpExchange exchange = this.servletHttpExchange;
if (exchange != null && !exchange.isChannelActive() && exchange.isAsyncStartIng()) {
throw new IOException("connection was forcibly closed by the remote host. " + exchange.getChannelHandlerContext().channel());
}
}
/**
* Reset buffer
*/
protected void resetBuffer() {
if (isClosed()) {
return;
}
ServletHttpExchange exchange = getServletHttpExchange();
ChannelHandlerContext channelHandlerContext = exchange.getChannelHandlerContext();
ChannelHandlerContext context = channelHandlerContext.pipeline().context(ChunkedWriteHandler.class);
if (context != null) {
ChunkedWriteHandler handler = (ChunkedWriteHandler) context.handler();
int unWriteSize = handler.unWriteSize();
if (unWriteSize > 0) {
handler.discard(RESET_BUFFER_EXCEPTION);
}
}
writeBytes.set(0);
}
/**
* Get servlet object
*
* @return ServletHttpExchange
*/
protected ServletHttpExchange getServletHttpExchange() {
return servletHttpExchange;
}
/**
* Put in the servlet object
*
* @param servletHttpExchange servletHttpExchange
*/
protected void setServletHttpExchange(ServletHttpExchange servletHttpExchange) {
this.servletHttpExchange = servletHttpExchange;
}
/**
* Whether to shut down
*
* @return True = closed,false= not closed
*/
public boolean isClosed() {
return isClosed.get();
}
@Override
public void recycle(Consumer consumer) {
if (isClosed()) {
return;
}
this.closeListenerWrapper.addRecycleConsumer(consumer);
close();
}
/**
* Closing the listening wrapper class (for data collection)
*/
public class CloseListener implements ChannelFutureListener {
private final Queue recycleConsumerQueue = new LinkedList<>();
private ChannelFutureListener closeListener;
public void addRecycleConsumer(Consumer consumer) {
recycleConsumerQueue.add(consumer);
}
public void setCloseListener(ChannelFutureListener closeListener) {
this.closeListener = closeListener;
}
@Override
public void operationComplete(ChannelFuture future) throws Exception {
ChannelFutureListener closeListener = this.closeListener;
if (closeListener != null) {
closeListener.operationComplete(future);
}
Consumer recycleConsumer;
while ((recycleConsumer = recycleConsumerQueue.poll()) != null) {
recycleConsumer.accept(ServletOutputStream.this);
}
blockPromise = null;
lastContentPromise = null;
writeListener = null;
servletHttpExchange = null;
this.closeListener = null;
RECYCLER.recycleInstance(ServletOutputStream.this);
}
}
}