org.apache.tajo.plan.function.stream.ByteBufLineReader.java Source code

Introduction

Here is the source code for org.apache.tajo.plan.function.stream.ByteBufLineReader.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.tajo.plan.function.stream;

import io.netty.buffer.ByteBuf;
import io.netty.util.CharsetUtil;

import java.io.Closeable;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicInteger;

public class ByteBufLineReader implements Closeable {
    private static int DEFAULT_BUFFER = 64 * 1024;

    private int bufferSize;
    private long readBytes;
    private int startIndex;
    private boolean eof = false;
    private ByteBuf buffer;
    private final ByteBufInputChannel channel;
    private final AtomicInteger lineReadBytes = new AtomicInteger();
    private final LineSplitProcessor processor = new LineSplitProcessor();

    public ByteBufLineReader(ByteBufInputChannel channel) {
        this(channel, BufferPool.directBuffer(DEFAULT_BUFFER));
    }

    public ByteBufLineReader(ByteBufInputChannel channel, ByteBuf buf) {
        this.readBytes = 0;
        this.channel = channel;
        this.buffer = buf;
        this.bufferSize = buf.capacity();
    }

    public long readBytes() {
        return readBytes - buffer.readableBytes();
    }

    @Override
    public void close() throws IOException {
        if (this.buffer.refCnt() > 0) {
            this.buffer.release();
        }
        this.channel.close();
    }

    public String readLine() throws IOException {
        ByteBuf buf = readLineBuf(lineReadBytes);
        if (buf != null) {
            return buf.toString(CharsetUtil.UTF_8);
        }
        return null;
    }

    private void fillBuffer() throws IOException {

        int tailBytes = 0;
        if (this.readBytes > 0) {
            //startIndex = 0, readIndex = tailBytes length, writable = (buffer capacity - tailBytes)
            this.buffer.markReaderIndex();
            this.buffer.discardReadBytes(); // compact the buffer
            tailBytes = this.buffer.writerIndex();
            if (!this.buffer.isWritable()) {
                // a line bytes is large than the buffer
                BufferPool.ensureWritable(buffer, bufferSize * 2);
                this.bufferSize = buffer.capacity();
            }
            this.startIndex = 0;
        }

        boolean release = true;
        try {
            int readBytes = tailBytes;
            // read only once
            int localReadBytes = buffer.writeBytes(channel, this.bufferSize - readBytes);
            if (localReadBytes < 0) {
                if (buffer.isWritable()) {
                    //if read bytes is less than the buffer capacity,  there is no more bytes in the channel
                    eof = true;
                }
            }
            readBytes += localReadBytes;

            this.readBytes += (readBytes - tailBytes);
            release = false;

            this.buffer.readerIndex(this.buffer.readerIndex() + tailBytes); //skip past buffer (tail)
        } finally {
            if (release) {
                buffer.release();
            }
        }
    }

    /**
     * Read a line terminated by one of CR, LF, or CRLF.
     */
    public ByteBuf readLineBuf(AtomicInteger reads) throws IOException {
        int readBytes = 0; // newline + text line bytes
        int newlineLength = 0; //length of terminating newline
        int readable;

        this.startIndex = buffer.readerIndex();

        loop: while (true) {
            readable = buffer.readableBytes();
            if (readable <= 0) {
                buffer.readerIndex(this.startIndex);
                fillBuffer(); //compact and fill buffer

                //if buffer.writerIndex() is zero, there is no bytes in buffer
                if (!buffer.isReadable() && buffer.writerIndex() == 0) {
                    reads.set(0);
                    return null;
                } else {
                    //skip first newLine
                    if (processor.isPrevCharCR() && buffer.getByte(buffer.readerIndex()) == LineSplitProcessor.LF) {
                        buffer.skipBytes(1);
                        if (eof && !buffer.isReadable()) {
                            reads.set(1);
                            return null;
                        }

                        newlineLength++;
                        readBytes++;
                        startIndex = buffer.readerIndex();
                    }
                }
                readable = buffer.readableBytes();
            }

            int endIndex = buffer.forEachByte(buffer.readerIndex(), readable, processor);
            if (endIndex < 0) {
                //does not appeared terminating newline
                buffer.readerIndex(buffer.writerIndex()); // set to end buffer
                if (eof) {
                    readBytes += (buffer.readerIndex() - startIndex);
                    break loop;
                }
            } else {
                buffer.readerIndex(endIndex + 1);
                readBytes += (buffer.readerIndex() - startIndex); //past newline + text line

                //appeared terminating CRLF
                if (processor.isPrevCharCR() && buffer.isReadable()
                        && buffer.getByte(buffer.readerIndex()) == LineSplitProcessor.LF) {
                    buffer.skipBytes(1);
                    readBytes++;
                    newlineLength += 2;
                } else {
                    newlineLength += 1;
                }
                break loop;
            }
        }
        reads.set(readBytes);
        return buffer.slice(startIndex, readBytes - newlineLength);
    }
}