From 8f0ef8ecc874e3461be826c19be10254481314ff Mon Sep 17 00:00:00 2001
From: Sergii Tkachenko <sergiitk@google.com>
Date: Thu, 3 Oct 2024 11:41:07 -0700
Subject: [PATCH] Dynamic bucket id builder processing initial logic.

---
 .../grpc/xds/internal/rlqs/RlqsBucketId.java  | 22 ++++++----
 .../xds/internal/rlqs/RlqsBucketSettings.java | 44 ++++++++++++++++---
 .../io/grpc/xds/internal/rlqs/RlqsEngine.java |  4 +-
 3 files changed, 54 insertions(+), 16 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketId.java b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketId.java
index 65360ce0e23..e78b36ebbd3 100644
--- a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketId.java
+++ b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketId.java
@@ -24,19 +24,25 @@
 
 @AutoValue
 public abstract class RlqsBucketId {
+  // No class loading deadlock, see
+  // https://github.com/google/error-prone/issues/2062#issuecomment-1566253739
+  public static final RlqsBucketId EMPTY = create(ImmutableMap.of());
+
   public abstract ImmutableMap<String, String> bucketId();
 
-  public static RlqsBucketId create(ImmutableMap<String, String> bucketId) {
-    return new AutoValue_RlqsBucketId(bucketId);
+  public static RlqsBucketId create(Map<String, String> bucketIdMap) {
+    if (bucketIdMap.isEmpty()) {
+      return EMPTY;
+    }
+    return new AutoValue_RlqsBucketId(ImmutableMap.copyOf(bucketIdMap));
   }
 
-  public static RlqsBucketId fromEnvoyProto(BucketId envoyProto) {
-    ImmutableMap.Builder<String, String> bucketId = ImmutableMap.builder();
-    for (Map.Entry<String, String> entry : envoyProto.getBucketMap().entrySet()) {
-      bucketId.put(entry.getKey(), entry.getValue());
-    }
-    return RlqsBucketId.create(bucketId.build());
+  public final boolean isEmpty() {
+    return bucketId().isEmpty();
+  }
 
+  public static RlqsBucketId fromEnvoyProto(BucketId envoyProto) {
+    return RlqsBucketId.create(ImmutableMap.copyOf(envoyProto.getBucketMap().entrySet()));
   }
 
   @Memoized
diff --git a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketSettings.java b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketSettings.java
index 763675a9af0..da202de011c 100644
--- a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketSettings.java
+++ b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsBucketSettings.java
@@ -17,22 +17,31 @@
 package io.grpc.xds.internal.rlqs;
 
 import com.google.auto.value.AutoValue;
-import com.google.common.base.Function;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Duration;
 import com.google.protobuf.util.Durations;
 import io.grpc.xds.internal.datatype.RateLimitStrategy;
 import io.grpc.xds.internal.matchers.HttpMatchInput;
 import io.grpc.xds.internal.rlqs.RlqsRateLimitResult.DenyResponse;
+import java.util.function.Function;
+import javax.annotation.Nullable;
 
 @AutoValue
 public abstract class RlqsBucketSettings {
   // TODO(sergiitk): [IMPL] this misses most of the parsing and implementation.
 
+  @Nullable
   public abstract ImmutableMap<String, Function<HttpMatchInput, String>> bucketIdBuilder();
 
-  public RlqsBucketId toBucketId(HttpMatchInput input) {
-    return null;
+  abstract RlqsBucketId staticBucketId();
+
+  public abstract long reportingIntervalMillis();
+
+  public final RlqsBucketId toBucketId(HttpMatchInput input) {
+    if (bucketIdBuilder() == null) {
+      return staticBucketId();
+    }
+    return processBucketBuilder(bucketIdBuilder(), input);
   }
 
   public RateLimitStrategy noAssignmentStrategy() {
@@ -47,11 +56,34 @@ public RateLimitStrategy expiredAssignmentStrategy() {
     return null;
   }
 
-  public abstract long reportingIntervalMillis();
-
   public static RlqsBucketSettings create(
       ImmutableMap<String, Function<HttpMatchInput, String>> bucketIdBuilder,
       Duration reportingInterval) {
-    return new AutoValue_RlqsBucketSettings(bucketIdBuilder, Durations.toMillis(reportingInterval));
+    // TODO(sergiitk): instead of create, use Builder pattern.
+    RlqsBucketId staticBucketId = processBucketBuilder(bucketIdBuilder, null);
+    return new AutoValue_RlqsBucketSettings(
+        staticBucketId.isEmpty() ? bucketIdBuilder : null,
+        staticBucketId,
+        Durations.toMillis(reportingInterval));
+  }
+
+  private static RlqsBucketId processBucketBuilder(
+      ImmutableMap<String, Function<HttpMatchInput, String>> bucketIdBuilder,
+      HttpMatchInput input) {
+    ImmutableMap.Builder<String, String> bucketIdMapBuilder = ImmutableMap.builder();
+    if (input == null) {
+      // TODO(sergiitk): [IMPL] calculate static map
+      return RlqsBucketId.EMPTY;
+    }
+    for (String key : bucketIdBuilder.keySet()) {
+      Function<HttpMatchInput, String> fn = bucketIdBuilder.get(key);
+      String value = null;
+      if (fn != null) {
+        value = fn.apply(input);
+      }
+      bucketIdMapBuilder.put(key, value != null ? value : "");
+    }
+    ImmutableMap<String, String> bucketIdMap = bucketIdMapBuilder.build();
+    return bucketIdMap.isEmpty() ? RlqsBucketId.EMPTY : RlqsBucketId.create(bucketIdMap);
   }
 }
diff --git a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsEngine.java b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsEngine.java
index 8b1f159ec07..d0b2021fa93 100644
--- a/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsEngine.java
+++ b/xds/src/main/java/io/grpc/xds/internal/rlqs/RlqsEngine.java
@@ -56,8 +56,8 @@ public RlqsEngine(
   public RlqsRateLimitResult rateLimit(HttpMatchInput input) {
     RlqsBucketSettings bucketSettings = bucketMatchers.match(input);
     RlqsBucketId bucketId = bucketSettings.toBucketId(input);
-    // Special case when bucket id builder not set.
-    if (bucketId == null) {
+    // Special case when bucket id builder not set, or has no values.
+    if (bucketId.isEmpty()) {
       return rateLimitWithoutReports(bucketSettings);
     }
     RlqsBucket bucket = bucketCache.getOrCreate(bucketId, bucketSettings, newBucket -> {