Examples with MockRateLimiter com.google.common.util.concurrent.MockRateLimiter used on opensource projects

Example 1 with MockRateLimiter

use of com.google.common.util.concurrent.MockRateLimiter in project alluxio by Alluxio.

the class FileDataManagerTest method persistFileRateLimiting.

/**
   * Tests the rate limiting functionality for asynchronous persistence.
   */
@Test
public void persistFileRateLimiting() throws Exception {
    Configuration.set(PropertyKey.WORKER_FILE_PERSIST_RATE_LIMIT_ENABLED, "true");
    Configuration.set(PropertyKey.WORKER_FILE_PERSIST_RATE_LIMIT, "100");
    mMockRateLimiter = new MockRateLimiter(Configuration.getBytes(PropertyKey.WORKER_FILE_PERSIST_RATE_LIMIT));
    mManager = new FileDataManager(mBlockWorker, mMockRateLimiter.getGuavaRateLimiter());
    long fileId = 1;
    List<Long> blockIds = Lists.newArrayList(1L, 2L, 3L);
    FileInfo fileInfo = new FileInfo();
    fileInfo.setPath("test");
    Mockito.when(mBlockWorker.getFileInfo(fileId)).thenReturn(fileInfo);
    BlockReader reader = Mockito.mock(BlockReader.class);
    for (long blockId : blockIds) {
        Mockito.when(mBlockWorker.lockBlock(Sessions.CHECKPOINT_SESSION_ID, blockId)).thenReturn(blockId);
        Mockito.when(mBlockWorker.readBlockRemote(Sessions.CHECKPOINT_SESSION_ID, blockId, blockId)).thenReturn(reader);
        BlockMeta mockedBlockMeta = PowerMockito.mock(BlockMeta.class);
        Mockito.when(mockedBlockMeta.getBlockSize()).thenReturn(100L);
        Mockito.when(mBlockWorker.getBlockMeta(Sessions.CHECKPOINT_SESSION_ID, blockId, blockId)).thenReturn(mockedBlockMeta);
    }
    String ufsRoot = Configuration.get(PropertyKey.UNDERFS_ADDRESS);
    Mockito.when(mUfs.isDirectory(ufsRoot)).thenReturn(true);
    OutputStream outputStream = Mockito.mock(OutputStream.class);
    // mock BufferUtils
    PowerMockito.mockStatic(BufferUtils.class);
    String dstPath = PathUtils.concatPath(ufsRoot, fileInfo.getPath());
    fileInfo.setUfsPath(dstPath);
    Mockito.when(mUfs.create(dstPath)).thenReturn(outputStream);
    Mockito.when(mUfs.create(Mockito.anyString(), Mockito.any(CreateOptions.class))).thenReturn(outputStream);
    Mockito.when(mMockFileSystem.getStatus(Mockito.any(AlluxioURI.class))).thenReturn(new URIStatus(fileInfo));
    mManager.lockBlocks(fileId, blockIds);
    mManager.persistFile(fileId, blockIds);
    List<String> expectedEvents = Lists.newArrayList("R0.00", "R1.00", "R1.00");
    assertEquals(expectedEvents, mMockRateLimiter.readEventsAndClear());
    // Simulate waiting for 1 second.
    mMockRateLimiter.sleepMillis(1000);
    mManager.lockBlocks(fileId, blockIds);
    mManager.persistFile(fileId, blockIds);
    // The first write will go through immediately without throttling.
    expectedEvents = Lists.newArrayList("U1.00", "R0.00", "R1.00", "R1.00");
    assertEquals(expectedEvents, mMockRateLimiter.readEventsAndClear());
    // Repeat persistence without sleeping.
    mManager.lockBlocks(fileId, blockIds);
    mManager.persistFile(fileId, blockIds);
    expectedEvents = Lists.newArrayList("R1.00", "R1.00", "R1.00");
    assertEquals(expectedEvents, mMockRateLimiter.readEventsAndClear());
}

Example 2 with MockRateLimiter

use of com.google.common.util.concurrent.MockRateLimiter in project alluxio by Alluxio.

the class FileDataManagerTest method before.

@Before
public void before() throws Exception {
    mUfs = Mockito.mock(UnderFileSystem.class);
    mBlockWorker = Mockito.mock(BlockWorker.class);
    mMockRateLimiter = new MockRateLimiter(Configuration.getBytes(PropertyKey.WORKER_FILE_PERSIST_RATE_LIMIT));
    mManager = new FileDataManager(mBlockWorker, mMockRateLimiter.getGuavaRateLimiter());
    mMockFileSystem = PowerMockito.mock(FileSystem.class);
    PowerMockito.mockStatic(FileSystem.Factory.class);
    Mockito.when(FileSystem.Factory.get()).thenReturn(mMockFileSystem);
    Mockito.when(mUfs.isDirectory(Mockito.anyString())).thenReturn(true);
    mUfsFactory = Mockito.mock(UnderFileSystemFactory.class);
    Mockito.when(mUfsFactory.supportsPath(Mockito.anyString())).thenReturn(true);
    Mockito.when(mUfsFactory.create(Mockito.anyString(), Mockito.anyObject())).thenReturn(mUfs);
    UnderFileSystemRegistry.register(mUfsFactory);
}