Add ratelimiter framework for grpc server interceptor #73

pavan-traceable · 2025-05-19T12:27:56Z

Description

Add grpc rate limiter utils which contains grpc server interceptor helps in rate limiting grpc methods based on configuration.

Example configuration should be defined like

  public static List<RateLimiterConfiguration> loadRateLimitConfigs() {
    return List.of(
        // Single upsert: 1 token per request
        RateLimiterConfiguration.builder()
            .enabled(true)
            .method("com.example.EntityDataService/upsert")
            .matchAttributes(Map.of("tenant_id", "default", "entity_type", "order"))
            .rateLimit(RateLimiterConfiguration.RateLimit.builder()
                .tokens(10)
                .refreshPeriodSeconds(60)
                .build())
            .attributeExtractor((requestContext, message) -> {
              Map<String, String> attrs = new HashMap<>();
              if (message instanceof Entity entity) {
                attrs.put("tenant_id", entity.getTenantId());
                attrs.put("entity_type", entity.getEntityType());
              }
              return attrs;
            })
            .tokenCostFunction((ctx, req) -> 1) // static cost
            .build(),
        // Batch upsert: cost = number of entities
        RateLimiterConfiguration.builder()
            .enabled(true)
            .method("com.example.EntityDataService/upsertEntities")
            .matchAttributes(Map.of("tenant_id", "default"))
            .rateLimit(RateLimiterConfiguration.RateLimit.builder()
                .tokens(50)
                .refreshPeriodSeconds(60)
                .build())
            .attributeExtractor((requestContext, message) -> {
              Map<String, String> attrs = new HashMap<>();
              if (message instanceof Entities entities && !entities.getEntityList().isEmpty()) {
                Entity first = entities.getEntity(0);
                attrs.put("tenant_id", first.getTenantId());
              }
              return attrs;
            })
            .tokenCostFunction((ctx, req) -> {
              if (req instanceof Entities entities) {
                return Math.max(1, entities.getEntityCount());
              }
              return 1;
            })
            .build()
    );
  }

codecov · 2025-05-19T12:30:23Z

Codecov Report

Attention: Patch coverage is 62.90323% with 23 lines in your changes missing coverage. Please review.

Project coverage is 60.93%. Comparing base (65e4f61) to head (a81ac2c).

Files with missing lines	Patch %	Lines
...grpcutils/bucket4j/Bucket4jRateLimiterFactory.java	0.00%	18 Missing ⚠️
...atelimiter/grpcutils/RateLimiterConfiguration.java	83.33%	2 Missing ⚠️
.../hypertrace/ratelimiter/grpcutils/RateLimiter.java	0.00%	1 Missing ⚠️
...elimiter/grpcutils/RateLimiterFactoryProvider.java	0.00%	1 Missing ⚠️
.../ratelimiter/grpcutils/RateLimiterInterceptor.java	96.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main      #73      +/-   ##
============================================
+ Coverage     60.77%   60.93%   +0.16%     
- Complexity      193      206      +13     
============================================
  Files            36       41       +5     
  Lines           752      814      +62     
  Branches         45       50       +5     
============================================
+ Hits            457      496      +39     
- Misses          265      287      +22     
- Partials         30       31       +1

Flag	Coverage Δ
unit	`60.93% <62.90%> (+0.16%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2025-05-19T12:30:24Z

Test Results

100 tests +10 100 ✅ +10 30s ⏱️ +5s
18 suites + 2 0 💤 ± 0
18 files + 2 0 ❌ ± 0

Results for commit a81ac2c. ± Comparison against base commit 65e4f61.

aaron-steinfeld · 2025-05-19T17:08:46Z

grpc-ratelimiter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiter.java

@@ -0,0 +1,10 @@
+package org.hypertrace.ratelimiter.grpcutils;
+
+public interface RateLimiter {


There's nothing in this interface tied to gRPC (which is great), but then we go ahead an implement such that it can only be used as a grpc interceptor. Suggest splitting this into a stand alone rate limit package (not in this repo) and then we can provide a utility here to build a gRPC interceptor given a rate limiter instance.

That way we can use it in non grpc use cases too.

aaron-steinfeld · 2025-05-19T17:10:11Z

...ter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterFactoryProvider.java

+    // Prevent instantiation
+  }
+
+  public static Bucket4jRateLimiterFactory bucket4j() {


Ideally not public. The point of the abstraction of a provider is so the caller is not concerned with the impl.

aaron-steinfeld · 2025-05-19T17:11:17Z

...miter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterConfiguration.java

+@Builder
+public class RateLimiterConfiguration {
+  boolean enabled;
+  String method;


What's this? It's not clear from the name

aaron-steinfeld · 2025-05-19T17:15:40Z

...miter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterConfiguration.java

+  Map<String, String> matchAttributes;
+
+  // Extract attributes from gRPC request
+  BiFunction<RequestContext, Object, Map<String, String>> attributeExtractor;


Following on the earlier comment, we could decouple this from grpc by dropping the attribute parts (because in the general case the caller is only passing something to the rate limiter if it wants to rate limit it) and taking the token cost from the call too.

Then, for the interceptor wrapper you'd have something like

BiPredicate<RequestContext, T> rateLimitMatcher; // Should we rate limit? Alternatively could be combined with permit calculator as a 0 response BiFunction<RequestContext, T, Object> rateLimitKeyCalculator; // What should they key be? We don't care about the type unless we have a need for human readable keys in which case swap to string BiFunction<RequestContext, T, Integer> permitCalculator; // Same as before, just renamed/typed

aaron-steinfeld · 2025-05-19T17:15:56Z

...miter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterConfiguration.java

+  @Builder
+  public static class RateLimit {
+    int tokens;
+    int refreshPeriodSeconds;


nit - use duration unless the underlying lib only has second granularity.

aaron-steinfeld · 2025-05-19T17:19:49Z

...limiter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterInterceptor.java

+import java.util.stream.Collectors;
+import org.hypertrace.core.grpcutils.context.RequestContext;
+
+public class RateLimiterInterceptor implements ServerInterceptor {


name should indicate server vs client as we likely want both eventually.

aaron-steinfeld · 2025-05-19T17:21:12Z

...limiter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterInterceptor.java

+      public void onMessage(ReqT message) {
+        RequestContext requestContext = RequestContext.fromMetadata(headers);
+        for (RateLimiterConfiguration config : rateLimitConfigs) {
+          if (!config.getMethod().equals(method)) continue;


missing braces here + below. Even for single statement if clauses our code style is to include braces.

aaron-steinfeld · 2025-05-19T17:25:09Z

...limiter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterInterceptor.java

+      public void onMessage(ReqT message) {
+        RequestContext requestContext = RequestContext.fromMetadata(headers);
+        for (RateLimiterConfiguration config : rateLimitConfigs) {
+          if (!config.getMethod().equals(method)) continue;


So this also answers my earlier question about what method is and it's the name being used for string comparison. Then, we use that to do an unchecked cast (well technically we make the rate limit config author do the unchecked cast by typing it as Object) of the message. Instead, if you accept the method descriptor itself in the config rather than just it's name you can use that directly and get type safety as it will specify your type parameters for the other functions.

aaron-steinfeld · 2025-05-19T17:28:53Z

...limiter-utils/src/main/java/org/hypertrace/ratelimiter/grpcutils/RateLimiterInterceptor.java

+  }
+
+  private boolean matches(Map<String, String> match, Map<String, String> actual) {
+    return match.entrySet().stream()


If we adjust the config as described in an earlier comment this would go away in favor of an explicit predicate.

aaron-steinfeld · 2025-05-19T17:36:53Z

.../src/main/java/org/hypertrace/ratelimiter/grpcutils/bucket4j/Bucket4jRateLimiterFactory.java

+import org.hypertrace.ratelimiter.grpcutils.RateLimiterConfiguration;
+import org.hypertrace.ratelimiter.grpcutils.RateLimiterFactory;
+
+public class Bucket4jRateLimiterFactory implements RateLimiterFactory {


Shouldn't be public, must be a singleton

Add ratelimiter framework for grpc server interceptor

a81ac2c

pavan-traceable requested a review from a team as a code owner May 19, 2025 12:27

pavan-traceable requested review from aman-bansal, aaron-steinfeld and skjindal93 May 19, 2025 13:33

aaron-steinfeld reviewed May 19, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add ratelimiter framework for grpc server interceptor #73

Add ratelimiter framework for grpc server interceptor #73

Uh oh!

pavan-traceable commented May 19, 2025 •

edited

Loading

Uh oh!

codecov bot commented May 19, 2025

Uh oh!

github-actions bot commented May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

aaron-steinfeld May 19, 2025

Uh oh!

Uh oh!

		@@ -0,0 +1,10 @@
		package org.hypertrace.ratelimiter.grpcutils;

		public interface RateLimiter {

Add ratelimiter framework for grpc server interceptor #73

Are you sure you want to change the base?

Add ratelimiter framework for grpc server interceptor #73

Uh oh!

Conversation

pavan-traceable commented May 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

codecov bot commented May 19, 2025

Codecov Report

Uh oh!

github-actions bot commented May 19, 2025

Test Results

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pavan-traceable commented May 19, 2025 •

edited

Loading