Add grpc circuit breaker utility using interceptors

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParser.java

@@ -6,8 +6,8 @@
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class CircuitBreakerConfigParser {
 
   // Percentage of failures to trigger OPEN state
   private static final String FAILURE_RATE_THRESHOLD = "failureRateThreshold";
@@ -28,81 +28,84 @@
   private static final String SLIDING_WINDOW_TYPE = "slidingWindowType";
   public static final String ENABLED = "enabled";
   public static final String DEFAULT_THRESHOLDS = "defaultThresholds";
-  private static final Set<String> NON_THRESHOLD_KEYS = Set.of(ENABLED, DEFAULT_THRESHOLDS);
+  private static final Set<String> NON_THRESHOLD_KEYS = Set.of(ENABLED);
 
   public static <T> CircuitBreakerConfiguration.CircuitBreakerConfigurationBuilder<T> parseConfig(
       Config config) {
     CircuitBreakerConfiguration.CircuitBreakerConfigurationBuilder<T> builder =
         CircuitBreakerConfiguration.builder();
     if (config.hasPath(ENABLED)) {
       builder.enabled(config.getBoolean(ENABLED));
     }
 
-    if (config.hasPath(DEFAULT_THRESHOLDS)) {
-      builder.defaultThresholds(
-          buildCircuitBreakerThresholds(config.getConfig(DEFAULT_THRESHOLDS)));
-    } else {
-      builder.defaultThresholds(buildCircuitBreakerDefaultThresholds());
-    }
-
     Map<String, CircuitBreakerThresholds> circuitBreakerThresholdsMap =
         config.root().keySet().stream()
             .filter(key -> !NON_THRESHOLD_KEYS.contains(key)) // Filter out non-threshold keys
             .collect(
                 Collectors.toMap(
                     key -> key, // Circuit breaker key
                     key -> buildCircuitBreakerThresholds(config.getConfig(key))));
+
+    if (!config.hasPath(DEFAULT_THRESHOLDS)) {
+      builder.defaultThresholds(buildCircuitBreakerDefaultThresholds());
+      circuitBreakerThresholdsMap.put(DEFAULT_THRESHOLDS, buildCircuitBreakerDefaultThresholds());
+    }
+
     builder.circuitBreakerThresholdsMap(circuitBreakerThresholdsMap);
     log.debug("Loaded circuit breaker configs: {}", builder);
     return builder;
   }
 
   private static CircuitBreakerThresholds buildCircuitBreakerThresholds(Config config) {
     CircuitBreakerThresholds.CircuitBreakerThresholdsBuilder builder =
         CircuitBreakerThresholds.builder();
 
     if (config.hasPath(FAILURE_RATE_THRESHOLD)) {
       builder.failureRateThreshold((float) config.getDouble(FAILURE_RATE_THRESHOLD));
     }
 
     if (config.hasPath(SLOW_CALL_RATE_THRESHOLD)) {
       builder.slowCallRateThreshold((float) config.getDouble(SLOW_CALL_RATE_THRESHOLD));
     }
 
     if (config.hasPath(SLOW_CALL_DURATION_THRESHOLD)) {
       builder.slowCallDurationThreshold(config.getDuration(SLOW_CALL_DURATION_THRESHOLD));
     }
 
     if (config.hasPath(SLIDING_WINDOW_TYPE)) {
       builder.slidingWindowType(getSlidingWindowType(config.getString(SLIDING_WINDOW_TYPE)));
     }
 
     if (config.hasPath(SLIDING_WINDOW_SIZE)) {
       builder.slidingWindowSize(config.getInt(SLIDING_WINDOW_SIZE));
     }
 
     if (config.hasPath(WAIT_DURATION_IN_OPEN_STATE)) {
       builder.waitDurationInOpenState(config.getDuration(WAIT_DURATION_IN_OPEN_STATE));
     }
 
     if (config.hasPath(PERMITTED_NUMBER_OF_CALLS_IN_HALF_OPEN_STATE)) {
       builder.permittedNumberOfCallsInHalfOpenState(
           config.getInt(PERMITTED_NUMBER_OF_CALLS_IN_HALF_OPEN_STATE));
     }
 
     if (config.hasPath(MINIMUM_NUMBER_OF_CALLS)) {
       builder.minimumNumberOfCalls(config.getInt(MINIMUM_NUMBER_OF_CALLS));
     }
 
+    if (config.hasPath(ENABLED)) {
+      builder.enabled(config.getBoolean(ENABLED));
+    }
+
     return builder.build();
   }
 
   public static CircuitBreakerThresholds buildCircuitBreakerDefaultThresholds() {
     return CircuitBreakerThresholds.builder().build();
   }
 
   private static CircuitBreakerThresholds.SlidingWindowType getSlidingWindowType(
       String slidingWindowType) {
     return CircuitBreakerThresholds.SlidingWindowType.valueOf(slidingWindowType);
   }
 }

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfiguration.java

@@ -1,6 +1,7 @@
 package org.hypertrace.circuitbreaker.grpcutils;
 
 import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -11,18 +12,16 @@
 @Value
 @Builder
 public class CircuitBreakerConfiguration<T> {
   Class<T> requestClass;
   BiFunction<RequestContext, T, String> keyFunction;
   @Builder.Default boolean enabled = false;
-  // Default value be "global" if not override.
-  String defaultCircuitBreakerKey = "global";
   // Standard/default thresholds
   CircuitBreakerThresholds defaultThresholds;
   // Custom overrides for specific cases (less common)
   @Builder.Default Map<String, CircuitBreakerThresholds> circuitBreakerThresholdsMap = Map.of();
 
   // New exception builder logic
   @Builder.Default
-  Function<String, RuntimeException> exceptionBuilder =
+  Function<String, StatusRuntimeException> exceptionBuilder =
       reason -> Status.RESOURCE_EXHAUSTED.withDescription(reason).asRuntimeException();
 }

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerThresholds.java

@@ -23,6 +23,7 @@
   // Calls allowed in HALF_OPEN state before deciding to
   // CLOSE or OPEN again
   @Builder.Default int permittedNumberOfCallsInHalfOpenState = 5;
+  @Builder.Default boolean enabled = true;
 
   public enum SlidingWindowType {
     COUNT_BASED,

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/resilience/ResilienceCircuitBreakerFactory.java

@@ -6,20 +6,22 @@
 import java.util.Map;
 import org.hypertrace.circuitbreaker.grpcutils.CircuitBreakerConfiguration;
 
 public class ResilienceCircuitBreakerFactory {
   public static ResilienceCircuitBreakerInterceptor getResilienceCircuitBreakerInterceptor(
       CircuitBreakerConfiguration<?> circuitBreakerConfiguration, Clock clock) {
     Map<String, CircuitBreakerConfig> resilienceCircuitBreakerConfigMap =
         ResilienceCircuitBreakerConfigConverter.getCircuitBreakerConfigs(
             circuitBreakerConfiguration.getCircuitBreakerThresholdsMap());
     CircuitBreakerRegistry resilicenceCircuitBreakerRegistry =
         new ResilienceCircuitBreakerRegistryProvider(
                 circuitBreakerConfiguration.getDefaultThresholds())
             .getCircuitBreakerRegistry();
     ResilienceCircuitBreakerProvider resilienceCircuitBreakerProvider =
         new ResilienceCircuitBreakerProvider(
-            resilicenceCircuitBreakerRegistry, resilienceCircuitBreakerConfigMap);
+            resilicenceCircuitBreakerRegistry,
+            resilienceCircuitBreakerConfigMap,
+            circuitBreakerConfiguration.getCircuitBreakerThresholdsMap());
     return new ResilienceCircuitBreakerInterceptor(
         circuitBreakerConfiguration, clock, resilienceCircuitBreakerProvider);
   }
 }

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/resilience/ResilienceCircuitBreakerInterceptor.java

@@ -13,6 +13,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.circuitbreaker.grpcutils.CircuitBreakerConfiguration;
@@ -37,7 +38,7 @@
 
   @Override
   protected boolean isCircuitBreakerEnabled() {
     return circuitBreakerConfiguration.isEnabled();
   }
 
   @Override
@@ -45,7 +46,7 @@
       MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
     return new ForwardingClientCall.SimpleForwardingClientCall<>(
         next.newCall(method, callOptions)) {
-      CircuitBreaker circuitBreaker;
+      Optional<CircuitBreaker> circuitBreaker;
       String circuitBreakerKey;
 
       @Override
@@ -64,20 +65,24 @@
             (CircuitBreakerConfiguration<ReqT>) circuitBreakerConfiguration;
         // Type check for message class compatibility
         if (config.getRequestClass() != null && !config.getRequestClass().isInstance(message)) {
           super.sendMessage(message);
           return;
         }
-        if (config.getKeyFunction() == null) {
-          log.debug("Circuit breaker will apply to all requests as keyFunction config is not set");
-          circuitBreakerKey = config.getDefaultCircuitBreakerKey();
-        } else {
+        if (config.getKeyFunction() != null) {
           circuitBreakerKey = config.getKeyFunction().apply(RequestContext.CURRENT.get(), message);
+          circuitBreaker = resilienceCircuitBreakerProvider.getCircuitBreaker(circuitBreakerKey);
+        } else {
+          log.debug("Circuit breaker will apply to all requests as keyFunction config is not set");
+          circuitBreaker = resilienceCircuitBreakerProvider.getDefaultCircuitBreaker();
+        }
+        if (circuitBreaker.isEmpty()) {
+          super.sendMessage(message);
+          return;
         }
-        circuitBreaker = resilienceCircuitBreakerProvider.getCircuitBreaker(circuitBreakerKey);
-        if (!circuitBreaker.tryAcquirePermission()) {
-          logCircuitBreakerRejection(circuitBreakerKey, circuitBreaker);
+        if (!circuitBreaker.get().tryAcquirePermission()) {
+          logCircuitBreakerRejection(circuitBreakerKey, circuitBreaker.get());
           String rejectionReason =
-              circuitBreaker.getState() == CircuitBreaker.State.HALF_OPEN
+              circuitBreaker.get().getState() == CircuitBreaker.State.HALF_OPEN
                   ? "Circuit Breaker is HALF-OPEN and rejecting excess requests"
                   : "Circuit Breaker is OPEN and blocking requests";
           throw config.getExceptionBuilder().apply(rejectionReason);
@@ -89,18 +94,21 @@
           wrapListenerWithCircuitBreaker(Listener<RespT> responseListener, Instant startTime) {
         return new ForwardingClientCallListener.SimpleForwardingClientCallListener<>(
             responseListener) {
+          @SuppressWarnings("OptionalGetWithoutIsPresent")
           @Override
           public void onClose(Status status, Metadata trailers) {
             long duration = Duration.between(startTime, clock.instant()).toNanos();
             if (status.isOk()) {
-              circuitBreaker.onSuccess(duration, TimeUnit.NANOSECONDS);
+              circuitBreaker.get().onSuccess(duration, TimeUnit.NANOSECONDS);
             } else {
               log.debug(
                   "Circuit Breaker '{}' detected failure. Status: {}, Description: {}",
-                  circuitBreaker.getName(),
+                  circuitBreaker.get().getName(),
                   status.getCode(),
                   status.getDescription());
-              circuitBreaker.onError(duration, TimeUnit.NANOSECONDS, status.asRuntimeException());
+              circuitBreaker
+                  .get()
+                  .onError(duration, TimeUnit.NANOSECONDS, status.asRuntimeException());
             }
             super.onClose(status, trailers);
           }

grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/resilience/ResilienceCircuitBreakerProvider.java

@@ -7,53 +7,67 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import lombok.extern.slf4j.Slf4j;
+import org.hypertrace.circuitbreaker.grpcutils.CircuitBreakerConfigParser;
+import org.hypertrace.circuitbreaker.grpcutils.CircuitBreakerThresholds;
 
 /** Utility class to provide Resilience4j CircuitBreaker */
 @Slf4j
 class ResilienceCircuitBreakerProvider {
 
   private final CircuitBreakerRegistry circuitBreakerRegistry;
   private final Map<String, CircuitBreakerConfig> circuitBreakerConfigMap;
+  private final Map<String, CircuitBreakerThresholds> circuitBreakerThresholdsMap;
   private final Map<String, CircuitBreaker> circuitBreakerCache = new ConcurrentHashMap<>();
 
   public ResilienceCircuitBreakerProvider(
       CircuitBreakerRegistry circuitBreakerRegistry,
-      Map<String, CircuitBreakerConfig> circuitBreakerConfigMap) {
+      Map<String, CircuitBreakerConfig> circuitBreakerConfigMap,
+      Map<String, CircuitBreakerThresholds> circuitBreakerThresholdsMap) {
     this.circuitBreakerRegistry = circuitBreakerRegistry;
     this.circuitBreakerConfigMap = circuitBreakerConfigMap;
+    this.circuitBreakerThresholdsMap = circuitBreakerThresholdsMap;
   }
 
-  public CircuitBreaker getCircuitBreaker(String circuitBreakerKey) {
-    return circuitBreakerCache.computeIfAbsent(
-        circuitBreakerKey,
-        key -> {
-          CircuitBreaker circuitBreaker = getCircuitBreakerFromConfigMap(circuitBreakerKey);
-          circuitBreaker
-              .getEventPublisher()
-              .onStateTransition(
-                  event ->
-                      log.info(
-                          "State transition: {} for circuit breaker {}",
-                          event.getStateTransition(),
-                          event.getCircuitBreakerName()))
-              .onCallNotPermitted(
-                  event ->
-                      log.debug(
-                          "Call not permitted: Circuit is OPEN for circuit breaker {}",
-                          event.getCircuitBreakerName()))
-              .onEvent(
-                  event ->
-                      log.debug(
-                          "Circuit breaker event type {} for circuit breaker name {}",
-                          event.getEventType(),
-                          event.getCircuitBreakerName()));
-          return circuitBreaker;
-        });
+  public Optional<CircuitBreaker> getCircuitBreaker(String circuitBreakerKey) {
+    if (circuitBreakerThresholdsMap.containsKey(circuitBreakerKey)
+        && !circuitBreakerThresholdsMap.get(circuitBreakerKey).isEnabled()) {
+      return Optional.empty();
+    }
+    return Optional.of(
+        circuitBreakerCache.computeIfAbsent(
+            circuitBreakerKey,
+            key -> {
+              CircuitBreaker circuitBreaker = getCircuitBreakerFromConfigMap(circuitBreakerKey);
+              circuitBreaker
+                  .getEventPublisher()
+                  .onStateTransition(
+                      event ->
+                          log.info(
+                              "State transition: {} for circuit breaker {}",
+                              event.getStateTransition(),
+                              event.getCircuitBreakerName()))
+                  .onCallNotPermitted(
+                      event ->
+                          log.debug(
+                              "Call not permitted: Circuit is OPEN for circuit breaker {}",
+                              event.getCircuitBreakerName()))
+                  .onEvent(
+                      event ->
+                          log.debug(
+                              "Circuit breaker event type {} for circuit breaker name {}",
+                              event.getEventType(),
+                              event.getCircuitBreakerName()));
+              return circuitBreaker;
+            }));
+  }
+
+  public Optional<CircuitBreaker> getDefaultCircuitBreaker() {
+    return getCircuitBreaker(CircuitBreakerConfigParser.DEFAULT_THRESHOLDS);
   }
 
   private CircuitBreaker getCircuitBreakerFromConfigMap(String circuitBreakerKey) {
     return Optional.ofNullable(circuitBreakerConfigMap.get(circuitBreakerKey))
         .map(config -> circuitBreakerRegistry.circuitBreaker(circuitBreakerKey, config))
         .orElseGet(() -> circuitBreakerRegistry.circuitBreaker(circuitBreakerKey));
   }
 }

grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/resilience/ResilienceCircuitBreakerInterceptorTest.java

@@ -2,7 +2,6 @@
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.*;
 
 import io.github.resilience4j.circuitbreaker.CircuitBreaker;
@@ -11,6 +10,7 @@
 import java.time.Clock;
 import java.time.Instant;
 import java.time.ZoneOffset;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import org.hypertrace.circuitbreaker.grpcutils.CircuitBreakerConfiguration;
 import org.junit.jupiter.api.BeforeEach;
@@ -40,14 +40,11 @@ void setUp() {
 
     fixedClock = Clock.fixed(Instant.now(), ZoneOffset.UTC);
     when(mockChannel.newCall(any(), any())).thenReturn(mockClientCall);
-    when(mockCircuitBreakerProvider.getCircuitBreaker(anyString())).thenReturn(mockCircuitBreaker);
-    when(mockCircuitBreakerConfig.getDefaultCircuitBreakerKey()).thenReturn("global");
   }
 
   @Test
   void testSendMessage_CallsSuperSendMessage_Success() {
     doNothing().when(mockClientCall).sendMessage(any());
-    when(mockCircuitBreaker.tryAcquirePermission()).thenReturn(true);
 
     ResilienceCircuitBreakerInterceptor interceptor =
         new ResilienceCircuitBreakerInterceptor(
@@ -67,6 +64,8 @@ void testSendMessage_CallsSuperSendMessage_Success() {
   void testSendMessage_CircuitBreakerRejectsRequest() {
     when(mockCircuitBreaker.tryAcquirePermission()).thenReturn(false);
     when(mockCircuitBreaker.getState()).thenReturn(CircuitBreaker.State.OPEN);
+    when(mockCircuitBreakerProvider.getDefaultCircuitBreaker())
+        .thenReturn(Optional.of(mockCircuitBreaker));
     when(mockCircuitBreakerConfig.getExceptionBuilder())
         .thenReturn(
             reason ->
@@ -94,6 +93,8 @@ void testSendMessage_CircuitBreakerRejectsRequest() {
   void testSendMessage_CircuitBreakerInHalfOpenState() {
     when(mockCircuitBreaker.tryAcquirePermission()).thenReturn(false);
     when(mockCircuitBreaker.getState()).thenReturn(CircuitBreaker.State.HALF_OPEN);
+    when(mockCircuitBreakerProvider.getDefaultCircuitBreaker())
+        .thenReturn(Optional.of(mockCircuitBreaker));
     when(mockCircuitBreakerConfig.getExceptionBuilder())
         .thenReturn(
             reason ->
@@ -120,6 +121,8 @@ void testSendMessage_CircuitBreakerInHalfOpenState() {
   @Test
   void testWrapListenerWithCircuitBreaker_Success() {
     when(mockCircuitBreaker.tryAcquirePermission()).thenReturn(true);
+    when(mockCircuitBreakerProvider.getDefaultCircuitBreaker())
+        .thenReturn(Optional.of(mockCircuitBreaker));
     ResilienceCircuitBreakerInterceptor interceptor =
         new ResilienceCircuitBreakerInterceptor(
             mockCircuitBreakerConfig, fixedClock, mockCircuitBreakerProvider);
@@ -144,6 +147,8 @@ void testWrapListenerWithCircuitBreaker_Success() {
   @Test
   void testWrapListenerWithCircuitBreaker_Failure() {
     when(mockCircuitBreaker.tryAcquirePermission()).thenReturn(true);
+    when(mockCircuitBreakerProvider.getDefaultCircuitBreaker())
+        .thenReturn(Optional.of(mockCircuitBreaker));
     ResilienceCircuitBreakerInterceptor interceptor =
         new ResilienceCircuitBreakerInterceptor(
             mockCircuitBreakerConfig, fixedClock, mockCircuitBreakerProvider);