org.apache.drill.exec.rpc.RpcDecoder.java Source code

Introduction

Here is the source code for org.apache.drill.exec.rpc.RpcDecoder.java

Source

/**
 * 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.drill.exec.rpc;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufInputStream;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.CorruptedFrameException;
import io.netty.handler.codec.MessageToMessageDecoder;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.drill.exec.proto.GeneralRPCProtos.RpcHeader;

/**
 * Converts a previously length adjusted buffer into an RpcMessage.
 */
class RpcDecoder extends MessageToMessageDecoder<ByteBuf> {
    final org.slf4j.Logger logger;

    private final AtomicLong messageCounter = new AtomicLong();

    public RpcDecoder(String name) {
        this.logger = org.slf4j.LoggerFactory.getLogger(RpcDecoder.class.getCanonicalName() + "-" + name);
    }

    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> out) throws Exception {
        if (!ctx.channel().isOpen()) {
            return;
        }

        if (RpcConstants.EXTRA_DEBUGGING) {
            logger.debug("Inbound rpc message received.");
        }

        // now, we know the entire message is in the buffer and the buffer is constrained to this message. Additionally,
        // this process should avoid reading beyond the end of this buffer so we inform the ByteBufInputStream to throw an
        // exception if be go beyond readable bytes (as opposed to blocking).
        final ByteBufInputStream is = new ByteBufInputStream(buffer, buffer.readableBytes());

        // read the rpc header, saved in delimited format.
        checkTag(is, RpcEncoder.HEADER_TAG);
        final RpcHeader header = RpcHeader.parseDelimitedFrom(is);

        if (RpcConstants.EXTRA_DEBUGGING) {
            logger.debug(" post header read index {}", buffer.readerIndex());
        }

        // read the protobuf body into a buffer.
        checkTag(is, RpcEncoder.PROTOBUF_BODY_TAG);
        final int pBodyLength = readRawVarint32(is);
        final ByteBuf pBody = buffer.slice(buffer.readerIndex(), pBodyLength);
        buffer.skipBytes(pBodyLength);
        pBody.retain(1);
        if (RpcConstants.EXTRA_DEBUGGING) {
            logger.debug("Read protobuf body of length {} into buffer {}.", pBodyLength, pBody);
        }

        if (RpcConstants.EXTRA_DEBUGGING) {
            logger.debug("post protobufbody read index {}", buffer.readerIndex());
        }

        ByteBuf dBody = null;
        int dBodyLength = 0;

        // read the data body.
        if (buffer.readableBytes() > 0) {

            if (RpcConstants.EXTRA_DEBUGGING) {
                logger.debug("Reading raw body, buffer has {} bytes available, is available {}.",
                        buffer.readableBytes(), is.available());
            }
            checkTag(is, RpcEncoder.RAW_BODY_TAG);
            dBodyLength = readRawVarint32(is);
            if (buffer.readableBytes() != dBodyLength) {
                throw new CorruptedFrameException(String.format(
                        "Expected to receive a raw body of %d bytes but received a buffer with %d bytes.",
                        dBodyLength, buffer.readableBytes()));
            }
            dBody = buffer.slice();
            dBody.retain(1);
            if (RpcConstants.EXTRA_DEBUGGING) {
                logger.debug("Read raw body of {}", dBody);
            }

        } else {
            if (RpcConstants.EXTRA_DEBUGGING) {
                logger.debug("No need to read raw body, no readable bytes left.");
            }
        }

        // return the rpc message.
        InboundRpcMessage m = new InboundRpcMessage(header.getMode(), header.getRpcType(),
                header.getCoordinationId(), pBody, dBody);

        // move the reader index forward so the next rpc call won't try to work with it.
        buffer.skipBytes(dBodyLength);
        messageCounter.incrementAndGet();
        if (RpcConstants.SOME_DEBUGGING) {
            logger.debug("Inbound Rpc Message Decoded {}.", m);
        }
        out.add(m);

    }

    private void checkTag(ByteBufInputStream is, int expectedTag) throws IOException {
        int actualTag = readRawVarint32(is);
        if (actualTag != expectedTag) {
            throw new CorruptedFrameException(String.format(
                    "Expected to read a tag of %d but actually received a value of %d.  Happened after reading %d message.",
                    expectedTag, actualTag, messageCounter.get()));
        }
    }

    // Taken from CodedInputStream and modified to enable ByteBufInterface.
    public static int readRawVarint32(ByteBufInputStream is) throws IOException {
        byte tmp = is.readByte();
        if (tmp >= 0) {
            return tmp;
        }
        int result = tmp & 0x7f;
        if ((tmp = is.readByte()) >= 0) {
            result |= tmp << 7;
        } else {
            result |= (tmp & 0x7f) << 7;
            if ((tmp = is.readByte()) >= 0) {
                result |= tmp << 14;
            } else {
                result |= (tmp & 0x7f) << 14;
                if ((tmp = is.readByte()) >= 0) {
                    result |= tmp << 21;
                } else {
                    result |= (tmp & 0x7f) << 21;
                    result |= (tmp = is.readByte()) << 28;
                    if (tmp < 0) {
                        // Discard upper 32 bits.
                        for (int i = 0; i < 5; i++) {
                            if (is.readByte() >= 0) {
                                return result;
                            }
                        }
                        throw new CorruptedFrameException("Encountered a malformed varint.");
                    }
                }
            }
        }
        return result;
    }

}