org.apache.spark.network.util.TransportFrameDecoder Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 org.apache.spark.network.util;
import java.util.LinkedList;
import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.CompositeByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
/**
* A customized frame decoder that allows intercepting raw data.
*
* This behaves like Netty's frame decoder (with hard coded parameters that match this library's
* needs), except it allows an interceptor to be installed to read data directly before it's
* framed.
*
* Unlike Netty's frame decoder, each frame is dispatched to child handlers as soon as it's
* decoded, instead of building as many frames as the current buffer allows and dispatching
* all of them. This allows a child handler to install an interceptor if needed.
*
* If an interceptor is installed, framing stops, and data is instead fed directly to the
* interceptor. When the interceptor indicates that it doesn't need to read any more data,
* framing resumes. Interceptors should not hold references to the data buffers provided
* to their handle() method.
*/
public class TransportFrameDecoder extends ChannelInboundHandlerAdapter {
public static final String HANDLER_NAME = "frameDecoder";
private static final int LENGTH_SIZE = 8;
private static final int MAX_FRAME_SIZE = Integer.MAX_VALUE;
private static final int UNKNOWN_FRAME_SIZE = -1;
private final LinkedList buffers = new LinkedList<>();
private final ByteBuf frameLenBuf = Unpooled.buffer(LENGTH_SIZE, LENGTH_SIZE);
private long totalSize = 0;
private long nextFrameSize = UNKNOWN_FRAME_SIZE;
private volatile Interceptor interceptor;
@Override
public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception {
ByteBuf in = (ByteBuf) data;
buffers.add(in);
totalSize += in.readableBytes();
while (!buffers.isEmpty()) {
// First, feed the interceptor, and if it's still, active, try again.
if (interceptor != null) {
ByteBuf first = buffers.getFirst();
int available = first.readableBytes();
if (feedInterceptor(first)) {
assert !first.isReadable() : "Interceptor still active but buffer has data.";
}
int read = available - first.readableBytes();
if (read == available) {
buffers.removeFirst().release();
}
totalSize -= read;
} else {
// Interceptor is not active, so try to decode one frame.
ByteBuf frame = decodeNext();
if (frame == null) {
break;
}
ctx.fireChannelRead(frame);
}
}
}
private long decodeFrameSize() {
if (nextFrameSize != UNKNOWN_FRAME_SIZE || totalSize < LENGTH_SIZE) {
return nextFrameSize;
}
// We know there's enough data. If the first buffer contains all the data, great. Otherwise,
// hold the bytes for the frame length in a composite buffer until we have enough data to read
// the frame size. Normally, it should be rare to need more than one buffer to read the frame
// size.
ByteBuf first = buffers.getFirst();
if (first.readableBytes() >= LENGTH_SIZE) {
nextFrameSize = first.readLong() - LENGTH_SIZE;
totalSize -= LENGTH_SIZE;
if (!first.isReadable()) {
buffers.removeFirst().release();
}
return nextFrameSize;
}
while (frameLenBuf.readableBytes() < LENGTH_SIZE) {
ByteBuf next = buffers.getFirst();
int toRead = Math.min(next.readableBytes(), LENGTH_SIZE - frameLenBuf.readableBytes());
frameLenBuf.writeBytes(next, toRead);
if (!next.isReadable()) {
buffers.removeFirst().release();
}
}
nextFrameSize = frameLenBuf.readLong() - LENGTH_SIZE;
totalSize -= LENGTH_SIZE;
frameLenBuf.clear();
return nextFrameSize;
}
private ByteBuf decodeNext() {
long frameSize = decodeFrameSize();
if (frameSize == UNKNOWN_FRAME_SIZE || totalSize < frameSize) {
return null;
}
// Reset size for next frame.
nextFrameSize = UNKNOWN_FRAME_SIZE;
Preconditions.checkArgument(frameSize < MAX_FRAME_SIZE, "Too large frame: %s", frameSize);
Preconditions.checkArgument(frameSize > 0, "Frame length should be positive: %s", frameSize);
// If the first buffer holds the entire frame, return it.
int remaining = (int) frameSize;
if (buffers.getFirst().readableBytes() >= remaining) {
return nextBufferForFrame(remaining);
}
// Otherwise, create a composite buffer.
CompositeByteBuf frame = buffers.getFirst().alloc().compositeBuffer(Integer.MAX_VALUE);
while (remaining > 0) {
ByteBuf next = nextBufferForFrame(remaining);
remaining -= next.readableBytes();
frame.addComponent(next).writerIndex(frame.writerIndex() + next.readableBytes());
}
assert remaining == 0;
return frame;
}
/**
* Takes the first buffer in the internal list, and either adjust it to fit in the frame
* (by taking a slice out of it) or remove it from the internal list.
*/
private ByteBuf nextBufferForFrame(int bytesToRead) {
ByteBuf buf = buffers.getFirst();
ByteBuf frame;
if (buf.readableBytes() > bytesToRead) {
frame = buf.retain().readSlice(bytesToRead);
totalSize -= bytesToRead;
} else {
frame = buf;
buffers.removeFirst();
totalSize -= frame.readableBytes();
}
return frame;
}
@Override
public void channelInactive(ChannelHandlerContext ctx) throws Exception {
for (ByteBuf b : buffers) {
b.release();
}
if (interceptor != null) {
interceptor.channelInactive();
}
frameLenBuf.release();
super.channelInactive(ctx);
}
@Override
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
if (interceptor != null) {
interceptor.exceptionCaught(cause);
}
super.exceptionCaught(ctx, cause);
}
public void setInterceptor(Interceptor interceptor) {
Preconditions.checkState(this.interceptor == null, "Already have an interceptor.");
this.interceptor = interceptor;
}
/**
* @return Whether the interceptor is still active after processing the data.
*/
private boolean feedInterceptor(ByteBuf buf) throws Exception {
if (interceptor != null && !interceptor.handle(buf)) {
interceptor = null;
}
return interceptor != null;
}
public interface Interceptor {
/**
* Handles data received from the remote end.
*
* @param data Buffer containing data.
* @return "true" if the interceptor expects more data, "false" to uninstall the interceptor.
*/
boolean handle(ByteBuf data) throws Exception;
/** Called if an exception is thrown in the channel pipeline. */
void exceptionCaught(Throwable cause) throws Exception;
/** Called if the channel is closed and the interceptor is still installed. */
void channelInactive() throws Exception;
}
}