refactoring/cleanup

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) |
 

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<String> ignorableMethods = Sets.newHashSet("isOpen", "close", "set_ugi", "flushCache");
+  private static final Set<String> 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);
   }
 }

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();
   }