Search in sources :

Example 1 with EmbeddedTaggedRateLimiter

use of org.apache.samza.util.EmbeddedTaggedRateLimiter 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());
}
Also used : ThreadPoolExecutor(java.util.concurrent.ThreadPoolExecutor) TableDescriptor(org.apache.samza.table.descriptors.TableDescriptor) TaskModel(org.apache.samza.job.model.TaskModel) HashMap(java.util.HashMap) Mockito.spy(org.mockito.Mockito.spy) TableReadFunction(org.apache.samza.table.remote.TableReadFunction) CreditFunction(org.apache.samza.table.remote.TableRateLimiter.CreditFunction) TestUtils(org.apache.samza.testUtils.TestUtils) RemoteTable(org.apache.samza.table.remote.RemoteTable) Counter(org.apache.samza.metrics.Counter) AsyncRetriableTable(org.apache.samza.table.retry.AsyncRetriableTable) Map(java.util.Map) TableWriteFunction(org.apache.samza.table.remote.TableWriteFunction) MapConfig(org.apache.samza.config.MapConfig) Mockito.anyString(org.mockito.Mockito.anyString) JobModel(org.apache.samza.job.model.JobModel) Table(org.apache.samza.table.Table) AsyncRateLimitedTable(org.apache.samza.table.ratelimit.AsyncRateLimitedTable) TaskName(org.apache.samza.container.TaskName) EmbeddedTaggedRateLimiter(org.apache.samza.util.EmbeddedTaggedRateLimiter) ImmutableMap(com.google.common.collect.ImmutableMap) JobContext(org.apache.samza.context.JobContext) Timer(org.apache.samza.metrics.Timer) ContainerContext(org.apache.samza.context.ContainerContext) TableRetryPolicy(org.apache.samza.table.retry.TableRetryPolicy) Test(org.junit.Test) Mockito.times(org.mockito.Mockito.times) JavaTableConfig(org.apache.samza.config.JavaTableConfig) MetricsRegistry(org.apache.samza.metrics.MetricsRegistry) Mockito.when(org.mockito.Mockito.when) AsyncRemoteTable(org.apache.samza.table.remote.AsyncRemoteTable) Mockito.verify(org.mockito.Mockito.verify) Context(org.apache.samza.context.Context) ContainerModel(org.apache.samza.job.model.ContainerModel) RateLimiter(org.apache.samza.util.RateLimiter) RemoteTableProvider(org.apache.samza.table.remote.RemoteTableProvider) Config(org.apache.samza.config.Config) TablePart(org.apache.samza.table.remote.TablePart) Mockito.withSettings(org.mockito.Mockito.withSettings) TaskContextImpl(org.apache.samza.context.TaskContextImpl) AsyncReadWriteUpdateTable(org.apache.samza.table.AsyncReadWriteUpdateTable) Assert(org.junit.Assert) Mockito.any(org.mockito.Mockito.any) RemoteTableDescriptor(org.apache.samza.table.descriptors.RemoteTableDescriptor) TableRateLimiter(org.apache.samza.table.remote.TableRateLimiter) Mockito.mock(org.mockito.Mockito.mock) RemoteTable(org.apache.samza.table.remote.RemoteTable) AsyncRetriableTable(org.apache.samza.table.retry.AsyncRetriableTable) Table(org.apache.samza.table.Table) AsyncRateLimitedTable(org.apache.samza.table.ratelimit.AsyncRateLimitedTable) AsyncRemoteTable(org.apache.samza.table.remote.AsyncRemoteTable) AsyncReadWriteUpdateTable(org.apache.samza.table.AsyncReadWriteUpdateTable) AsyncRemoteTable(org.apache.samza.table.remote.AsyncRemoteTable) HashMap(java.util.HashMap) TableRetryPolicy(org.apache.samza.table.retry.TableRetryPolicy) RemoteTableDescriptor(org.apache.samza.table.descriptors.RemoteTableDescriptor) Mockito.anyString(org.mockito.Mockito.anyString) RemoteTableProvider(org.apache.samza.table.remote.RemoteTableProvider) EmbeddedTaggedRateLimiter(org.apache.samza.util.EmbeddedTaggedRateLimiter) RateLimiter(org.apache.samza.util.RateLimiter) TableRateLimiter(org.apache.samza.table.remote.TableRateLimiter) EmbeddedTaggedRateLimiter(org.apache.samza.util.EmbeddedTaggedRateLimiter) AsyncRetriableTable(org.apache.samza.table.retry.AsyncRetriableTable) TableRateLimiter(org.apache.samza.table.remote.TableRateLimiter) RemoteTable(org.apache.samza.table.remote.RemoteTable) AsyncRemoteTable(org.apache.samza.table.remote.AsyncRemoteTable) AsyncReadWriteUpdateTable(org.apache.samza.table.AsyncReadWriteUpdateTable) ThreadPoolExecutor(java.util.concurrent.ThreadPoolExecutor) AsyncRateLimitedTable(org.apache.samza.table.ratelimit.AsyncRateLimitedTable)

Aggregations

ImmutableMap (com.google.common.collect.ImmutableMap)1 HashMap (java.util.HashMap)1 Map (java.util.Map)1 ThreadPoolExecutor (java.util.concurrent.ThreadPoolExecutor)1 Config (org.apache.samza.config.Config)1 JavaTableConfig (org.apache.samza.config.JavaTableConfig)1 MapConfig (org.apache.samza.config.MapConfig)1 TaskName (org.apache.samza.container.TaskName)1 ContainerContext (org.apache.samza.context.ContainerContext)1 Context (org.apache.samza.context.Context)1 JobContext (org.apache.samza.context.JobContext)1 TaskContextImpl (org.apache.samza.context.TaskContextImpl)1 ContainerModel (org.apache.samza.job.model.ContainerModel)1 JobModel (org.apache.samza.job.model.JobModel)1 TaskModel (org.apache.samza.job.model.TaskModel)1 Counter (org.apache.samza.metrics.Counter)1 MetricsRegistry (org.apache.samza.metrics.MetricsRegistry)1 Timer (org.apache.samza.metrics.Timer)1 AsyncReadWriteUpdateTable (org.apache.samza.table.AsyncReadWriteUpdateTable)1 Table (org.apache.samza.table.Table)1