changed metrics a bit

ExpediaGroup · Apr 17, 2024 · b2a8502 · b2a8502
1 parent 533bd23
commit b2a8502
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 18 deletions.
diff --git a/.../main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/BucketKeyGenerator.java b/.../main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/BucketKeyGenerator.java
@@ -27,7 +27,7 @@ public String generateKey(String key) {
     if (prefix != null && !prefix.isEmpty()) {
       return prefix + "_" + key;
     }
-    return key;
+    return ""+key;
   }
 
 }
diff --git a/...rc/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitMetrics.java b/...rc/main/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitMetrics.java
@@ -16,20 +16,20 @@
 package com.hotels.bdp.waggledance.extensions.client.ratelimit;
 
 public enum RateLimitMetrics {
-
-  CALLS("calls"),
+
   EXCEEDED("exceeded"),
-  ERRORS("errors");
+  ERRORS("errors"),
+  WITHIN_LIMIT("within_limit");
 
   private final static String METRIC_BASE_NAME = "com.hotels.bdp.waggledance.extensions.client.ratelimit";
   private String metricName;
-  
+
   private RateLimitMetrics(String name) {
     this.metricName = METRIC_BASE_NAME + "." + name;
   }
-  
+
   public String getMetricName() {
     return metricName;
   }
-  
+
 }
diff --git a/...com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java b/...com/hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandler.java
@@ -76,7 +76,6 @@ private Object doRateLimitCall(CloseableThriftHiveMetastoreIface client, Method
     if (shouldProceedWithCall(method)) {
       return doRealCall(client, method, args);
     } else {
-      meterRegistry.counter(RateLimitMetrics.EXCEEDED.getMetricName()).increment();
       log.info("User '{}' made too many requests.", user);
       // HTTP status would be 429, so using same for Thrift.
       throw new WaggleDanceServerException("[STATUS=429] Too many requests.");
@@ -90,7 +89,13 @@ private boolean shouldProceedWithCall(Method method) {
       log
           .info("RateLimitCall:[User:{}, method:{}, source_ip:{}, tokens_remaining:{}, metastoreName:{}]", user,
               method.getName(), HMSHandler.getThreadLocalIpAddress(), probe.getRemainingTokens(), metastoreName);
-      return probe.isConsumed();
+      boolean isConsumed = probe.isConsumed();
+      if (isConsumed) {
+        meterRegistry.counter(RateLimitMetrics.WITHIN_LIMIT.getMetricName()).increment();     
+      } else {
+        meterRegistry.counter(RateLimitMetrics.EXCEEDED.getMetricName()).increment();
+      }
+      return isConsumed;
     } catch (Exception e) {
       meterRegistry.counter(RateLimitMetrics.ERRORS.getMetricName()).increment();
       if (log.isDebugEnabled()) {
@@ -101,8 +106,6 @@ private boolean shouldProceedWithCall(Method method) {
                 e.getMessage());
       }
       return true;
-    } finally {
-      meterRegistry.counter(RateLimitMetrics.CALLS.getMetricName()).increment();
     }
   }
 

diff --git a/...t/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/BucketKeyGeneratorTest.java b/...t/java/com/hotels/bdp/waggledance/extensions/client/ratelimit/BucketKeyGeneratorTest.java
@@ -23,12 +23,28 @@
 public class BucketKeyGeneratorTest {
 
   @Test
-    public void testGenerateKey() {
-      BucketKeyGenerator bucketKeyGenerator = new BucketKeyGenerator("prefix");
-      String key = bucketKeyGenerator.generateKey("key");
-      assertThat(key, is("prefix_key"));
-    }
+  public void testGenerateKey() {
+    BucketKeyGenerator bucketKeyGenerator = new BucketKeyGenerator("prefix");
+    String key = bucketKeyGenerator.generateKey("key");
+    assertThat(key, is("prefix_key"));
+  }
+
+  @Test
+  public void testGenerateNullKey() {
+    BucketKeyGenerator bucketKeyGenerator = new BucketKeyGenerator("prefix");
+    String key = bucketKeyGenerator.generateKey(null);
+    assertThat(key, is("prefix_null"));
+  }
+
+
+  @Test
+  public void testGenerateNullKeyNullPrefix() {
+    BucketKeyGenerator bucketKeyGenerator = new BucketKeyGenerator(null);
+    String key = bucketKeyGenerator.generateKey(null);
+    assertThat(key, is("null"));
+  }
 
+
   @Test
     public void testGenerateKeyNullPrefix() {
       BucketKeyGenerator bucketKeyGenerator = new BucketKeyGenerator(null);

diff --git a/...hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java b/...hotels/bdp/waggledance/extensions/client/ratelimit/RateLimitingInvocationHandlerTest.java
@@ -89,7 +89,7 @@ public void testLimitDifferentUsers() throws Exception {
 
     verify(client, times(3)).get_table("db", "table");
     verify(client).set_ugi(USER, null);
-    assertThat(meterRegistry.counter(RateLimitMetrics.CALLS.getMetricName()).count(), is(4.0));
+    assertThat(meterRegistry.counter(RateLimitMetrics.WITHIN_LIMIT.getMetricName()).count(), is(3.0));
     assertThat(meterRegistry.counter(RateLimitMetrics.ERRORS.getMetricName()).count(), is(0.0));
     assertThat(meterRegistry.counter(RateLimitMetrics.EXCEEDED.getMetricName()).count(), is(1.0));
   }
@@ -105,7 +105,7 @@ public void testBucketExceptionStillDoCall() throws Exception {
 
     Table result = proxy.get_table("db", "table");
     assertThat(result, is(table));
-    assertThat(meterRegistry.counter(RateLimitMetrics.CALLS.getMetricName()).count(), is(1.0));
+    assertThat(meterRegistry.counter(RateLimitMetrics.WITHIN_LIMIT.getMetricName()).count(), is(0.0));
     assertThat(meterRegistry.counter(RateLimitMetrics.ERRORS.getMetricName()).count(), is(1.0));
     assertThat(meterRegistry.counter(RateLimitMetrics.EXCEEDED.getMetricName()).count(), is(0.0));