/** * Throttle a request with a key argument if necessary. * @param key key used for the table request */ public void throttle(K key) { throttle(getCredits(key, null)); }
/** * Throttle a request with a collection of table records as the argument if necessary. * @param records collection of records used for the table request */ public void throttleRecords(Collection<Entry<K, V>> records) { throttle(getEntryCredits(records)); } }
/** * Execute an async request given a table key * @param rateLimiter helper for rate limiting * @param key key of the table record * @param method method to be executed * @param timer latency metric to be updated * @param <T> return type * @return CompletableFuture of the operation */ protected <T> CompletableFuture<T> execute(TableRateLimiter<K, V> rateLimiter, K key, Function<K, CompletableFuture<T>> method, Timer timer) { final long startNs = System.nanoTime(); CompletableFuture<T> ioFuture = rateLimiter.isRateLimited() ? CompletableFuture .runAsync(() -> rateLimiter.throttle(key), tableExecutor) .thenCompose((r) -> method.apply(key)) : method.apply(key); if (callbackExecutor != null) { ioFuture.thenApplyAsync(r -> { timer.update(System.nanoTime() - startNs); return r; }, callbackExecutor); } else { ioFuture.thenApply(r -> { timer.update(System.nanoTime() - startNs); return r; }); } return ioFuture; }
@Override public CompletableFuture<Map<K, V>> getAllAsync(List<K> keys) { return isReadRateLimited() ? CompletableFuture .runAsync(() -> readRateLimiter.throttle(keys), rateLimitingExecutor) .thenCompose((r) -> table.getAllAsync(keys)) : table.getAllAsync(keys); }
final long startNs = System.nanoTime(); CompletableFuture<Void> ioFuture; if (rateLimiter.isRateLimited()) { ioFuture = CompletableFuture .runAsync(() -> rateLimiter.throttleRecords(records), tableExecutor) .thenCompose((r) -> method.apply(records)); } else {
TableRateLimiter readRateLimiter = new TableRateLimiter(tableSpec.getId(), rateLimiter, readCreditFn, RemoteTableDescriptor.RL_READ_TAG); boolean isRateLimited = readRateLimiter.isRateLimited(); if (!readOnly) { writeCreditFn = deserializeObject(WRITE_CREDIT_FN); writeRateLimiter = new TableRateLimiter(tableSpec.getId(), rateLimiter, writeCreditFn, RemoteTableDescriptor.RL_WRITE_TAG); isRateLimited |= writeRateLimiter.isRateLimited(); writeRetryPolicy = deserializeObject(WRITE_RETRY_POLICY); if (writeRetryPolicy != null) {
verify(writeRateLimiter, times(1)).throttle(anyString(), anyString()); verify(writeRateLimiter, times(0)).throttleRecords(anyList()); verify(writeRateLimiter, times(0)).throttle(anyString()); verify(writeRateLimiter, times(0)).throttle(anyList()); verify(writeRateLimiter, times(1)).throttle(anyString(), anyString()); verify(writeRateLimiter, times(1)).throttleRecords(anyList()); verify(writeRateLimiter, times(0)).throttle(anyString()); verify(writeRateLimiter, times(0)).throttle(anyList()); verify(writeRateLimiter, times(1)).throttle(anyString(), anyString()); verify(writeRateLimiter, times(1)).throttleRecords(anyList()); verify(writeRateLimiter, times(1)).throttle(anyString()); verify(writeRateLimiter, times(0)).throttle(anyList()); verify(writeRateLimiter, times(1)).throttle(anyString(), anyString()); verify(writeRateLimiter, times(1)).throttleRecords(anyList()); verify(writeRateLimiter, times(1)).throttle(anyString()); verify(writeRateLimiter, times(1)).throttle(anyList());
/** * {@inheritDoc} */ @Override public void init(Context context) { readMetrics = new DefaultTableReadMetrics(context, this, tableId); TableMetricsUtil tableMetricsUtil = new TableMetricsUtil(context, this, tableId); readRateLimiter.setTimerMetric(tableMetricsUtil.newTimer("get-throttle-ns")); }
@Test public void testThrottle() { TableRateLimiter<String, String> rateLimitHelper = getThrottler(); Timer timer = mock(Timer.class); rateLimitHelper.setTimerMetric(timer); rateLimitHelper.throttle("foo"); verify(rateLimitHelper.rateLimiter, times(1)).acquire(anyMap()); verify(timer, times(1)).update(anyLong()); }
public TableRateLimiter<String, String> getThrottler(String tag) { TableRateLimiter.CreditFunction<String, String> credFn = (TableRateLimiter.CreditFunction<String, String>) (key, value) -> { int credits = key == null ? 0 : 3; credits += value == null ? 0 : 3; return credits; }; RateLimiter rateLimiter = mock(RateLimiter.class); doReturn(Collections.singleton(tag)).when(rateLimiter).getSupportedTags(); TableRateLimiter<String, String> rateLimitHelper = new TableRateLimiter<>("foo", rateLimiter, credFn, tag); Timer timer = mock(Timer.class); rateLimitHelper.setTimerMetric(timer); return rateLimitHelper; }
@Test public void testCreditKeys() { TableRateLimiter<String, String> rateLimitHelper = getThrottler(); Assert.assertEquals(9, rateLimitHelper.getCredits(Arrays.asList("abc", "efg", "hij"))); }
@Override public CompletableFuture<Void> putAllAsync(List<Entry<K, V>> entries) { return isWriteRateLimited() ? CompletableFuture .runAsync(() -> writeRateLimiter.throttleRecords(entries), rateLimitingExecutor) .thenCompose((r) -> table.putAllAsync(entries)) : table.putAllAsync(entries); }
? new TableRateLimiter(tableId, rateLimiter, readCreditFn, RemoteTableDescriptor.RL_READ_TAG) : null; TableRetryPolicy readRetryPolicy = deserializeObject(tableConfig, RemoteTableDescriptor.READ_RETRY_POLICY); TableRateLimiter.CreditFunction<?, ?> writeCreditFn = deserializeObject(tableConfig, RemoteTableDescriptor.WRITE_CREDIT_FN); writeRateLimiter = rateLimiter != null && rateLimiter.getSupportedTags().contains(RemoteTableDescriptor.RL_WRITE_TAG) ? new TableRateLimiter(tableId, rateLimiter, writeCreditFn, RemoteTableDescriptor.RL_WRITE_TAG) : null; writeRetryPolicy = deserializeObject(tableConfig, RemoteTableDescriptor.WRITE_RETRY_POLICY);
@Test public void testCreditEntries() { TableRateLimiter<String, String> rateLimitHelper = getThrottler(); Assert.assertEquals(12, rateLimitHelper.getEntryCredits( Arrays.asList(new Entry<>("abc", "efg"), new Entry<>("hij", "lmn")))); }
/** * Execute an async request given a collection of table keys * @param rateLimiter helper for rate limiting * @param keys collection of keys * @param method method to be executed * @param timer latency metric to be updated * @param <T> return type * @return CompletableFuture of the operation */ protected <T> CompletableFuture<T> execute(TableRateLimiter<K, V> rateLimiter, Collection<K> keys, Function<Collection<K>, CompletableFuture<T>> method, Timer timer) { final long startNs = System.nanoTime(); CompletableFuture<T> ioFuture = rateLimiter.isRateLimited() ? CompletableFuture .runAsync(() -> rateLimiter.throttle(keys), tableExecutor) .thenCompose((r) -> method.apply(keys)) : method.apply(keys); if (callbackExecutor != null) { ioFuture.thenApplyAsync(r -> { timer.update(System.nanoTime() - startNs); return r; }, callbackExecutor); } else { ioFuture.thenApply(r -> { timer.update(System.nanoTime() - startNs); return r; }); } return ioFuture; }
@Override public CompletableFuture<Void> putAsync(K key, V value) { return isWriteRateLimited() ? CompletableFuture .runAsync(() -> writeRateLimiter.throttle(key, value), rateLimitingExecutor) .thenCompose((r) -> table.putAsync(key, value)) : table.putAsync(key, value); }
final long startNs = System.nanoTime(); CompletableFuture<Void> ioFuture; if (rateLimiter.isRateLimited()) { ioFuture = CompletableFuture .runAsync(() -> rateLimiter.throttleRecords(records), tableExecutor) .thenCompose((r) -> method.apply(records)); } else {
TableRateLimiter readRateLimiter = new TableRateLimiter(tableSpec.getId(), rateLimiter, readCreditFn, RemoteTableDescriptor.RL_READ_TAG); boolean isRateLimited = readRateLimiter.isRateLimited(); if (!readOnly) { writeCreditFn = deserializeObject(WRITE_CREDIT_FN); writeRateLimiter = new TableRateLimiter(tableSpec.getId(), rateLimiter, writeCreditFn, RemoteTableDescriptor.RL_WRITE_TAG); isRateLimited |= writeRateLimiter.isRateLimited(); writeRetryPolicy = deserializeObject(WRITE_RETRY_POLICY); if (writeRetryPolicy != null) {
Assert.assertEquals(1, argCaptor.getValue().size()); Assert.assertEquals("foo1", ((Entry) argCaptor.getValue().get(0)).getKey()); verify(table.writeRateLimiter, times(1)).throttle(anyCollection()); } else { Assert.assertEquals(entries, argCaptor.getValue()); verify(table.writeRateLimiter, times(1)).throttleRecords(anyCollection());
/** * {@inheritDoc} */ @Override public void init(Context context) { readMetrics = new DefaultTableReadMetrics(context, this, tableId); TableMetricsUtil tableMetricsUtil = new TableMetricsUtil(context, this, tableId); readRateLimiter.setTimerMetric(tableMetricsUtil.newTimer("get-throttle-ns")); }