Examples with RateLimiter org.apache.samza.util.RateLimiter used on opensource projects

Example 1 with RateLimiter

use of org.apache.samza.util.RateLimiter in project samza by apache.

the class RemoteTableDescriptor method writeRateLimiterConfig.

// Handle rate limiter
private void writeRateLimiterConfig(Config jobConfig, Map<String, String> tableConfig) {
    if (!tagCreditsMap.isEmpty()) {
        RateLimiter defaultRateLimiter;
        try {
            @SuppressWarnings("unchecked") Class<? extends RateLimiter> clazz = (Class<? extends RateLimiter>) Class.forName(DEFAULT_RATE_LIMITER_CLASS_NAME);
            Constructor<? extends RateLimiter> ctor = clazz.getConstructor(Map.class);
            defaultRateLimiter = ctor.newInstance(tagCreditsMap);
        } catch (Exception ex) {
            throw new SamzaException("Failed to create default rate limiter", ex);
        }
        addTableConfig(RATE_LIMITER, SerdeUtils.serialize("rate limiter", defaultRateLimiter), tableConfig);
        if (defaultRateLimiter instanceof TablePart) {
            addTablePartConfig(RATE_LIMITER, (TablePart) defaultRateLimiter, jobConfig, tableConfig);
        }
    } else if (rateLimiter != null) {
        addTableConfig(RATE_LIMITER, SerdeUtils.serialize("rate limiter", rateLimiter), tableConfig);
        if (rateLimiter instanceof TablePart) {
            addTablePartConfig(RATE_LIMITER, (TablePart) rateLimiter, jobConfig, tableConfig);
        }
    }
    // Write table api read/write rate limit
    if (this.enableReadRateLimiter && tagCreditsMap.containsKey(RL_READ_TAG)) {
        addTableConfig(READ_CREDITS, String.valueOf(tagCreditsMap.get(RL_READ_TAG)), tableConfig);
    }
    if (this.enableWriteRateLimiter && tagCreditsMap.containsKey(RL_WRITE_TAG)) {
        addTableConfig(WRITE_CREDITS, String.valueOf(tagCreditsMap.get(RL_WRITE_TAG)), tableConfig);
    }
}

Example 2 with RateLimiter

use of org.apache.samza.util.RateLimiter 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 3 with RateLimiter

use of org.apache.samza.util.RateLimiter in project samza by apache.

the class TestRemoteTableDescriptor method testTablePartToConfig.

@Test
public void testTablePartToConfig() {
    int key = 0;
    TableReadFunction readFn = createMockTableReadFunction();
    when(readFn.toConfig(any(), any())).thenReturn(createConfigPair(key));
    TableWriteFunction writeFn = createMockTableWriteFunction();
    when(writeFn.toConfig(any(), any())).thenReturn(createConfigPair(key));
    RateLimiter rateLimiter = createMockRateLimiter();
    when(((TablePart) rateLimiter).toConfig(any(), any())).thenReturn(createConfigPair(key));
    CreditFunction readCredFn = createMockCreditFunction();
    when(readCredFn.toConfig(any(), any())).thenReturn(createConfigPair(key));
    CreditFunction writeCredFn = createMockCreditFunction();
    when(writeCredFn.toConfig(any(), any())).thenReturn(createConfigPair(key));
    TableRetryPolicy readRetryPolicy = createMockTableRetryPolicy();
    when(readRetryPolicy.toConfig(any(), any())).thenReturn(createConfigPair(key));
    TableRetryPolicy writeRetryPolicy = createMockTableRetryPolicy();
    when(writeRetryPolicy.toConfig(any(), any())).thenReturn(createConfigPair(key));
    Map<String, String> tableConfig = new RemoteTableDescriptor("1").withReadFunction(readFn).withWriteFunction(writeFn).withRateLimiter(rateLimiter, readCredFn, writeCredFn).withReadRetryPolicy(readRetryPolicy).withWriteRetryPolicy(writeRetryPolicy).toConfig(new MapConfig());
    verify(readFn, times(1)).toConfig(any(), any());
    verify(writeFn, times(1)).toConfig(any(), any());
    verify((TablePart) rateLimiter, times(1)).toConfig(any(), any());
    verify(readCredFn, times(1)).toConfig(any(), any());
    verify(writeCredFn, times(1)).toConfig(any(), any());
    verify(readRetryPolicy, times(1)).toConfig(any(), any());
    verify(writeRetryPolicy, times(1)).toConfig(any(), any());
    Assert.assertEquals(tableConfig.get("tables.1.io.read.func.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.write.func.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.ratelimiter.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.read.credit.func.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.write.credit.func.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.read.retry.policy.k0"), "v0");
    Assert.assertEquals(tableConfig.get("tables.1.io.write.retry.policy.k0"), "v0");
}

Example 4 with RateLimiter

use of org.apache.samza.util.RateLimiter in project samza by apache.

the class TestRemoteTableWithBatchEndToEnd method doTestStreamTableJoinRemoteTable.

private void doTestStreamTableJoinRemoteTable(String testName, boolean batchRead, boolean batchWrite) throws Exception {
    final InMemoryWriteFunction writer = new InMemoryWriteFunction(testName);
    BATCH_READS.put(testName, new AtomicInteger());
    BATCH_WRITES.put(testName, new AtomicInteger());
    WRITTEN_RECORDS.put(testName, new HashMap<>());
    int count = 16;
    int batchSize = 4;
    String profiles = Base64Serializer.serialize(generateProfiles(count));
    final RateLimiter readRateLimiter = mock(RateLimiter.class, withSettings().serializable());
    final RateLimiter writeRateLimiter = mock(RateLimiter.class, withSettings().serializable());
    final TableRateLimiter.CreditFunction creditFunction = (k, v, args) -> 1;
    final StreamApplication app = appDesc -> {
        RemoteTableDescriptor<Integer, Profile, Void> inputTableDesc = new RemoteTableDescriptor<>("profile-table-1");
        inputTableDesc.withReadFunction(InMemoryReadFunction.getInMemoryReadFunction(testName, profiles)).withRateLimiter(readRateLimiter, creditFunction, null);
        if (batchRead) {
            inputTableDesc.withBatchProvider(new CompactBatchProvider().withMaxBatchSize(batchSize).withMaxBatchDelay(Duration.ofHours(1)));
        }
        // dummy reader
        TableReadFunction readFn = new MyReadFunction();
        RemoteTableDescriptor<Integer, EnrichedPageView, EnrichedPageView> outputTableDesc = new RemoteTableDescriptor<>("enriched-page-view-table-1");
        outputTableDesc.withReadFunction(readFn).withWriteFunction(writer).withRateLimiter(writeRateLimiter, creditFunction, creditFunction);
        if (batchWrite) {
            outputTableDesc.withBatchProvider(new CompactBatchProvider().withMaxBatchSize(batchSize).withMaxBatchDelay(Duration.ofHours(1)));
        }
        Table outputTable = appDesc.getTable(outputTableDesc);
        Table<KV<Integer, Profile>> inputTable = appDesc.getTable(inputTableDesc);
        DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor("test");
        GenericInputDescriptor<PageView> isd = ksd.getInputDescriptor("PageView", new NoOpSerde<>());
        appDesc.getInputStream(isd).map(pv -> new KV<>(pv.getMemberId(), pv)).join(inputTable, new PageViewToProfileJoinFunction()).map(m -> new KV<>(m.getMemberId(), m)).sendTo(outputTable);
    };
    InMemorySystemDescriptor isd = new InMemorySystemDescriptor("test");
    InMemoryInputDescriptor<PageView> inputDescriptor = isd.getInputDescriptor("PageView", new NoOpSerde<>());
    TestRunner.of(app).addInputStream(inputDescriptor, Arrays.asList(generatePageViewsWithDistinctKeys(count))).addConfig("task.max.concurrency", String.valueOf(count)).addConfig("task.async.commit", String.valueOf(true)).run(Duration.ofSeconds(10));
    Assert.assertEquals(count, WRITTEN_RECORDS.get(testName).size());
    Assert.assertNotNull(WRITTEN_RECORDS.get(testName).get(0));
    if (batchWrite) {
        Assert.assertEquals(count / batchSize, BATCH_WRITES.get(testName).get());
    }
}

Example 5 with RateLimiter

use of org.apache.samza.util.RateLimiter in project samza by apache.

the class TestRemoteTableEndToEnd method testSendToUpdatesWithoutUpdateOptions.

// Test will fail as we use sendTo with KV<K, UpdateMessage> stream without UpdateOptions
@Test(expected = SamzaException.class)
public void testSendToUpdatesWithoutUpdateOptions() throws Exception {
    // max member id for page views is 10
    final String profiles = Base64Serializer.serialize(generateProfiles(10));
    final RateLimiter readRateLimiter = mock(RateLimiter.class, withSettings().serializable());
    final TableRateLimiter.CreditFunction creditFunction = (k, v, args) -> 1;
    final StreamApplication app = appDesc -> {
        final RemoteTableDescriptor joinTableDesc = new RemoteTableDescriptor<Integer, TestTableData.Profile, Void>("profile-table-1").withReadFunction(InMemoryProfileReadFunction.getInMemoryReadFunction(profiles)).withRateLimiter(readRateLimiter, creditFunction, null);
        final RemoteTableDescriptor outputTableDesc = new RemoteTableDescriptor<Integer, EnrichedPageView, EnrichedPageView>("enriched-page-view-table-1").withReadFunction(new NoOpTableReadFunction<>()).withReadRateLimiterDisabled().withWriteFunction(new InMemoryEnrichedPageViewWriteFunction2("testUpdateWithoutUpdateOptions", false)).withWriteRateLimit(1000);
        final Table<KV<Integer, Profile>> outputTable = appDesc.getTable(outputTableDesc);
        final Table<KV<Integer, Profile>> joinTable = appDesc.getTable(joinTableDesc);
        final DelegatingSystemDescriptor ksd = new DelegatingSystemDescriptor("test");
        final GenericInputDescriptor<PageView> isd = ksd.getInputDescriptor("PageView", new NoOpSerde<>());
        appDesc.getInputStream(isd).map(pv -> new KV<>(pv.getMemberId(), pv)).join(joinTable, new PageViewToProfileJoinFunction()).map(m -> new KV(m.getMemberId(), UpdateMessage.of(m, m))).sendTo(outputTable);
    };
    int numPageViews = 40;
    InMemorySystemDescriptor isd = new InMemorySystemDescriptor("test");
    InMemoryInputDescriptor<PageView> inputDescriptor = isd.getInputDescriptor("PageView", new NoOpSerde<>());
    TestRunner.of(app).addInputStream(inputDescriptor, TestTableData.generatePartitionedPageViews(numPageViews, 4)).run(Duration.ofSeconds(10));
}