alluxio-master-2.x/core/server/worker/src/main/java/alluxio/worker/block/UfsInputStreamCache.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.Constants;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.exception.runtime.AlluxioRuntimeException;
import alluxio.underfs.SeekableUnderFileInputStream;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.options.OpenOptions;
import alluxio.util.IdUtils;
import alluxio.util.executor.ExecutorServiceFactories;
import alluxio.util.logging.SamplingLogger;
import alluxio.worker.block.io.BlockReader;

import com.google.common.base.Preconditions;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalListeners;
import com.google.common.cache.RemovalNotification;
import com.google.common.util.concurrent.UncheckedExecutionException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import javax.annotation.concurrent.ThreadSafe;

/**
 * This class implements a {@link BlockReader} to read a block directly from UFS, and
 * optionally cache the block to the Alluxio worker if the whole block it is read.
 */
@ThreadSafe
public final class UfsInputStreamCache {
  private static final Logger LOG = LoggerFactory.getLogger(UfsInputStreamCache.class);
  private static final Logger SAMPLING_LOG = new SamplingLogger(LOG, 10L * Constants.MINUTE_MS);
  private static final boolean CACHE_ENABLED =
      Configuration.getBoolean(PropertyKey.WORKER_UFS_INSTREAM_CACHE_ENABLED);

  /**
   * A map from the ufs file id to the metadata of the input streams. Synchronization on this map
   * before access.
   */
  private final Map<Long, StreamIdSet> mFileIdToStreamIds;
  /** Cache of the input streams, from the input stream id to the input stream. */
  private final Cache<Long, CachedSeekableInputStream> mStreamCache;

  /**
   * Constructs a new UFS input stream cache.
   */
  public UfsInputStreamCache() {
    mFileIdToStreamIds = new ConcurrentHashMap<>();
    /* Thread pool for asynchronously removing the expired input streams. */
    ExecutorService removalThreadPool =
        ExecutorServiceFactories.fixedThreadPool(Constants.UFS_INPUT_STREAM_CACHE_EXPIRATION, 2)
            .create();

    // A listener to the input stream removal.
    RemovalListener<Long, CachedSeekableInputStream> listener =
        (RemovalNotification<Long, CachedSeekableInputStream> removal) -> {
          CachedSeekableInputStream inputStream = removal.getValue();
          final long fileId = inputStream.getFileId();
          final long resourceId = removal.getKey();
          boolean shouldClose = false;

          StreamIdSet streamIds = mFileIdToStreamIds.get(fileId);
          if (streamIds == null) {
            LOG.warn(
                "Removed UFS input stream (fileId: {} resourceId: {}) but does not exist",
                fileId, resourceId);
          } else {
            synchronized (streamIds) {
              // remove the key
              if (streamIds.removeInUse(resourceId)) {
                LOG.warn("Removed in-use UFS input stream (fileId: {} resourceId: {})", fileId,
                    resourceId);
              }
              if (streamIds.removeAvailable(resourceId)) {
                // close the resource
                LOG.debug("Removed available UFS input stream (fileId: {} resourceId: {})", fileId,
                    resourceId);
                shouldClose = true;
              }
              if (streamIds.isEmpty()) {
                // remove the value from the mapping
                mFileIdToStreamIds.remove(fileId);
              }
            }
          }

          if (shouldClose) {
            try {
              inputStream.close();
            } catch (IOException e) {
              LOG.warn("Failed to close UFS input stream resource of file {} with file id {}"
                      + " and resource id {}", inputStream.getFilePath(), inputStream.getFileId(),
                  resourceId);
            }
          }
        };
    mStreamCache = CacheBuilder.newBuilder()
        .maximumSize(Configuration.getInt(PropertyKey.WORKER_UFS_INSTREAM_CACHE_MAX_SIZE))
        .expireAfterAccess(
            Configuration.getMs(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME),
            TimeUnit.MILLISECONDS)
        .removalListener(RemovalListeners.asynchronous(listener, removalThreadPool)).build();
  }

  /**
   * Releases an input stream. The input stream is closed if it's already expired.
   *
   * @param inputStream the input stream to release
   * @throws IOException when input stream fails to close
   */
  public void release(InputStream inputStream) {
    if (!(inputStream instanceof CachedSeekableInputStream) || !CACHE_ENABLED) {
      // for non-seekable input stream, close and return
      close(inputStream);
      return;
    }

    CachedSeekableInputStream cachedStream = (CachedSeekableInputStream) inputStream;
    long fileId = cachedStream.getFileId();
    long resourceId = cachedStream.getResourceId();
    StreamIdSet streamIds = mFileIdToStreamIds.get(fileId);

    if (streamIds == null) {
      // the cache no longer tracks this input stream
      LOG.debug("UFS input stream (fileId: {} resourceId: {}) is already expired", fileId,
          resourceId);
      close(inputStream);
      return;
    }

    if (!streamIds.release(resourceId)) {
      LOG.debug("Close the expired UFS input stream (fileId: {} resourceId: {})", fileId,
          resourceId);
      // the input stream expired, close it
      close(inputStream);
    }
  }

  private void close(InputStream in) {
    try {
      in.close();
    } catch (IOException e) {
      throw AlluxioRuntimeException.from(e);
    }
  }

  /**
   * Acquires an input stream. For seekable input streams, if there is an available input stream in
   * the cache, reuse it and repositions the offset, otherwise the manager opens a new input stream.
   *
   * @param ufs the under file system
   * @param path the path to the under storage file
   * @param fileId the file id
   * @param openOptions the open options
   * @return the acquired input stream
   * @throws IOException if the input stream fails to open
   */
  public InputStream acquire(UnderFileSystem ufs, String path, long fileId, OpenOptions openOptions)
      throws IOException {
    if (!ufs.isSeekable() || !CACHE_ENABLED) {
      // only seekable UFSes are cachable/reusable, always return a new input stream
      return ufs.openExistingFile(path, openOptions);
    }

    // explicit cache cleanup
    try {
      mStreamCache.cleanUp();
    } catch (Throwable error) {
      SAMPLING_LOG.warn("Explicit cache removal failed.", error);
    }

    StreamIdSet streamIds = mFileIdToStreamIds.compute(fileId, (key, value) -> {
      if (value != null) {
        return value;
      }
      return new StreamIdSet();
    });

    // Try to acquire an existing id from the stream id set.
    // synchronized is required to be consistent between availableIds() and acquire(id).
    CachedSeekableInputStream inputStream = null;
    synchronized (streamIds) {
      // find the next available input stream from the cache
      for (long id : streamIds.availableIds()) {
        inputStream = mStreamCache.getIfPresent(id);
        if (inputStream != null) {
          // acquire it now while locked, so other threads cannot take it
          streamIds.acquire(id);
          break;
        }
      }
    }

    if (inputStream != null) {
      // for the cached ufs instream, seek (outside of critical section) to the requested position.
      LOG.debug("Reused the under file input stream resource of {}", inputStream.getResourceId());
      inputStream.seek(openOptions.getOffset());
      return inputStream;
    }

    // no cached input stream is available, acquire a new id and open a new stream
    final long newId = streamIds.acquireNewId();
    try {
      inputStream = mStreamCache.get(newId, () -> {
        SeekableUnderFileInputStream ufsStream = (SeekableUnderFileInputStream) ufs
            .openExistingFile(path,
                OpenOptions.defaults().setPositionShort(openOptions.getPositionShort())
                    .setOffset(openOptions.getOffset()));
        LOG.debug("Created the under file input stream resource of {}", newId);
        return new CachedSeekableInputStream(ufsStream, newId, fileId, path);
      });
    } catch (ExecutionException e) {
      LOG.warn("Failed to create a new cached ufs instream of file id {} and path {}", fileId,
          path, e);
      // fall back to an uncached ufs creation.
      return ufs.openExistingFile(path,
          OpenOptions.defaults().setOffset(openOptions.getOffset()));
    }
    catch (UncheckedExecutionException e) {
      throw AlluxioRuntimeException.from(e.getCause());
    }
    return inputStream;
  }

  /**
   * The metadata of the input streams associated with an under storage file that tracks which input
   * streams are in-use or available. Each input stream is identified by a unique id.
   */
  @ThreadSafe
  private static class StreamIdSet {
    private final Set<Long> mInUseStreamIds;
    private final Set<Long> mAvailableStreamIds;

    /**
     * Creates a new {@link StreamIdSet}.
     */
    StreamIdSet() {
      mInUseStreamIds = new HashSet<>();
      mAvailableStreamIds = new HashSet<>();
    }

    /**
     * @return a view of the available input stream ids
     */
    synchronized Set<Long> availableIds() {
      return Collections.unmodifiableSet(mAvailableStreamIds);
    }

    /**
     * Marks an input stream as acquired.
     *
     * @param id the id of the input stream
     */
    synchronized void acquire(long id) {
      Preconditions.checkArgument(!mInUseStreamIds.contains(id), "%s is already in use", id);
      mAvailableStreamIds.remove(id);
      mInUseStreamIds.add(id);
    }

    synchronized long acquireNewId() {
      while (true) {
        long newId = IdUtils.getRandomNonNegativeLong();
        if (mAvailableStreamIds.contains(newId) || mInUseStreamIds.contains(newId)) {
          // This id already managed, try again.
          continue;
        }

        acquire(newId);
        return newId;
      }
    }

    /**
     * @return if there is any outstanding input streams of the file
     */
    synchronized boolean isEmpty() {
      return mInUseStreamIds.isEmpty() && mAvailableStreamIds.isEmpty();
    }

    /**
     * Marks an input stream as not in use.
     * @param id the id of the input stream
     * @return true if removed from the in-use streams
     */
    synchronized boolean removeInUse(long id) {
      return mInUseStreamIds.remove(id);
    }

    /**
     * Removes the mark of the input stream as available.
     *
     * @param id the id of the input stream
     * @return <code>true</code> if the given input stream is available, <code>false</code> if the
     *         given input stream is not among available input streams
     */
    synchronized boolean removeAvailable(long id) {
      return mAvailableStreamIds.remove(id);
    }

    /**
     * Returns an id to the input stream pool. It marks the id from in use to available. If the
     * input stream is already removed from the cache, then do nothing.
     *
     * @param id id of the input stream
     * @return true if the id is marked from in use to available; false if the id no longer used for
     *         cache
     */
    synchronized boolean release(long id) {
      Preconditions.checkArgument(!mAvailableStreamIds.contains(id));
      if (mInUseStreamIds.contains(id)) {
        mInUseStreamIds.remove(id);
        mAvailableStreamIds.add(id);
        return true;
      }
      return false;
    }
  }
}
