Examples with AsyncRateLimitedTable org.apache.samza.table.ratelimit.AsyncRateLimitedTable used on opensource projects

Example 1 with AsyncRateLimitedTable

use of org.apache.samza.table.ratelimit.AsyncRateLimitedTable in project samza by apache.

the class TestRemoteTableDescriptor method doTestDeserializeReadFunctionAndLimiter.

private void doTestDeserializeReadFunctionAndLimiter(boolean rateOnly, boolean rlGets, boolean rlPuts) {
    int numRateLimitOps = (rlGets ? 1 : 0) + (rlPuts ? 1 : 0);
    RemoteTableDescriptor<String, String, String> desc = new RemoteTableDescriptor("1").withReadFunction(createMockTableReadFunction()).withReadRetryPolicy(new TableRetryPolicy().withRetryPredicate((ex) -> false)).withWriteFunction(createMockTableWriteFunction()).withAsyncCallbackExecutorPoolSize(10);
    if (rateOnly) {
        if (rlGets) {
            desc.withReadRateLimit(1000);
        } else {
            desc.withReadRateLimiterDisabled();
        }
        if (rlPuts) {
            desc.withWriteRateLimit(2000);
        } else {
            desc.withWriteRateLimiterDisabled();
        }
    } else {
        if (numRateLimitOps > 0) {
            Map<String, Integer> tagCredits = new HashMap<>();
            if (rlGets) {
                tagCredits.put(RemoteTableDescriptor.RL_READ_TAG, 1000);
            } else {
                desc.withReadRateLimiterDisabled();
            }
            if (rlPuts) {
                tagCredits.put(RemoteTableDescriptor.RL_WRITE_TAG, 2000);
            } else {
                desc.withWriteRateLimiterDisabled();
            }
            // Spy the rate limiter to verify call count
            RateLimiter rateLimiter = spy(new EmbeddedTaggedRateLimiter(tagCredits));
            desc.withRateLimiter(rateLimiter, new CountingCreditFunction(), new CountingCreditFunction());
        } else {
            desc.withRateLimiterDisabled();
        }
    }
    RemoteTableProvider provider = new RemoteTableProvider(desc.getTableId());
    provider.init(createMockContext(desc));
    Table table = provider.getTable();
    Assert.assertTrue(table instanceof RemoteTable);
    RemoteTable rwTable = (RemoteTable) table;
    AsyncReadWriteUpdateTable delegate = TestUtils.getFieldValue(rwTable, "asyncTable");
    Assert.assertTrue(delegate instanceof AsyncRetriableTable);
    if (rlGets || rlPuts) {
        delegate = TestUtils.getFieldValue(delegate, "table");
        Assert.assertTrue(delegate instanceof AsyncRateLimitedTable);
    }
    delegate = TestUtils.getFieldValue(delegate, "table");
    Assert.assertTrue(delegate instanceof AsyncRemoteTable);
    if (numRateLimitOps > 0) {
        TableRateLimiter readRateLimiter = TestUtils.getFieldValue(rwTable, "readRateLimiter");
        TableRateLimiter writeRateLimiter = TestUtils.getFieldValue(rwTable, "writeRateLimiter");
        Assert.assertTrue(!rlGets || readRateLimiter != null);
        Assert.assertTrue(!rlPuts || writeRateLimiter != null);
    }
    ThreadPoolExecutor callbackExecutor = TestUtils.getFieldValue(rwTable, "callbackExecutor");
    Assert.assertEquals(10, callbackExecutor.getCorePoolSize());
}

Example 2 with AsyncRateLimitedTable

use of org.apache.samza.table.ratelimit.AsyncRateLimitedTable in project samza by apache.

the class TestRemoteTable method testInit.

@Test
public void testInit() {
    String tableId = "testInit";
    TableReadFunction<String, String> readFn = mock(TableReadFunction.class);
    TableWriteFunction<String, String, String> writeFn = mock(TableWriteFunction.class);
    RemoteTable<String, String, String> table = getTable(tableId, readFn, writeFn, true);
    // AsyncRetriableTable
    AsyncReadWriteUpdateTable innerTable = TestUtils.getFieldValue(table, "asyncTable");
    Assert.assertTrue(innerTable instanceof AsyncRetriableTable);
    Assert.assertNotNull(TestUtils.getFieldValue(innerTable, "readRetryMetrics"));
    Assert.assertNotNull(TestUtils.getFieldValue(innerTable, "writeRetryMetrics"));
    // AsyncRateLimitedTable
    innerTable = TestUtils.getFieldValue(innerTable, "table");
    Assert.assertTrue(innerTable instanceof AsyncRateLimitedTable);
    // AsyncRemoteTable
    innerTable = TestUtils.getFieldValue(innerTable, "table");
    Assert.assertTrue(innerTable instanceof AsyncRemoteTable);
    // Verify table functions are initialized
    verify(readFn, times(1)).init(any(), any());
    verify(writeFn, times(1)).init(any(), any());
}