content: /*
   * 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.Constants;
  import alluxio.Sessions;
  import alluxio.client.file.FileSystemContext;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.exception.AlluxioException;
  import alluxio.exception.status.CancelledException;
  import alluxio.grpc.CacheRequest;
  import alluxio.metrics.MetricKey;
  import alluxio.metrics.MetricsSystem;
  import alluxio.proto.dataserver.Protocol;
  import alluxio.util.CommonUtils;
  import alluxio.util.WaitForOptions;
  import alluxio.util.io.BufferUtils;
  import alluxio.util.logging.SamplingLogger;
  import alluxio.util.network.NetworkAddressUtils;
  import alluxio.worker.block.io.BlockReader;
  import alluxio.worker.block.io.BlockWriter;
  import alluxio.worker.block.meta.TempBlockMeta;
  
  import com.codahale.metrics.Counter;
  import com.google.common.annotations.VisibleForTesting;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.RejectedExecutionException;
  import java.util.concurrent.TimeoutException;
  import java.util.concurrent.atomic.AtomicLong;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * Handles client requests to synchronously/asynchronously cache blocks. Responsible for
   * managing the local worker resources and intelligent pruning of duplicate or meaningless requests.
   */
  @ThreadSafe
  public class CacheRequestManager {
    private static final Logger LOG = LoggerFactory.getLogger(CacheRequestManager.class);
    private static final Logger SAMPLING_LOG = new SamplingLogger(LOG, 10L * Constants.MINUTE_MS);
    private static final int NETWORK_HOST_RESOLUTION_TIMEOUT =
        (int) Configuration.getMs(PropertyKey.NETWORK_HOST_RESOLUTION_TIMEOUT_MS);
  
    /** Executor service for execute the async cache tasks. */
    private final ExecutorService mCacheExecutor;
    /** The block worker. */
    private final DefaultBlockWorker mBlockWorker;
    private final ConcurrentHashMap<Long, CacheRequest> mActiveCacheRequests =
        new ConcurrentHashMap<>();
    private final FileSystemContext mFsContext;
    /** Keeps track of the number of rejected cache requests. */
    private final AtomicLong mNumRejected = new AtomicLong(0);
  
    /**
     * @param service thread pool to run the background caching work
     * @param blockWorker handler to the block worker
     * @param fsContext context
     */
    public CacheRequestManager(ExecutorService service, DefaultBlockWorker blockWorker,
        FileSystemContext fsContext) {
      mCacheExecutor = service;
      mBlockWorker = blockWorker;
      mFsContext = fsContext;
    }
  
    /**
     * Handles a request to cache a block. If it's async cache, it wouldn't throw exception.
     *
     * @param request the cache request fields will be available
     */
    public void submitRequest(CacheRequest request)
        throws AlluxioException, IOException {
      CACHE_REQUESTS.inc();
      long blockId = request.getBlockId();
      boolean async = request.getAsync();
      if (mActiveCacheRequests.putIfAbsent(blockId, request) != null) {
        // This block is already planned and just return.
        if (async) {
          LOG.debug("request already planned: {}", request);
        } else {
          try {
            CommonUtils.waitFor("block to be loaded",
                () -> !mActiveCacheRequests.containsKey(blockId),
                WaitForOptions.defaults().setTimeoutMs(30 * Constants.SECOND_MS));
          } catch (InterruptedException e) {
            throw new CancelledException("Fail to finish cache request synchronously. "
                + "Interrupted while waiting for block to be loaded by another request.", e);
          } catch (TimeoutException e) {
            throw new CancelledException("Fail to finish cache request synchronously due to timeout",
                e);
          }
        }
        return;
      }
  
      if (!async) {
        CACHE_REQUESTS_SYNC.inc();
      } else {
        CACHE_REQUESTS_ASYNC.inc();
      }
      Future<Void> future = null;
      try {
        future = mCacheExecutor.submit(new CacheTask(request));
      } catch (RejectedExecutionException e) {
        // RejectedExecutionException may be thrown in extreme cases when the
        // gRPC thread pool is drained due to highly concurrent caching workloads. In these cases,
        // return as async caching is at best effort.
        mNumRejected.incrementAndGet();
        CACHE_REJECTED_BLOCKS.inc();
        SAMPLING_LOG.warn(String.format(
            "Failed to cache block locally as the thread pool is at capacity."
                + " To increase, update the parameter '%s'. numRejected: {} error: {}",
            PropertyKey.Name.WORKER_NETWORK_ASYNC_CACHE_MANAGER_THREADS_MAX), mNumRejected.get(),
            e.toString());
        mActiveCacheRequests.remove(blockId);
        if (!async) {
          throw new CancelledException(String.format(
              "Fail to finish cache request synchronously as the thread pool is at capacity."
                  + " To increase the capacity, set the parameter '%s' and '%s' higher. ",
              PropertyKey.Name.WORKER_NETWORK_ASYNC_CACHE_MANAGER_THREADS_MAX,
              PropertyKey.Name.WORKER_NETWORK_ASYNC_CACHE_MANAGER_QUEUE_MAX), e);
        }
      }
      if (future != null && !async) {
        try {
          future.get();
        } catch (ExecutionException e) {
          CACHE_FAILED_BLOCKS.inc();
          Throwable cause = e.getCause();
          if (cause instanceof AlluxioException) {
            throw new AlluxioException(cause.getMessage(), cause);
          } else {
            throw new IOException(cause);
          }
        } catch (InterruptedException e) {
          throw new CancelledException(
              "Fail to finish cache request synchronously. Interrupted while waiting for response.",
              e);
        }
      }
    }
  
    /**
     * CacheTask is a callable task that can be considered equal if the blockId of the request is the
     * same.
     */
    @VisibleForTesting
    class CacheTask implements Callable<Void> {
      private final CacheRequest mRequest;
  
      /**
       * Constructor for an CacheTask.
       *
       * @param request an CacheRequest
       */
      CacheTask(CacheRequest request) {
        mRequest = request;
      }
  
      @Override
      public int hashCode() {
        // Only care about the block id being the same.
        // Do not care if the source host or port is different.
        return Objects.hash(mRequest.getBlockId());
      }
  
      @Override
      public boolean equals(Object obj) {
        if (this == obj) {
          return true;
        }
        if (!(obj instanceof CacheTask)) {
          return false;
        }
        CacheTask that = ((CacheTask) obj);
        if (mRequest == that.mRequest) {
          return true;
        }
        if (this.mRequest == null || that.mRequest == null) {
          return false;
        }
        // Only care about the block id being the same.
        // Do not care if the source host or port is different.
        return mRequest.getBlockId() == that.mRequest.getBlockId();
      }
  
      @Override
      public Void call() throws IOException, AlluxioException {
        long blockId = mRequest.getBlockId();
        long blockLength = mRequest.getLength();
        CacheResult result = CacheResult.FAILED;
        try {
          result = cacheBlock(mRequest);
        } finally {
          switch (result) {
            case SUCCEED:
              CACHE_BLOCKS_SIZE.inc(blockLength);
              CACHE_SUCCEEDED_BLOCKS.inc();
              break;
            case FAILED:
              CACHE_FAILED_BLOCKS.inc();
              break;
            default:
              break;
          }
          mActiveCacheRequests.remove(blockId);
        }
        return null;
      }
    }
  
    enum CacheResult {
  
      SUCCEED, FAILED, ALREADY_CACHED
    }
  
    private CacheResult cacheBlock(CacheRequest request) throws IOException, AlluxioException {
      CacheResult result;
      boolean isSourceLocal = NetworkAddressUtils.isLocalAddress(request.getSourceHost(),
          NETWORK_HOST_RESOLUTION_TIMEOUT);
      long blockId = request.getBlockId();
      long blockLength = request.getLength();
      // Check if the block has already been cached on this worker
      if (mBlockWorker.getBlockStore().hasBlockMeta(blockId)) {
        LOG.debug("block already cached: {}", blockId);
        return CacheResult.ALREADY_CACHED;
      }
      Protocol.OpenUfsBlockOptions openUfsBlockOptions = request.getOpenUfsBlockOptions();
      // Depends on the request, cache the target block from different sources
      if (isSourceLocal) {
        CACHE_UFS_BLOCKS.inc();
        result = cacheBlockFromUfs(blockId, blockLength, openUfsBlockOptions);
      } else {
        CACHE_REMOTE_BLOCKS.inc();
        InetSocketAddress sourceAddress =
            new InetSocketAddress(request.getSourceHost(), request.getSourcePort());
        result =
            cacheBlockFromRemoteWorker(blockId, blockLength, sourceAddress, openUfsBlockOptions);
      }
      LOG.debug("Result of caching block {}: {}", blockId, result);
      return result;
    }
  
    /**
     * Caches the block via the local worker to read from UFS.
     *
     * @param blockId block ID
     * @param blockSize block size
     * @param openUfsBlockOptions options to open the UFS file
     * @return cache result
     */
    private CacheResult cacheBlockFromUfs(long blockId, long blockSize,
        Protocol.OpenUfsBlockOptions openUfsBlockOptions) throws IOException {
      try (BlockReader reader = mBlockWorker.createUfsBlockReader(Sessions.CACHE_UFS_SESSION_ID,
          blockId, 0, false, openUfsBlockOptions)) {
        // Read the entire block, caching to block store will be handled internally in UFS block store
        // when close the reader.
        // Note that, we read from UFS with a smaller buffer to avoid high pressure on heap
        // memory when concurrent async requests are received and thus trigger GC.
        long offset = 0;
        while (offset < blockSize) {
          long bufferSize = Math.min(8L * Constants.MB, blockSize - offset);
          reader.read(offset, bufferSize);
          offset += bufferSize;
        }
        if (isTempBlockMetaPresent(blockId)) {
          mBlockWorker.commitBlock(Sessions.CACHE_UFS_SESSION_ID, blockId, false);
        }
      } catch (Exception e) {
        LOG.warn("Failed to async cache block {} from UFS on reading the block:", blockId, e);
        if (isTempBlockMetaPresent(blockId)) {
          mBlockWorker.abortBlock(Sessions.CACHE_UFS_SESSION_ID, blockId);
        }
        return CacheResult.FAILED;
      }
      return CacheResult.SUCCEED;
    }
  
    private boolean isTempBlockMetaPresent(long blockId) {
      Optional<TempBlockMeta> tempBlockMeta =
            mBlockWorker.getBlockStore().getTempBlockMeta(blockId);
      return tempBlockMeta.isPresent()
            && tempBlockMeta.get().getSessionId() == Sessions.CACHE_UFS_SESSION_ID;
    }
  
    /**
     * Caches the block at best effort from a remote worker (possibly from UFS indirectly).
     *
     * @param blockId block ID
     * @param blockSize block size
     * @param sourceAddress the source to read the block previously by client
     * @param openUfsBlockOptions options to open the UFS file
     * @return cache result
     */
    private CacheResult cacheBlockFromRemoteWorker(long blockId, long blockSize,
        InetSocketAddress sourceAddress, Protocol.OpenUfsBlockOptions openUfsBlockOptions)
        throws IOException {
      if (mBlockWorker.getBlockStore().hasBlockMeta(blockId)
          || mBlockWorker.getBlockStore().hasTempBlockMeta(blockId)) {
        // It is already cached
        return CacheResult.ALREADY_CACHED;
      }
      mBlockWorker.createBlock(Sessions.CACHE_WORKER_SESSION_ID, blockId, 0,
          new CreateBlockOptions(null, "", blockSize));
      try (
          BlockReader reader =
              getRemoteBlockReader(blockId, blockSize, sourceAddress, openUfsBlockOptions);
           BlockWriter writer = mBlockWorker
               .createBlockWriter(Sessions.CACHE_WORKER_SESSION_ID, blockId)) {
        BufferUtils.transfer(reader.getChannel(), writer.getChannel());
        mBlockWorker.commitBlock(Sessions.CACHE_WORKER_SESSION_ID, blockId, false);
        return CacheResult.SUCCEED;
      } catch (IllegalStateException | IOException e) {
        LOG.warn("Failed to async cache block {} from remote worker ({}) on copying the block: {}",
            blockId, sourceAddress, e);
        try {
          mBlockWorker.abortBlock(Sessions.CACHE_WORKER_SESSION_ID, blockId);
        } catch (IOException ee) {
          LOG.warn("Failed to abort block {}: {}", blockId, ee.toString());
        }
        throw e;
      }
    }
  
    /**
     * only public to test since we want to mock remoteBlockReader.
     * @param blockId block ID
     * @param blockSize block size
     * @param sourceAddress remote source address
     * @param openUfsBlockOptions options to open the UFS file
     * @return remote block reader
     */
    @VisibleForTesting
    public RemoteBlockReader getRemoteBlockReader(long blockId, long blockSize,
        InetSocketAddress sourceAddress, Protocol.OpenUfsBlockOptions openUfsBlockOptions) {
      return new RemoteBlockReader(mFsContext, blockId, blockSize, sourceAddress,
          openUfsBlockOptions);
    }
  
    // Metrics
    private static final Counter CACHE_REQUESTS =
        MetricsSystem.counter(MetricKey.WORKER_CACHE_REQUESTS.getName());
    private static final Counter CACHE_REQUESTS_ASYNC =
        MetricsSystem.counter(MetricKey.WORKER_CACHE_REQUESTS_ASYNC.getName());
    private static final Counter CACHE_REQUESTS_SYNC =
        MetricsSystem.counter(MetricKey.WORKER_CACHE_REQUESTS_SYNC.getName());
    private static final Counter CACHE_FAILED_BLOCKS =
            MetricsSystem.counter(MetricKey.WORKER_CACHE_FAILED_BLOCKS.getName());
    private static final Counter CACHE_REMOTE_BLOCKS =
            MetricsSystem.counter(MetricKey.WORKER_CACHE_REMOTE_BLOCKS.getName());
    private static final Counter CACHE_SUCCEEDED_BLOCKS =
            MetricsSystem.counter(MetricKey.WORKER_CACHE_SUCCEEDED_BLOCKS.getName());
    private static final Counter CACHE_REJECTED_BLOCKS =
        MetricsSystem.counter(MetricKey.WORKER_CACHE_REJECTED_BLOCKS.getName());
    private static final Counter CACHE_UFS_BLOCKS =
            MetricsSystem.counter(MetricKey.WORKER_CACHE_UFS_BLOCKS.getName());
    private static final Counter CACHE_BLOCKS_SIZE =
        MetricsSystem.counter(MetricKey.WORKER_CACHE_BLOCKS_SIZE.getName());
  }
  