alluxio-master-2.x/core/server/proxy/src/test/java/alluxio/proxy/s3/RateLimitInputStreamTest.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.proxy.s3;

import static alluxio.Constants.KB;
import static alluxio.Constants.MB;

import com.google.common.util.concurrent.RateLimiter;
import org.apache.commons.io.IOUtils;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.FutureTask;
import java.util.stream.Collectors;

public class RateLimitInputStreamTest {

  private byte[] mData;

  @Before
  public void init() throws IOException {
    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(MB);
    int count = 0;
    while (count < MB) {
      byte[] bytes = UUID.randomUUID().toString().getBytes();
      byteArrayOutputStream.write(bytes);
      count += bytes.length;
    }
    mData = Arrays.copyOf(byteArrayOutputStream.toByteArray(), MB);
  }

  @Test
  public void testSingleThreadRead() throws IOException {
    Random random = new Random();
    for (int i = 1; i <= 5; i++) {
      long rate1 = (random.nextInt(4) + 1) * 100 * KB;
      long rate2 = (random.nextInt(4) + 1) * 100 * KB;
      ByteArrayInputStream inputStream = new ByteArrayInputStream(mData);
      RateLimitInputStream rateLimitInputStream = new RateLimitInputStream(inputStream,
          RateLimiter.create(rate1), RateLimiter.create(rate2));
      ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(MB);
      long start = System.currentTimeMillis();
      IOUtils.copy(rateLimitInputStream, byteArrayOutputStream, KB);
      long end = System.currentTimeMillis();
      long duration = end - start;
      long expectedDuration = MB / Math.min(rate1, rate2) * 1000;
      Assert.assertTrue(duration >= expectedDuration && duration <= expectedDuration + 1000);
      Assert.assertArrayEquals(mData, byteArrayOutputStream.toByteArray());
    }
  }

  private void testMultiThreadRead(long globalRate, long rate, int threadNum) {
    long totalSize = (long) threadNum * mData.length;
    RateLimiter globalRateLimiter = RateLimiter.create(globalRate);
    ExecutorService threadPool = Executors.newFixedThreadPool(threadNum);
    List<FutureTask<byte[]>> tasks = new ArrayList<>();
    for (int i = 1; i <= threadNum; i++) {
      tasks.add(new FutureTask<>(() -> {
        ByteArrayInputStream inputStream = new ByteArrayInputStream(mData);
        RateLimitInputStream rateLimitInputStream = new RateLimitInputStream(inputStream,
            RateLimiter.create(rate), globalRateLimiter);
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(MB);
        IOUtils.copy(rateLimitInputStream, byteArrayOutputStream, KB);
        return byteArrayOutputStream.toByteArray();
      }));
    }
    long start = System.currentTimeMillis();
    tasks.forEach(threadPool::submit);
    List<byte[]> results;
    try {
      results = tasks.stream().map(task -> {
        try {
          return task.get();
        } catch (Exception e) {
          throw new RuntimeException(e);
        }
      }).collect(Collectors.toList());
    } finally {
      threadPool.shutdownNow();
    }
    long end = System.currentTimeMillis();
    long duration = end - start;
    long expectedDuration = totalSize / Math.min(globalRate, (long) threadNum * rate) * 1000;
    Assert.assertTrue(duration >= expectedDuration && duration <= expectedDuration + 1000);
    results.forEach(bytes -> Assert.assertArrayEquals(mData, bytes));
  }

  @Test
  public void testMultiThreadReadWithBiggerGlobalRate() {
    testMultiThreadRead(400 * KB, 100 * KB, 3);
  }

  @Test
  public void testMultiThreadReadWithSmallerGlobalRate() {
    testMultiThreadRead(100 * KB, 200 * KB, 3);
  }
}
