alluxio-master-2.x/core/client/fs/src/main/java/alluxio/client/block/stream/BlockOutStream.java
/*
 * 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.client.BoundedStream;
import alluxio.client.Cancelable;
import alluxio.client.file.FileSystemContext;
import alluxio.client.file.options.OutStreamOptions;
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 io.netty.buffer.Unpooled;

import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * Provides an {@link OutputStream} implementation that is based on {@link DataWriter} which
 * streams data chunk by chunk.
 */
@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 final WorkerNetAddress mAddress;
  private ByteBuf mCurrentChunk = null;

  private final List<DataWriter> mDataWriters;
  private boolean mClosed;

  /**
   * Constructs a new {@link BlockOutStream} with only one {@link DataWriter}.
   *
   * @param dataWriter the data writer
   * @param length the length of the stream
   * @param address the Alluxio worker address
   */
  public BlockOutStream(DataWriter dataWriter, long length, WorkerNetAddress address) {
    mCloser = Closer.create();
    mLength = length;
    Objects.requireNonNull(address);
    mAddress = address;
    mDataWriters = new ArrayList<>(1);
    Objects.requireNonNull(dataWriter);
    mDataWriters.add(dataWriter);
    mCloser.register(dataWriter);
    mClosed = false;
  }

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

  /**
   * Creates a new remote block output stream.
   *
   * @param context the file system context
   * @param blockId the block id
   * @param blockSize the block size
   * @param workerNetAddresses the worker network addresses
   * @param options the options
   * @return the {@link BlockOutStream} instance created
   */
  public static BlockOutStream createReplicatedBlockOutStream(FileSystemContext context,
      long blockId, long blockSize, java.util.List<WorkerNetAddress> workerNetAddresses,
      OutStreamOptions options) throws IOException {
    List<DataWriter> dataWriters = new ArrayList<>();
    for (WorkerNetAddress address: workerNetAddresses) {
      DataWriter dataWriter =
            DataWriter.Factory.create(context, blockId, blockSize, address, options);
      dataWriters.add(dataWriter);
    }
    return new BlockOutStream(dataWriters, blockSize, workerNetAddresses);
  }

  /**
   * Constructs a new {@link BlockOutStream} with only one {@link DataWriter}.
   *
   * @param dataWriters the data writer
   * @param length the length of the stream
   * @param workerNetAddresses the worker network addresses
   */
  protected BlockOutStream(List<DataWriter> dataWriters, long length,
      java.util.List<WorkerNetAddress> workerNetAddresses) {
    mCloser = Closer.create();
    mLength = length;
    mAddress = workerNetAddresses.get(0);
    mDataWriters = dataWriters;
    for (DataWriter dataWriter : dataWriters) {
      mCloser.register(dataWriter);
    }
    mClosed = false;
  }

  @Override
  public void write(int b) throws IOException {
    Preconditions.checkState(remaining() > 0, "Cannot write past end of block");
    updateCurrentChunk(false);
    mCurrentChunk.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 {
    writeInternal(Unpooled.wrappedBuffer(b), off, len);
  }

  /**
   * Writes the data in the specified byte buf to this output stream.
   *
   * @param buf the buffer
   * @throws IOException exception
   */
  public void write(io.netty.buffer.ByteBuf buf) throws IOException {
    write(buf, 0, buf.readableBytes());
  }

  /**
   * Writes len bytes from the specified byte buf starting at offset off to this output stream.
   *
   * @param buf the buffer
   * @param off the offset
   * @param len the length
   */
  public void write(io.netty.buffer.ByteBuf buf, int off, int len) throws IOException {
    writeInternal(buf, off, len);
  }

  private void writeInternal(ByteBuf b, int off, int len) throws IOException {
    if (len == 0) {
      return;
    }

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

  @Override
  public void flush() throws IOException {
    if (mClosed) {
      return;
    }
    updateCurrentChunk(true);
    for (DataWriter dataWriter : mDataWriters) {
      dataWriter.flush();
    }
  }

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

    List<Exception> exceptions = new LinkedList<>();
    for (DataWriter dataWriter : mDataWriters) {
      try {
        dataWriter.cancel();
      } catch (IOException e) {
        exceptions.add(e);
      }
    }
    if (exceptions.size() > 0) {
      IOException ex = new IOException("Failed to cancel all block write attempts");
      exceptions.forEach(ex::addSuppressed);
      throw ex;
    }

    close();
  }

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

  /**
   * @return the worker address for this stream
   */
  public WorkerNetAddress getAddress() {
    return mAddress;
  }

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

    if (mCurrentChunk == null) {
      if (!lastChunk) {
        mCurrentChunk = allocateBuffer();
      }
      return;
    }

    if (mCurrentChunk.writableBytes() == 0 || lastChunk) {
      try {
        if (mCurrentChunk.readableBytes() > 0) {
          for (DataWriter dataWriter : mDataWriters) {
            mCurrentChunk.retain();
            dataWriter.writeChunk(mCurrentChunk.duplicate());
          }
        } else {
          Preconditions.checkState(lastChunk);
        }
      } 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.
        mCurrentChunk.release();
        mCurrentChunk = null;
      }
    }
    if (!lastChunk) {
      mCurrentChunk = allocateBuffer();
    }
  }

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

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