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.client.block;
  
  import static java.util.stream.Collectors.toList;
  import static java.util.stream.Collectors.toSet;
  
  import alluxio.client.WriteType;
  import alluxio.client.block.policy.BlockLocationPolicy;
  import alluxio.client.block.policy.options.GetWorkerOptions;
  import alluxio.client.block.stream.BlockInStream;
  import alluxio.client.block.stream.BlockInStream.BlockInStreamSource;
  import alluxio.client.block.stream.BlockOutStream;
  import alluxio.client.block.stream.DataWriter;
  import alluxio.client.block.util.BlockLocationUtils;
  import alluxio.client.file.FileSystemContext;
  import alluxio.client.file.URIStatus;
  import alluxio.client.file.options.InStreamOptions;
  import alluxio.client.file.options.OutStreamOptions;
  import alluxio.collections.Pair;
  import alluxio.exception.ExceptionMessage;
  import alluxio.exception.PreconditionMessage;
  import alluxio.exception.status.UnavailableException;
  import alluxio.network.TieredIdentityFactory;
  import alluxio.resource.CloseableResource;
  import alluxio.wire.BlockInfo;
  import alluxio.wire.BlockLocation;
  import alluxio.wire.TieredIdentity;
  import alluxio.wire.WorkerNetAddress;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.base.Preconditions;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Lists;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.text.MessageFormat;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * Alluxio Block Store client. This is an internal client for all block level operations in Alluxio.
   * An instance of this class can be obtained via {@link BlockStoreClient} constructors.
   */
  @ThreadSafe
  public final class BlockStoreClient {
    private static final Logger LOG = LoggerFactory.getLogger(BlockStoreClient.class);
  
    private final FileSystemContext mContext;
    private final TieredIdentity mTieredIdentity;
  
    /**
     * Creates an Alluxio block store with default local hostname.
     *
     * @param context the file system context
     * @return the {@link BlockStoreClient} created
     */
    public static BlockStoreClient create(FileSystemContext context) {
      return new BlockStoreClient(context,
          TieredIdentityFactory.localIdentity(context.getClusterConf()));
    }
  
    /**
     * Creates an Alluxio block store.
     *
     * @param context the file system context
     * @param tieredIdentity the tiered identity
     */
    @VisibleForTesting
    BlockStoreClient(FileSystemContext context, TieredIdentity tieredIdentity) {
      mContext = context;
      mTieredIdentity = tieredIdentity;
    }
  
    /**
     * Gets the block info of a block, if it exists.
     *
     * @param blockId the blockId to obtain information about
     * @return a {@link BlockInfo} containing the metadata of the block
     */
    public BlockInfo getInfo(long blockId) throws IOException {
      try (CloseableResource<BlockMasterClient> masterClientResource =
          mContext.acquireBlockMasterClientResource()) {
        return masterClientResource.get().getBlockInfo(blockId);
      }
    }
  
    /**
     * Gets a stream to read the data of a block. This method is primarily responsible for
     * determining the data source and type of data source. The latest BlockInfo will be fetched
     * from the master to ensure the locations are up-to-date.
     *
     * @param blockId the id of the block to read
     * @param options the options associated with the read request
     * @return a stream which reads from the beginning of the block
     */
    public BlockInStream getInStream(long blockId, InStreamOptions options) throws IOException {
      return getInStream(blockId, options, ImmutableMap.of());
    }
  
    /**
     * Gets a stream to read the data of a block. This method is primarily responsible for
     * determining the data source and type of data source. The latest BlockInfo will be fetched
     * from the master to ensure the locations are up-to-date. It takes a map of failed workers and
     * their most recently failed time and tries to update it when BlockInStream created failed,
     * attempting to avoid reading from a recently failed worker.
     *
     * @param blockId the id of the block to read
     * @param options the options associated with the read request
     * @param failedWorkers the map of worker's addresses to most recent failure time
     * @return a stream which reads from the beginning of the block
     */
    public BlockInStream getInStream(long blockId, InStreamOptions options,
        Map<WorkerNetAddress, Long> failedWorkers) throws IOException {
      // Get the latest block info from master
      BlockInfo info = getInfo(blockId);
      return getInStream(info, options, failedWorkers);
    }
  
    /**
     * {@link #getInStream(long, InStreamOptions, Map)}.
     *
     * @param info the block info
     * @param options the options associated with the read request
     * @param failedWorkers the map of worker's addresses to most recent failure time
     * @return a stream which reads from the beginning of the block
     */
    public BlockInStream getInStream(BlockInfo info, InStreamOptions options,
        Map<WorkerNetAddress, Long> failedWorkers) throws IOException {
      Pair<WorkerNetAddress, BlockInStreamSource> dataSourceAndType = getDataSourceAndType(info,
          options.getStatus(), options.getUfsReadLocationPolicy(), failedWorkers);
      WorkerNetAddress dataSource = dataSourceAndType.getFirst();
      BlockInStreamSource dataSourceType = dataSourceAndType.getSecond();
      try {
        return BlockInStream.create(mContext, info, dataSource, dataSourceType, options);
      } catch (UnavailableException e) {
        //When BlockInStream created failed, it will update the passed-in failedWorkers
        //to attempt to avoid reading from this failed worker in next try.
        LOG.info("Added {} to failedWorkers for {}", dataSource, e.toString());
        failedWorkers.put(dataSource, System.currentTimeMillis());
        throw e;
      }
    }
  
    /**
     * Gets the data source and type of data source of a block. This method is primarily responsible
     * for determining the data source and type of data source. It takes a map of failed workers and
     * their most recently failed time and tries to update it when BlockInStream created failed,
     * attempting to avoid reading from a recently failed worker.
     *
     * @param info the info of the block to read
     * @param status the URIStatus associated with the read request
     * @param policy the policy determining the Alluxio worker location
     * @param failedWorkers the map of worker's addresses to most recent failure time
     * @return the data source and type of data source of the block
     */
    public Pair<WorkerNetAddress, BlockInStreamSource> getDataSourceAndType(BlockInfo info,
        URIStatus status, BlockLocationPolicy policy, Map<WorkerNetAddress, Long> failedWorkers)
        throws IOException {
      List<BlockLocation> locations = info.getLocations();
      List<BlockWorkerInfo> blockWorkerInfo = Collections.emptyList();
      // Initial target workers to read the block given the block locations.
      Set<WorkerNetAddress> workerPool;
      // Note that, it is possible that the blocks have been written as UFS blocks
      if (status.isPersisted()
          || status.getPersistenceState().equals("TO_BE_PERSISTED")) {
        blockWorkerInfo = mContext.getCachedWorkers();
        if (blockWorkerInfo.isEmpty()) {
          throw new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage());
        }
        workerPool = blockWorkerInfo.stream().map(BlockWorkerInfo::getNetAddress).collect(toSet());
      } else {
        if (locations.isEmpty()) {
          blockWorkerInfo = mContext.getCachedWorkers();
          if (blockWorkerInfo.isEmpty()) {
            throw new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage());
          }
          throw new UnavailableException(MessageFormat
              .format("Block {0} is unavailable in both Alluxio and UFS.", info.getBlockId()));
        }
        workerPool = locations.stream().map(BlockLocation::getWorkerAddress).collect(toSet());
      }
      // Workers to read the block, after considering failed workers.
      Set<WorkerNetAddress> workers = handleFailedWorkers(workerPool, failedWorkers);
      // TODO(calvin, jianjian): Consider containing these two variables in one object
      BlockInStreamSource dataSourceType = null;
      WorkerNetAddress dataSource = null;
      locations = locations.stream()
          .filter(location -> workers.contains(location.getWorkerAddress())).collect(toList());
      // First try to read data from Alluxio
      if (!locations.isEmpty()) {
        // TODO(calvin): Get location via a policy
        List<WorkerNetAddress> tieredLocations =
            locations.stream().map(BlockLocation::getWorkerAddress)
                .collect(toList());
        Collections.shuffle(tieredLocations);
        Optional<Pair<WorkerNetAddress, Boolean>> nearest =
            BlockLocationUtils.nearest(mTieredIdentity, tieredLocations, mContext.getClusterConf());
        if (nearest.isPresent()) {
          dataSource = nearest.get().getFirst();
          dataSourceType = nearest.get().getSecond() ? mContext.hasProcessLocalWorker()
              ? BlockInStreamSource.PROCESS_LOCAL : BlockInStreamSource.NODE_LOCAL
              : BlockInStreamSource.REMOTE;
        }
      }
      // Can't get data from Alluxio, get it from the UFS instead
      if (dataSource == null) {
        dataSourceType = BlockInStreamSource.UFS;
        Preconditions.checkNotNull(policy, "The UFS read location policy is not specified");
        blockWorkerInfo = blockWorkerInfo.stream()
            .filter(workerInfo -> workers.contains(workerInfo.getNetAddress())).collect(toList());
        GetWorkerOptions getWorkerOptions = GetWorkerOptions.defaults()
            .setBlockInfo(new BlockInfo()
                .setBlockId(info.getBlockId())
                .setLength(info.getLength())
                .setLocations(locations))
            .setBlockWorkerInfos(blockWorkerInfo);
        dataSource = policy.getWorker(getWorkerOptions).orElseThrow(
            () -> new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage())
        );
        if (mContext.hasProcessLocalWorker()
            && dataSource.equals(mContext.getNodeLocalWorker())) {
          dataSourceType = BlockInStreamSource.PROCESS_LOCAL;
          LOG.debug("Create BlockInStream to read data from UFS through process local worker {}",
              dataSource);
        } else {
          LOG.debug("Create BlockInStream to read data from UFS through worker {} "
                  + "(client embedded in local worker process: {},"
                  + "client co-located with worker in different processes: {}, "
                  + "local worker address: {})",
              dataSource, mContext.hasProcessLocalWorker(), mContext.hasNodeLocalWorker(),
              mContext.hasNodeLocalWorker() ? mContext.getNodeLocalWorker() : "N/A");
        }
      }
      return new Pair<>(dataSource, dataSourceType);
    }
  
    private Set<WorkerNetAddress> handleFailedWorkers(Set<WorkerNetAddress> workers,
        Map<WorkerNetAddress, Long> failedWorkers) {
      if (workers.isEmpty()) {
        return Collections.emptySet();
      }
      Set<WorkerNetAddress> nonFailed =
          workers.stream().filter(worker -> !failedWorkers.containsKey(worker)).collect(toSet());
      if (nonFailed.isEmpty()) {
        return Collections.singleton(workers.stream()
            .min(Comparator.comparingLong(failedWorkers::get)).get());
      }
      return nonFailed;
    }
  
    /**
     * Gets a stream to write data to a block. The stream can only be backed by Alluxio storage.
     *
     * @param blockId the block to write
     * @param blockSize the standard block size to write
     * @param address the address of the worker to write the block to, fails if the worker cannot
     *        serve the request
     * @param options the output stream options
     * @return an {@link BlockOutStream} which can be used to write data to the block in a streaming
     *         fashion
     */
    public BlockOutStream getOutStream(long blockId, long blockSize, WorkerNetAddress address,
        OutStreamOptions options) throws IOException {
      // No specified location to write to.
      Preconditions.checkNotNull(address, "address");
      LOG.debug("Create BlockOutStream for {} of block size {} at address {}, using options: {}",
          blockId, blockSize, address, options);
      DataWriter dataWriter =
          DataWriter.Factory.create(mContext, blockId, blockSize, address, options);
      return new BlockOutStream(dataWriter, blockSize, address);
    }
  
    /**
     * Gets a stream to write data to a block based on the options. The stream can only be backed by
     * Alluxio storage.
     *
     * @param blockId the block to write
     * @param blockSize the standard block size to write
     * @param options the output stream option
     * @return a {@link BlockOutStream} which can be used to write data to the block in a streaming
     *         fashion
     */
    public BlockOutStream getOutStream(long blockId, long blockSize, OutStreamOptions options)
        throws IOException {
      WorkerNetAddress address;
      BlockLocationPolicy locationPolicy = Preconditions.checkNotNull(options.getLocationPolicy(),
          PreconditionMessage.BLOCK_WRITE_LOCATION_POLICY_UNSPECIFIED);
      GetWorkerOptions workerOptions = GetWorkerOptions.defaults()
          .setBlockInfo(new BlockInfo().setBlockId(blockId).setLength(blockSize))
          .setBlockWorkerInfos(new ArrayList<>(mContext.getCachedWorkers()));
  
      // The number of initial copies depends on the write type: if ASYNC_THROUGH, it is the property
      // "alluxio.user.file.replication.durable" before data has been persisted; otherwise
      // "alluxio.user.file.replication.min"
      int initialReplicas = (options.getWriteType() == WriteType.ASYNC_THROUGH
          && options.getReplicationDurable() > options.getReplicationMin())
          ? options.getReplicationDurable() : options.getReplicationMin();
      if (initialReplicas <= 1) {
        address = locationPolicy.getWorker(workerOptions).orElseThrow(
            () -> {
              try {
                if (mContext.getCachedWorkers().isEmpty()) {
                  return new UnavailableException(ExceptionMessage.NO_WORKER_AVAILABLE.getMessage());
                }
              } catch (IOException e) {
                return e;
              }
              return new UnavailableException(
                  ExceptionMessage.NO_SPACE_FOR_BLOCK_ON_WORKER.getMessage(blockSize));
            }
        );
        // TODO(ggezer): Retry on another worker if this has no storage.
        return getOutStream(blockId, blockSize, address, options);
      }
  
      // Group different block workers by their hostnames
      Map<String, Set<BlockWorkerInfo>> blockWorkersByHost = new HashMap<>();
      for (BlockWorkerInfo blockWorker : workerOptions.getBlockWorkerInfos()) {
        String hostName = blockWorker.getNetAddress().getHost();
        if (blockWorkersByHost.containsKey(hostName)) {
          blockWorkersByHost.get(hostName).add(blockWorker);
        } else {
          blockWorkersByHost.put(hostName, com.google.common.collect.Sets.newHashSet(blockWorker));
        }
      }
  
      // Select N workers on different hosts where N is the value of initialReplicas for this block
      List<WorkerNetAddress> workerAddressList = new ArrayList<>();
      List<BlockWorkerInfo> updatedInfos = Lists.newArrayList(workerOptions.getBlockWorkerInfos());
      for (int i = 0; i < initialReplicas; i++) {
        locationPolicy.getWorker(workerOptions).ifPresent(workerAddress -> {
          workerAddressList.add(workerAddress);
          updatedInfos.removeAll(blockWorkersByHost.get(workerAddress.getHost()));
          workerOptions.setBlockWorkerInfos(updatedInfos);
        });
      }
      if (workerAddressList.size() < initialReplicas) {
        throw new alluxio.exception.status.ResourceExhaustedException(String.format(
            "Not enough workers for replications, %d workers selected but %d required",
            workerAddressList.size(), initialReplicas));
      }
      return BlockOutStream
          .createReplicatedBlockOutStream(mContext, blockId, blockSize, workerAddressList, options);
    }
  
    /**
     * Gets the total capacity of Alluxio's BlockStore.
     *
     * @return the capacity in bytes
     */
    public long getCapacityBytes() throws IOException {
      try (CloseableResource<BlockMasterClient> blockMasterClientResource =
          mContext.acquireBlockMasterClientResource()) {
        return blockMasterClientResource.get().getCapacityBytes();
      }
    }
  
    /**
     * Gets the used bytes of Alluxio's BlockStore.
     *
     * @return the used bytes of Alluxio's BlockStore
     */
    public long getUsedBytes() throws IOException {
      try (CloseableResource<BlockMasterClient> blockMasterClientResource =
          mContext.acquireBlockMasterClientResource()) {
        return blockMasterClientResource.get().getUsedBytes();
      }
    }
  }
  