/*
* 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 com.addthis.hydra.query.web;
import java.nio.CharBuffer;
import com.addthis.basis.util.Parameter;
import com.addthis.bundle.core.Bundle;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.ByteBufUtil;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelOutboundHandlerAdapter;
import io.netty.channel.ChannelPromise;
import io.netty.handler.codec.http.DefaultHttpContent;
import io.netty.handler.codec.http.DefaultHttpResponse;
import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.HttpResponse;
import io.netty.handler.codec.http.HttpResponseStatus;
import io.netty.handler.codec.http.HttpVersion;
import io.netty.util.CharsetUtil;
/**
* parent of all streaming response classes
*/
abstract class AbstractBufferingHttpBundleEncoder extends ChannelOutboundHandlerAdapter {
private static final Logger log = LoggerFactory.getLogger(AbstractBufferingHttpBundleEncoder.class);
private static final int DEFAULT_INITIAL_BUFFER_SIZE = Parameter.intValue("qmaster.http.buffer.initial", 100);
private static final int DEFAULT_BATCH_BUFFER_SIZE = Parameter.intValue("qmaster.http.buffer.batch", 100000);
protected final HttpResponse responseStart = new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK);
private final StringBuilder sendBuffer;
private final int batchBufferSize;
private boolean writeStarted = false;
private boolean responseWritten = false;
AbstractBufferingHttpBundleEncoder(int initialBufferSize, int batchBufferSize) {
this.batchBufferSize = batchBufferSize;
HttpHeaders.setTransferEncodingChunked(responseStart);
sendBuffer = new StringBuilder(initialBufferSize);
}
AbstractBufferingHttpBundleEncoder() {
this(DEFAULT_INITIAL_BUFFER_SIZE, DEFAULT_BATCH_BUFFER_SIZE);
}
private static ByteBuf encodeString(ByteBufAllocator alloc, CharSequence msg) {
return ByteBufUtil.encodeString(alloc, CharBuffer.wrap(msg), CharsetUtil.UTF_8);
}
@Override
public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
if (msg instanceof Bundle) {
send(ctx, (Bundle) msg);
} else if (msg == DataChannelOutputToNettyBridge.SEND_COMPLETE) {
sendComplete(ctx);
ctx.pipeline().remove(this);
} else {
super.write(ctx, msg, promise); // forward write to next handler
}
}
@Override
public void flush(ChannelHandlerContext ctx) {
if (sendBuffer.length() > 0) {
flushStringBuilder(ctx);
} else {
ctx.flush();
}
}
private boolean maybeWriteStart(ChannelHandlerContext ctx, Bundle row) {
if (!writeStarted) {
appendResponseStartToString(sendBuffer);
if (row != null) {
appendInitialBundleToString(row, sendBuffer);
}
writeStarted = true;
return true;
}
return false;
}
public abstract void appendBundleToString(Bundle row, StringBuilder sendBuffer);
/**
* Called before any bundles are written.
*/
protected void appendResponseStartToString(StringBuilder sendBuffer) {
// override in subclasses if desired
}
/**
* The bundle firstRow is passed in and should be written in this method.
* It is provided due to the frequent case of the first row requiring special
* logic. If (first/ not first) is the only context needed to encode a bundle,
* then subsequent bundles can be encoded concurrently.
*/
protected void appendInitialBundleToString(Bundle firstRow, StringBuilder sendBuffer) {
appendBundleToString(firstRow, sendBuffer);
}
/**
* Called after all bundles are written.
*/
protected void appendResponseEndToString(StringBuilder sendBuffer) {
// override in subclasses if desired
}
public void send(ChannelHandlerContext ctx, Bundle row) {
if (!maybeWriteStart(ctx, row)) {
appendBundleToString(row, sendBuffer);
}
if (sendBuffer.length() >= batchBufferSize) {
flushStringBuilder(ctx);
}
}
protected void flushStringBuilder(ChannelHandlerContext ctx) {
if (!responseWritten) {
ctx.write(responseStart);
}
if (sendBuffer.length() > 0) {
ByteBuf msg = encodeString(ctx.alloc(), sendBuffer);
sendBuffer.setLength(0);
ctx.writeAndFlush(new DefaultHttpContent(msg), ctx.voidPromise());
}
}
public void sendComplete(ChannelHandlerContext ctx) {
maybeWriteStart(ctx, null);
appendResponseEndToString(sendBuffer);
flushStringBuilder(ctx);
}
}