From 9648fa26a26d794495bb3acd66b9ea2812920823 Mon Sep 17 00:00:00 2001 From: paduin Date: Wed, 17 Apr 2024 00:29:40 +0200 Subject: [PATCH] refactoring/cleanup --- waggle-dance-extensions/README.md | 2 +- .../RateLimitingInvocationHandler.java | 16 ++++------ .../RateLimitingInvocationHandlerTest.java | 32 ++++++++++++++++--- 3 files changed, 35 insertions(+), 15 deletions(-) diff --git a/waggle-dance-extensions/README.md b/waggle-dance-extensions/README.md index 7311e2be..cb5bd223 100644 --- a/waggle-dance-extensions/README.md +++ b/waggle-dance-extensions/README.md @@ -16,7 +16,7 @@ To enable and configure see the following table, you can add these properties to | waggledance.extensions.ratelimit.keyPrefix | no | Optional prefix for the bucket keys. Default is (empty string) `` | | waggledance.extensions.ratelimit.storage | yes (if `enabled: true`) | The storage backend for the rate limiter, possible values `MEMORY` or `REDIS` | | waggledance.extensions.ratelimit.capacity | no | The capacity of the bucket. Default `2000` | - | waggledance.extensions.ratelimit.refillType | no | The refill type, possible values `GREEDY` or `INTERVALLY`. Default is `GREEDY` | + | waggledance.extensions.ratelimit.refillType | no | The refill type, possible values `GREEDY` or `INTERVALLY`. See [Bucket4j](https://bucket4j.com/8.9.0/toc.html#refill-types) for explanation. Default is `GREEDY` | | waggledance.extensions.ratelimit.tokensPerMinute | no | The number of tokens to add to the bucket per minute. Default `1000` | | waggledance.extensions.ratelimit.reddison.embedded.config | yes (if `storage: REDIS`) | The configuration for Redisson client, can be added in a similar way as described [here](https://github.com/redisson/redisson/tree/master/redisson-spring-boot-starter#2-add-settings-into-applicationsettings-file) | diff --git a/waggle-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java b/waggle-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java index f83b3ebe..60e8ca84 100644 --- a/waggle-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java +++ b/waggle-dance-extensions/src/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java @@ -34,9 +34,9 @@ class RateLimitingInvocationHandler implements InvocationHandler { private static Logger log = LoggerFactory.getLogger(RateLimitingInvocationHandler.class); - + static final String UNKNOWN_USER = "_UNKNOWN_USER_"; - private static final Set ignorableMethods = Sets.newHashSet("isOpen", "close", "set_ugi", "flushCache"); + private static final Set IGNORABLE_METHODS = Sets.newHashSet("isOpen", "close", "set_ugi", "flushCache"); private String metastoreName; private CloseableThriftHiveMetastoreIface client; private String user = UNKNOWN_USER; @@ -59,7 +59,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl if (method.getName().equals("set_ugi")) { user = (String) args[0]; } - if (isIgnoredMethod(method)) { + if (isIgnoredMethod(method.getName())) { return doRealCall(client, method, args); } else { return doRateLimitCall(client, method, args); @@ -68,8 +68,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl private Object doRateLimitCall(CloseableThriftHiveMetastoreIface client, Method method, Object[] args) throws IllegalAccessException, Throwable { - boolean proceedWithCall = proceedWithCall(method); - if (proceedWithCall) { + if (shouldProceedWithCall(method)) { return doRealCall(client, method, args); } else { log.info("User '{}' made too many requests.", user); @@ -78,7 +77,7 @@ private Object doRateLimitCall(CloseableThriftHiveMetastoreIface client, Method } } - private boolean proceedWithCall(Method method) { + private boolean shouldProceedWithCall(Method method) { try { Bucket bucket = bucketService.getBucket(bucketKeyGenerator.generateKey(user)); ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1); @@ -114,8 +113,7 @@ private Object doRealCall(CloseableThriftHiveMetastoreIface client, Method metho * @param method * @return true if the method should be ignored for rate limiting purposes. */ - private boolean isIgnoredMethod(Method method) { - String name = method.getName(); - return ignorableMethods.contains(name); + private boolean isIgnoredMethod(String methodName) { + return IGNORABLE_METHODS.contains(methodName); } } diff --git a/waggle-dance-extensions/src/test/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java b/waggle-dance-extensions/src/test/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java index 941fe685..9fa22ca9 100644 --- a/waggle-dance-extensions/src/test/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java +++ b/waggle-dance-extensions/src/test/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java @@ -111,17 +111,39 @@ public void testInvocationHandlerThrowsCause() throws Exception { } @Test - public void testIgnoredMethods() throws Exception { + public void testIgnoreSetUgi() throws Exception { assertTokens(2, 2); handlerProxy.set_ugi(USER, null); - handlerProxy.isOpen(); - handlerProxy.flushCache(); - handlerProxy.close(); assertTokens(2, 2); verify(client).set_ugi(USER, null); - verify(client).isOpen(); + } + + @Test + public void testIgnoreFlushCache() throws Exception { + assertTokens(2, 2); + handlerProxy.flushCache(); + assertTokens(2, 2); + verify(client).flushCache(); + } + + @Test + public void testIgnoreIsOpen() throws Exception { + assertTokens(2, 2); + + handlerProxy.isOpen(); + assertTokens(2, 2); + + verify(client).isOpen(); + } + + @Test + public void testIgnoreClose() throws Exception { + assertTokens(2, 2); + handlerProxy.close(); + assertTokens(2, 2); + verify(client).close(); }