alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/block/UfsIOManager.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.worker.block;

import alluxio.AlluxioURI;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.runtime.AlluxioRuntimeException;
import alluxio.exception.runtime.OutOfRangeRuntimeException;
import alluxio.exception.runtime.ResourceExhaustedRuntimeException;
import alluxio.grpc.UfsReadOptions;
import alluxio.metrics.MetricInfo;
import alluxio.metrics.MetricKey;
import alluxio.metrics.MetricsSystem;
import alluxio.resource.CloseableResource;
import alluxio.underfs.UfsManager;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.options.OpenOptions;
import alluxio.util.IdUtils;
import alluxio.util.ThreadFactoryUtils;

import com.codahale.metrics.Meter;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

import java.io.Closeable;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;

/**
 * Control UFS IO.
 */
public class UfsIOManager implements Closeable {
  private static final int READ_CAPACITY = 1024;
  private final UfsManager.UfsClient mUfsClient;
  private final ConcurrentMap<String, Long> mThroughputQuota = new ConcurrentHashMap<>();
  private final UfsInputStreamCache mUfsInstreamCache = new UfsInputStreamCache();
  private final LinkedBlockingQueue<ReadTask> mReadQueue = new LinkedBlockingQueue<>(READ_CAPACITY);
  private final ConcurrentMap<AlluxioURI, Meter> mUfsBytesReadThroughputMetrics =
      new ConcurrentHashMap<>();
  private final ExecutorService mUfsIoExecutor =
      Executors.newFixedThreadPool(Configuration.getInt(PropertyKey.UNDERFS_IO_THREADS),
          ThreadFactoryUtils.build("UfsIOManager-IO-%d", false));
  private final ExecutorService mScheduleExecutor = Executors
      .newSingleThreadExecutor(ThreadFactoryUtils.build("UfsIOManager-Scheduler-%d", true));

  /**
   * @param ufsClient ufs client
   */
  public UfsIOManager(UfsManager.UfsClient ufsClient) {
    mUfsClient = ufsClient;
  }

  /**
   * Start schedule thread. Main reason is for test
   */
  public void start() {
    mScheduleExecutor.submit(this::schedule);
  }

  /**
   * Close.
   */
  @Override
  public void close() {
    mScheduleExecutor.shutdownNow();
    mUfsIoExecutor.shutdownNow();
  }

  /**
   * Set throughput quota for tag.
   * @param tag the client name or tag
   * @param throughput throughput limit in bytes
   */
  public void setQuota(String tag, long throughput) {
    Preconditions.checkArgument(throughput > 0, "throughput should be positive");
    mThroughputQuota.put(tag, throughput);
  }

  private void schedule() {
    while (!Thread.currentThread().isInterrupted()) {
      try {
        ReadTask task = mReadQueue.take();
        if (mThroughputQuota.containsKey(task.mOptions.getTag())
            && mThroughputQuota.get(task.mOptions.getTag()) < getUsedThroughput(task.mMeter)) {
          // resubmit to queue
          mReadQueue.put(task);
        } else {
          try {
            mUfsIoExecutor.submit(task);
          } catch (RejectedExecutionException e) {
            // should not reach here since we use unbounded queue in thread pool and use a bounded
            // blocking queue upfront to control number of read tasks
            mReadQueue.put(task);
          }
        }
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
    }
  }

  /**
   * Get used throughput.
   * @param meter throughput meter
   * @return throughput
   */
  @VisibleForTesting
  public double getUsedThroughput(Meter meter) {
    return meter.getOneMinuteRate() / 60;
  }

  /**
   * Read from ufs.
   *
   * @param buf bytebuffer
   * @param offset  offset in ufs file
   * @param len  length to read
   * @param blockId block id
   * @param ufsPath ufs path
   * @param options read ufs options
   * @return content read
   * @throws ResourceExhaustedRuntimeException when too many read task happens
   * @throws OutOfRangeRuntimeException offset is negative, len is negative, or len > buf remaining
   * @throws AlluxioRuntimeException future complete exceptionally when having exception from ufs
   */
  public CompletableFuture<Integer> read(ByteBuffer buf, long offset, long len, long blockId,
      String ufsPath, UfsReadOptions options) {
    Objects.requireNonNull(buf);
    if (offset < 0 || len < 0 || len > buf.remaining()) {
      throw new OutOfRangeRuntimeException(String.format(
          "offset is negative, len is negative, or len is greater than buf remaining. "
              + "offset: %s, len: %s, buf remaining: %s", offset, len, buf.remaining()));
    }
    if (mReadQueue.size() >= READ_CAPACITY) {
      throw new ResourceExhaustedRuntimeException("UFS read at capacity", true);
    }
    CompletableFuture<Integer> future = new CompletableFuture<>();
    if (len == 0) {
      future.complete(0);
      return future;
    }
    Meter meter = mUfsBytesReadThroughputMetrics.computeIfAbsent(mUfsClient.getUfsMountPointUri(),
        uri -> MetricsSystem.meterWithTags(MetricKey.WORKER_BYTES_READ_UFS_THROUGHPUT.getName(),
            MetricKey.WORKER_BYTES_READ_UFS_THROUGHPUT.isClusterAggregated(), MetricInfo.TAG_UFS,
            MetricsSystem.escape(mUfsClient.getUfsMountPointUri()), MetricInfo.TAG_USER,
            options.getTag()));

    mReadQueue.add(new ReadTask(buf, ufsPath, IdUtils.fileIdFromBlockId(blockId), offset,
        len, options, future, meter));
    return future;
  }

  private class ReadTask implements Runnable {
    private final long mOffset;
    private final long mLength;
    private final CompletableFuture<Integer> mFuture;
    private final String mUfsPath;
    private final UfsReadOptions mOptions;
    private final Meter mMeter;
    private final long mFileId;
    private final ByteBuffer mBuffuer;

    private ReadTask(ByteBuffer buf, String ufsPath, long fileId, long offset, long length,
        UfsReadOptions options, CompletableFuture<Integer> future, Meter meter) {
      mOptions = options;
      mUfsPath = ufsPath;
      mFileId = fileId;
      mOffset = offset;
      mLength = length;
      mFuture = future;
      mMeter = meter;
      mBuffuer = buf;
    }

    public void run() {
      try {
        mFuture.complete(readInternal());
      } catch (RuntimeException e) {
        mFuture.completeExceptionally(e);
      }
    }

    private int readInternal() {
      int bytesRead = 0;
      InputStream inStream = null;
      try (CloseableResource<UnderFileSystem> ufsResource = mUfsClient.acquireUfsResource()) {
        if (mOptions.hasUser()) {
          // Before interacting with ufs manager, set the user.
          alluxio.security.authentication.AuthenticatedClientUser.set(mOptions.getUser());
        }
        inStream = mUfsInstreamCache.acquire(ufsResource.get(), mUfsPath, mFileId,
            OpenOptions.defaults().setOffset(mOffset)
                .setPositionShort(mOptions.getPositionShort()));
        while (bytesRead < mLength) {
          int read;
          read = Channels.newChannel(inStream).read(mBuffuer);
          if (read == -1) {
            break;
          }
          bytesRead += read;
        }
      } catch (Exception e) {
        throw AlluxioRuntimeException.from(e);
      } finally {
        if (inStream != null) {
          mUfsInstreamCache.release(inStream);
        }
      }
      mMeter.mark(bytesRead);
      return bytesRead;
    }
  }
}
