alluxio-master-2.x/core/server/worker/src/test/java/alluxio/worker/block/UfsInputStreamCacheTest.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 static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockingDetails;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import alluxio.ConfigurationRule;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.test.util.ConcurrencyUtils;
import alluxio.underfs.SeekableUnderFileInputStream;
import alluxio.underfs.UnderFileSystem;
import alluxio.underfs.options.OpenOptions;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.invocation.Invocation;

import java.io.Closeable;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;

public final class UfsInputStreamCacheTest {
  private static final String FILE_NAME = "/test";
  private static final long FILE_ID = 1;

  private UnderFileSystem mUfs;
  private SeekableUnderFileInputStream[] mSeekableInStreams;
  private UfsInputStreamCache mManager;
  private int mNumOfInputStreams = 20;

  @Before
  public void before() throws Exception {
    mSeekableInStreams = new SeekableUnderFileInputStream[mNumOfInputStreams];
    mUfs = mock(UnderFileSystem.class);
    when(mUfs.isSeekable()).thenReturn(true);
    for (int i = 0; i < mNumOfInputStreams; i++) {
      SeekableUnderFileInputStream instream = mock(SeekableUnderFileInputStream.class);
      mSeekableInStreams[i] = instream;
    }
    when(mUfs.openExistingFile(eq(FILE_NAME), any(OpenOptions.class))).thenReturn(
        mSeekableInStreams[0], Arrays.copyOfRange(mSeekableInStreams, 1, mNumOfInputStreams));
    mManager = new UfsInputStreamCache();
  }

  @Test
  public void notSeekable() throws Exception {
    when(mUfs.isSeekable()).thenReturn(false);

    SeekableUnderFileInputStream mockedStream = mock(SeekableUnderFileInputStream.class);
    when(mUfs.openExistingFile(eq(FILE_NAME), any(OpenOptions.class)))
        .thenReturn(mockedStream).thenThrow(new IllegalStateException("Should be called once"));

    // acquire a stream
    InputStream instream1 =
        mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
    // release
    mManager.release(instream1);

    // ensure the second time the released instream is the same one but repositioned
    verify(mockedStream).close();
  }

  @Test
  public void acquireAndRelease() throws Exception {
    SeekableUnderFileInputStream mockedStream = mock(SeekableUnderFileInputStream.class);
    when(mUfs.openExistingFile(eq(FILE_NAME), any(OpenOptions.class)))
        .thenReturn(mockedStream).thenThrow(new IllegalStateException("Should be called once"));

    // acquire a stream
    InputStream instream1 =
        mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
    // release
    mManager.release(instream1);
    // acquire a stream again
    InputStream instream2 =
        mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(4));

    Assert.assertEquals(instream1, instream2);
    // ensure the second time the released instream is the same one but repositioned
    verify(mockedStream).seek(4);
  }

  @Test
  public void multipleCheckIn() throws Exception {
    mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
    mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(4));
    mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(6));
    // 3 different input streams are acquired
    verify(mUfs, times(3)).openExistingFile(eq(FILE_NAME),
        any(OpenOptions.class));
  }

  @Test
  public void expire() throws Exception {
    try (Closeable r = new ConfigurationRule(new HashMap<PropertyKey, Object>() {
      {
        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME, "2");
      }
    }, Configuration.modifiableGlobal()).toResource()) {
      mManager = new UfsInputStreamCache();
      // check out a stream
      InputStream instream =
          mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
      // check in back
      mManager.release(instream);
      Thread.sleep(10);
      // check out another stream should trigger the timeout
      mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(4));
      verify(mSeekableInStreams[0], timeout(2000).times(1)).close();
    }
  }

  @Test
  public void releaseExpiredSameFile() throws Exception {
    try (Closeable r = new ConfigurationRule(new HashMap<PropertyKey, Object>() {
      {
        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME, "2");
      }
    }, Configuration.modifiableGlobal()).toResource()) {
      mManager = new UfsInputStreamCache();
      // check out a stream
      InputStream instream =
          mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
      // wait a bit for a stream to be expired
      Thread.sleep(100);
      // check out another stream should trigger the timeout
      mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(4));

      // wait a bit so release occurs after removal listener
      Thread.sleep(100);
      mManager.release(instream);

      // verify the stream was closed once
      verify(mSeekableInStreams[0], timeout(2000).times(1)).close();
    }
  }

  @Test
  public void releaseExpiredDifferentFile() throws Exception {
    try (Closeable r = new ConfigurationRule(new HashMap<PropertyKey, Object>() {
      {
        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME, "2");
      }
    }, Configuration.modifiableGlobal()).toResource()) {
      mManager = new UfsInputStreamCache();
      // check out a stream
      InputStream instream =
          mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(2));
      // wait a bit for a stream to be expired
      Thread.sleep(100);
      // check out another stream should trigger the timeout
      mManager.acquire(mUfs, FILE_NAME + "2", FILE_ID + 1, OpenOptions.defaults().setOffset(4));

      // wait a bit so release occurs after removal listener
      Thread.sleep(100);
      mManager.release(instream);

      // verify the stream was closed once
      verify(mSeekableInStreams[0], timeout(2000).times(1)).close();
    }
  }

  @Test
  public void concurrency() throws Exception {
    try (Closeable r = new ConfigurationRule(new HashMap<PropertyKey, Object>() {
      {
        // use very large number
        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME, "200000");
      }
    }, Configuration.modifiableGlobal()).toResource()) {
      mManager = new UfsInputStreamCache();
      List<Thread> threads = new ArrayList<>();
      int numCheckOutPerThread = 10;
      for (int i = 0; i < mNumOfInputStreams / 2; i++) {
        Runnable runnable = () -> {
          for (int j = 0; j < numCheckOutPerThread; j++) {
            InputStream instream;
            try {
              instream =
                  mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(j));
              Thread.sleep(10);
              mManager.release(instream);
            } catch (Exception e) {
              // the input streams created should not be more than mNumOfInputStreams
              Assert.fail("input stream check in and out failed." + e);
            }
          }
        };
        threads.add(new Thread(runnable));
      }
      ConcurrencyUtils.assertConcurrent(threads, 30);
      // Each subsequent check out per thread should be a seek operation
      int numSeek = 0;
      for (int i = 0; i < mNumOfInputStreams; i++) {
        for (Invocation invocation : mockingDetails(mSeekableInStreams[i])
            .getInvocations()) {
          if (invocation.getMethod().getName().equals("seek")) {
            numSeek++;
          }
        }
      }
      Assert.assertEquals(mNumOfInputStreams / 2 * (numCheckOutPerThread - 1), numSeek);
    }
  }

  @Test
  public void concurrencyWithExpiration() throws Exception {
    try (Closeable r = new ConfigurationRule(new HashMap<PropertyKey, Object>() {
      {
        put(PropertyKey.WORKER_UFS_INSTREAM_CACHE_EXPIRARTION_TIME, "20");
      }
    }, Configuration.modifiableGlobal()).toResource()) {
      mManager = new UfsInputStreamCache();
      List<Thread> threads = new ArrayList<>();
      int numCheckOutPerThread = 4;
      for (int i = 0; i < mNumOfInputStreams / numCheckOutPerThread; i++) {
        Runnable runnable = () -> {
          for (int j = 0; j < numCheckOutPerThread; j++) {
            InputStream instream;
            try {
              instream =
                  mManager.acquire(mUfs, FILE_NAME, FILE_ID, OpenOptions.defaults().setOffset(j));
              mManager.release(instream);
              Thread.sleep(200);
            } catch (Exception e) {
              // the input streams created should not be more than mNumOfInputStreams
              Assert.fail("input stream check in and out failed." + e);
            }
          }
        };
        threads.add(new Thread(runnable));
      }
      ConcurrencyUtils.assertConcurrent(threads, 30);
      // ensure at least one expired in stream is closed
      verify(mSeekableInStreams[0], timeout(2000)).close();
    }
  }
}
