alluxio.client.block.stream.BlockOutStream.java Source code

Java tutorial

Introduction

Here is the source code for alluxio.client.block.stream.BlockOutStream.java

Source

/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.client.block.stream;

import alluxio.Configuration;
import alluxio.PropertyKey;
import alluxio.client.BoundedStream;
import alluxio.client.Cancelable;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.exception.PreconditionMessage;
import alluxio.proto.dataserver.Protocol;
import alluxio.util.CommonUtils;
import alluxio.util.network.NettyUtils;
import alluxio.wire.WorkerNetAddress;

import com.google.common.base.Preconditions;
import com.google.common.io.Closer;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.PooledByteBufAllocator;

import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

import javax.annotation.concurrent.NotThreadSafe;

/**
 * Provides an {@link OutputStream} implementation that is based on {@link PacketWriter} which
 * streams data packet by packet.
 */
@NotThreadSafe
public class BlockOutStream extends OutputStream implements BoundedStream, Cancelable {
    private final Closer mCloser;
    /** Length of the stream. If unknown, set to Long.MAX_VALUE. */
    private final long mLength;
    private ByteBuf mCurrentPacket = null;

    private final List<PacketWriter> mPacketWriters;
    private boolean mClosed;

    /**
     * Creates an {@link BlockOutStream}.
     *
     * @param context the file system context
     * @param blockId the block ID
     * @param blockSize the block size in bytes
     * @param address the Alluxio worker address
     * @param options the out stream options
     * @return the {@link OutputStream} object
     */
    public static BlockOutStream create(FileSystemContext context, long blockId, long blockSize,
            WorkerNetAddress address, OutStreamOptions options) throws IOException {
        if (CommonUtils.isLocalHost(address) && Configuration.getBoolean(PropertyKey.USER_SHORT_CIRCUIT_ENABLED)
                && !NettyUtils.isDomainSocketSupported(address)) {
            return createLocalBlockOutStream(context, address, blockId, blockSize, options);
        } else {
            Protocol.WriteRequest writeRequestPartial = Protocol.WriteRequest.newBuilder().setId(blockId)
                    .setTier(options.getWriteTier()).setType(Protocol.RequestType.ALLUXIO_BLOCK).buildPartial();
            return createNettyBlockOutStream(context, address, blockSize, writeRequestPartial, options);
        }
    }

    /**
     * Creates a {@link BlockOutStream} that writes to a local file.
     *
     * @param context the file system context
     * @param address the worker network address
     * @param id the ID
     * @param length the block or file length
     * @param options the out stream options
     * @return the {@link BlockOutStream} created
     */
    protected static BlockOutStream createLocalBlockOutStream(FileSystemContext context, WorkerNetAddress address,
            long id, long length, OutStreamOptions options) throws IOException {
        long packetSize = Configuration.getBytes(PropertyKey.USER_LOCAL_WRITER_PACKET_SIZE_BYTES);
        PacketWriter packetWriter = LocalFilePacketWriter.create(context, address, id, packetSize, options);
        return new BlockOutStream(packetWriter, length);
    }

    /**
     * Creates a {@link BlockOutStream} that writes to a netty data server.
     *
     * @param context the file system context
     * @param address the netty data server address
     * @param length the block or file length
     * @param partialRequest details of the write request which are constant for all requests
     * @param options the out stream options
     * @return the {@link BlockOutStream} created
     */
    protected static BlockOutStream createNettyBlockOutStream(FileSystemContext context, WorkerNetAddress address,
            long length, Protocol.WriteRequest partialRequest, OutStreamOptions options) throws IOException {
        long packetSize = Configuration.getBytes(PropertyKey.USER_NETWORK_NETTY_WRITER_PACKET_SIZE_BYTES);
        PacketWriter packetWriter = new NettyPacketWriter(context, address, length, partialRequest, packetSize);
        return new BlockOutStream(packetWriter, length);
    }

    /**
     * Constructs a new {@link BlockOutStream} with only one {@link PacketWriter}.
     *
     * @param packetWriter the packet writer
     * @param length the length of the stream
     */
    protected BlockOutStream(PacketWriter packetWriter, long length) {
        mCloser = Closer.create();
        mLength = length;
        mPacketWriters = new ArrayList<>(1);
        mPacketWriters.add(packetWriter);
        mCloser.register(packetWriter);
        mClosed = false;
    }

    /**
     * @return the remaining size of the block
     */
    @Override
    public long remaining() {
        long pos = Long.MAX_VALUE;
        for (PacketWriter packetWriter : mPacketWriters) {
            pos = Math.min(pos, packetWriter.pos());
        }
        return mLength - pos - (mCurrentPacket != null ? mCurrentPacket.readableBytes() : 0);
    }

    @Override
    public void write(int b) throws IOException {
        Preconditions.checkState(remaining() > 0, PreconditionMessage.ERR_END_OF_BLOCK);
        updateCurrentPacket(false);
        mCurrentPacket.writeByte(b);
    }

    @Override
    public void write(byte[] b) throws IOException {
        write(b, 0, b.length);
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        if (len == 0) {
            return;
        }

        while (len > 0) {
            updateCurrentPacket(false);
            int toWrite = Math.min(len, mCurrentPacket.writableBytes());
            mCurrentPacket.writeBytes(b, off, toWrite);
            off += toWrite;
            len -= toWrite;
        }
        updateCurrentPacket(false);
    }

    @Override
    public void flush() throws IOException {
        if (mClosed) {
            return;
        }
        updateCurrentPacket(true);
        for (PacketWriter packetWriter : mPacketWriters) {
            packetWriter.flush();
        }
    }

    @Override
    public void cancel() throws IOException {
        if (mClosed) {
            return;
        }
        releaseCurrentPacket();

        IOException exception = null;
        for (PacketWriter packetWriter : mPacketWriters) {
            try {
                packetWriter.cancel();
            } catch (IOException e) {
                if (exception != null) {
                    exception.addSuppressed(e);
                }
            }
        }
        if (exception != null) {
            throw exception;
        }

        close();
    }

    @Override
    public void close() throws IOException {
        try {
            updateCurrentPacket(true);
        } catch (Throwable t) {
            mCloser.rethrow(t);
        } finally {
            mClosed = true;
            mCloser.close();
        }
    }

    /**
     * Updates the current packet.
     *
     * @param lastPacket if the current packet is the last packet
     */
    private void updateCurrentPacket(boolean lastPacket) throws IOException {
        // Early return for the most common case.
        if (mCurrentPacket != null && mCurrentPacket.writableBytes() > 0 && !lastPacket) {
            return;
        }

        if (mCurrentPacket == null) {
            if (!lastPacket) {
                mCurrentPacket = allocateBuffer();
            }
            return;
        }

        if (mCurrentPacket.writableBytes() == 0 || lastPacket) {
            try {
                if (mCurrentPacket.readableBytes() > 0) {
                    for (PacketWriter packetWriter : mPacketWriters) {
                        mCurrentPacket.retain();
                        packetWriter.writePacket(mCurrentPacket.duplicate());
                    }
                } else {
                    Preconditions.checkState(lastPacket);
                }
            } finally {
                // If the packet has bytes to read, we increment its refcount explicitly for every packet
                // writer. So we need to release here. If the packet has no bytes to read, then it has
                // to be the last packet. It needs to be released as well.
                mCurrentPacket.release();
            }
            mCurrentPacket = null;
        }
        if (!lastPacket) {
            mCurrentPacket = allocateBuffer();
        }
    }

    /**
     * Releases the current packet.
     */
    private void releaseCurrentPacket() {
        if (mCurrentPacket != null) {
            mCurrentPacket.release();
            mCurrentPacket = null;
        }
    }

    /**
     * @return a newly allocated byte buffer of the user defined default size
     */
    private ByteBuf allocateBuffer() {
        return PooledByteBufAllocator.DEFAULT.buffer(mPacketWriters.get(0).packetSize());
    }
}