diff --git a/.gitignore b/.gitignore index 1f93dfc..f47b1ba 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ build .idea + +.superpowers/ +docs/superpowers/ diff --git a/grpc-circuitbreaker-utils/build.gradle.kts b/grpc-circuitbreaker-utils/build.gradle.kts index 58f6201..cbae719 100644 --- a/grpc-circuitbreaker-utils/build.gradle.kts +++ b/grpc-circuitbreaker-utils/build.gradle.kts @@ -7,11 +7,12 @@ plugins { dependencies { api(commonLibs.grpc.api) + api(commonLibs.typesafe.config) api(projects.grpcContextUtils) + api(projects.grpcClientUtils) implementation(commonLibs.slf4j2.api) implementation(localLibs.resilience4j.circuitbreaker) - implementation(commonLibs.typesafe.config) implementation(commonLibs.guava) implementation(commonLibs.jakarta.inject.api) diff --git a/grpc-circuitbreaker-utils/gradle.lockfile b/grpc-circuitbreaker-utils/gradle.lockfile index 7c54541..f768438 100644 --- a/grpc-circuitbreaker-utils/gradle.lockfile +++ b/grpc-circuitbreaker-utils/gradle.lockfile @@ -20,13 +20,16 @@ com.google.guava:guava:33.3.1-android=runtimeClasspath,testRuntimeClasspath com.google.guava:listenablefuture:9999.0-empty-to-avoid-conflict-with-guava=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.google.j2objc:j2objc-annotations:2.8=compileClasspath,testCompileClasspath com.google.j2objc:j2objc-annotations:3.0.0=runtimeClasspath,testRuntimeClasspath +com.google.protobuf:protobuf-java-util:3.25.8=testRuntimeClasspath +com.google.protobuf:protobuf-java:3.25.8=testRuntimeClasspath com.typesafe:config:1.4.2=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.github.resilience4j:resilience4j-circuitbreaker:1.7.1=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.github.resilience4j:resilience4j-core:1.7.1=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-api:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-bom:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath -io.grpc:grpc-context:1.75.0=runtimeClasspath,testRuntimeClasspath +io.grpc:grpc-context:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-core:1.75.0=runtimeClasspath,testRuntimeClasspath +io.grpc:grpc-inprocess:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.netty:netty-bom:4.1.132.Final=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.perfmark:perfmark-api:0.27.0=runtimeClasspath,testRuntimeClasspath io.vavr:vavr-match:0.10.2=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath diff --git a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistry.java b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistry.java new file mode 100644 index 0000000..2fea39e --- /dev/null +++ b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistry.java @@ -0,0 +1,31 @@ +package org.hypertrace.circuitbreaker.grpcutils; + +import com.typesafe.config.Config; +import java.util.Map; +import lombok.AccessLevel; +import lombok.NoArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.hypertrace.core.grpcutils.client.GrpcChannelRegistry; +import org.hypertrace.core.grpcutils.client.InterceptingChannelRegistry; + +@Slf4j +@NoArgsConstructor(access = AccessLevel.PRIVATE) +public final class CircuitBreakerChannelRegistry { + + static final String CIRCUIT_BREAKER_CONFIG = CircuitBreakerConfigParser.CIRCUIT_BREAKER_CONFIG; + + public static GrpcChannelRegistry wrap(final GrpcChannelRegistry registry, final Config config) { + return wrap(registry, CircuitBreakerConfigParser.parseServiceThresholds(config)); + } + + public static GrpcChannelRegistry wrap( + final GrpcChannelRegistry registry, + final Map thresholdsByService) { + if (registry == null || thresholdsByService == null || thresholdsByService.isEmpty()) { + log.debug("No enabled service circuit breakers; leaving the channel registry unchanged"); + return registry; + } + return InterceptingChannelRegistry.wrap( + registry, CircuitBreakerInterceptorFactory.forService(thresholdsByService)); + } +} diff --git a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParser.java b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParser.java index dc4867f..e7f9f25 100644 --- a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParser.java +++ b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParser.java @@ -1,6 +1,12 @@ package org.hypertrace.circuitbreaker.grpcutils; import com.typesafe.config.Config; +import com.typesafe.config.ConfigException; +import com.typesafe.config.ConfigUtil; +import io.grpc.Status; +import java.util.EnumSet; +import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -26,6 +32,8 @@ public class CircuitBreakerConfigParser { private static final String PERMITTED_NUMBER_OF_CALLS_IN_HALF_OPEN_STATE = "permittedNumberOfCallsInHalfOpenState"; private static final String SLIDING_WINDOW_TYPE = "slidingWindowType"; + private static final String FAILURE_STATUS_CODES = "failureStatusCodes"; + static final String CIRCUIT_BREAKER_CONFIG = "circuit.breaker.config"; public static final String ENABLED = "enabled"; public static final String DEFAULT_THRESHOLDS = "defaultThresholds"; private static final Set NON_THRESHOLD_KEYS = Set.of(ENABLED, DEFAULT_THRESHOLDS); @@ -56,9 +64,54 @@ public static CircuitBreakerConfiguration.CircuitBreakerConfigurationBuilder return builder; } + /** + * Service-keyed thresholds for {@code circuit.breaker.config}. An absent or disabled block is an + * empty map. Defaults come from {@link CircuitBreakerThresholds}; config only overlays fields + * that are present. + */ + public static Map parseServiceThresholds( + final Config rootConfig) { + if (rootConfig == null || !rootConfig.hasPath(CIRCUIT_BREAKER_CONFIG)) { + log.debug("circuit.breaker.config is absent; no service circuit breakers installed"); + return Map.of(); + } + final Config config = rootConfig.getConfig(CIRCUIT_BREAKER_CONFIG); + if (!config.hasPath(ENABLED) || !config.getBoolean(ENABLED)) { + log.debug("circuit.breaker.config is disabled; no service circuit breakers installed"); + return Map.of(); + } + final CircuitBreakerThresholds defaults = + config.hasPath(DEFAULT_THRESHOLDS) + ? overlay(buildCircuitBreakerDefaultThresholds(), config.getConfig(DEFAULT_THRESHOLDS)) + : buildCircuitBreakerDefaultThresholds(); + final Map thresholdsByService = new LinkedHashMap<>(); + for (final String serviceName : config.root().keySet()) { + if (ENABLED.equals(serviceName) || DEFAULT_THRESHOLDS.equals(serviceName)) { + continue; + } + // getConfig(key) treats dots as path separators, so service names are read with joinPath. + final CircuitBreakerThresholds thresholds = + overlay(defaults, config.getConfig(ConfigUtil.joinPath(serviceName))); + if (!thresholds.isEnabled()) { + log.debug("Service circuit breaker disabled for {}", serviceName); + continue; + } + thresholdsByService.put(serviceName, thresholds); + } + return Map.copyOf(thresholdsByService); + } + + public static CircuitBreakerThresholds overlay( + final CircuitBreakerThresholds base, final Config overrides) { + return applyOverrides(base.toBuilder(), overrides).build(); + } + private static CircuitBreakerThresholds buildCircuitBreakerThresholds(Config config) { - CircuitBreakerThresholds.CircuitBreakerThresholdsBuilder builder = - CircuitBreakerThresholds.builder(); + return overlay(buildCircuitBreakerDefaultThresholds(), config); + } + + private static CircuitBreakerThresholds.CircuitBreakerThresholdsBuilder applyOverrides( + final CircuitBreakerThresholds.CircuitBreakerThresholdsBuilder builder, final Config config) { if (config.hasPath(FAILURE_RATE_THRESHOLD)) { builder.failureRateThreshold((float) config.getDouble(FAILURE_RATE_THRESHOLD)); @@ -97,7 +150,25 @@ private static CircuitBreakerThresholds buildCircuitBreakerThresholds(Config con builder.enabled(config.getBoolean(ENABLED)); } - return builder.build(); + if (config.hasPath(FAILURE_STATUS_CODES)) { + builder.failureStatusCodes( + parseFailureStatusCodes(config.getStringList(FAILURE_STATUS_CODES))); + } + + return builder; + } + + // An unknown name must fail startup. Ignoring it would protect the wrong calls. + private static Set parseFailureStatusCodes(List names) { + Set codes = EnumSet.noneOf(Status.Code.class); + for (String name : names) { + try { + codes.add(Status.Code.valueOf(name)); + } catch (IllegalArgumentException exception) { + throw new ConfigException.BadValue(FAILURE_STATUS_CODES, name); + } + } + return Set.copyOf(codes); } public static CircuitBreakerThresholds buildCircuitBreakerDefaultThresholds() { diff --git a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerInterceptorFactory.java b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerInterceptorFactory.java index a353357..93df720 100644 --- a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerInterceptorFactory.java +++ b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerInterceptorFactory.java @@ -1,7 +1,10 @@ package org.hypertrace.circuitbreaker.grpcutils; +import com.typesafe.config.Config; +import io.grpc.ClientInterceptor; import jakarta.inject.Inject; import java.time.Clock; +import java.util.Map; import lombok.AllArgsConstructor; import org.hypertrace.circuitbreaker.grpcutils.resilience.ResilienceCircuitBreakerFactory; @@ -9,6 +12,15 @@ public class CircuitBreakerInterceptorFactory { private final Clock clock; + public static ClientInterceptor forService(final Config rootConfig) { + return forService(CircuitBreakerConfigParser.parseServiceThresholds(rootConfig)); + } + + public static ClientInterceptor forService( + final Map thresholdsByService) { + return ServiceCircuitBreakerInterceptor.create(thresholdsByService); + } + public CircuitBreakerInterceptor buildInterceptor(CircuitBreakerConfiguration configuration) { return ResilienceCircuitBreakerFactory.getCircuitBreakerInterceptor(configuration, clock); } diff --git a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerThresholds.java b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerThresholds.java index 3b7f89e..1bbeafd 100644 --- a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerThresholds.java +++ b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerThresholds.java @@ -1,11 +1,13 @@ package org.hypertrace.circuitbreaker.grpcutils; +import io.grpc.Status; import java.time.Duration; +import java.util.Set; import lombok.Builder; import lombok.Value; @Value -@Builder +@Builder(toBuilder = true) public class CircuitBreakerThresholds { // Percentage of failures to trigger OPEN state @Builder.Default float failureRateThreshold = 50f; @@ -25,6 +27,11 @@ public class CircuitBreakerThresholds { @Builder.Default int permittedNumberOfCallsInHalfOpenState = 5; @Builder.Default boolean enabled = true; + @Builder.Default + Set failureStatusCodes = + Set.of( + Status.Code.DEADLINE_EXCEEDED, Status.Code.UNAVAILABLE, Status.Code.RESOURCE_EXHAUSTED); + public enum SlidingWindowType { COUNT_BASED, TIME_BASED diff --git a/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptor.java b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptor.java new file mode 100644 index 0000000..355b282 --- /dev/null +++ b/grpc-circuitbreaker-utils/src/main/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptor.java @@ -0,0 +1,179 @@ +package org.hypertrace.circuitbreaker.grpcutils; + +import io.github.resilience4j.circuitbreaker.CircuitBreaker; +import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig; +import io.github.resilience4j.circuitbreaker.CircuitBreakerRegistry; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ForwardingClientCall; +import io.grpc.ForwardingClientCallListener; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public final class ServiceCircuitBreakerInterceptor implements ClientInterceptor { + + private final Map circuitBreakerConfigs; + private final Map> failureStatusCodesByService; + private final CircuitBreakerRegistry circuitBreakerRegistry; + + static ClientInterceptor create(final Map thresholdsByService) { + if (thresholdsByService == null || thresholdsByService.isEmpty()) { + log.debug("No service circuit breakers configured"); + return passthrough(); + } + final Map circuitBreakerConfigs = new HashMap<>(); + final Map> failureStatusCodesByService = new HashMap<>(); + for (final Map.Entry entry : thresholdsByService.entrySet()) { + final CircuitBreakerThresholds thresholds = entry.getValue(); + if (!thresholds.isEnabled()) { + log.debug("Service circuit breaker disabled for {}", entry.getKey()); + continue; + } + circuitBreakerConfigs.put(entry.getKey(), toCircuitBreakerConfig(thresholds)); + failureStatusCodesByService.put(entry.getKey(), thresholds.getFailureStatusCodes()); + } + if (circuitBreakerConfigs.isEmpty()) { + return passthrough(); + } + return new ServiceCircuitBreakerInterceptor( + Map.copyOf(circuitBreakerConfigs), Map.copyOf(failureStatusCodesByService)); + } + + private static ServiceCircuitBreakerInterceptor passthrough() { + return new ServiceCircuitBreakerInterceptor(Map.of(), Map.of()); + } + + private static CircuitBreakerConfig toCircuitBreakerConfig( + final CircuitBreakerThresholds thresholds) { + final Set failureStatusCodes = thresholds.getFailureStatusCodes(); + return CircuitBreakerConfig.custom() + .failureRateThreshold(thresholds.getFailureRateThreshold()) + .slowCallRateThreshold(thresholds.getSlowCallRateThreshold()) + .slowCallDurationThreshold(thresholds.getSlowCallDurationThreshold()) + .slidingWindowType(slidingWindowType(thresholds.getSlidingWindowType())) + .slidingWindowSize(thresholds.getSlidingWindowSize()) + .waitDurationInOpenState(thresholds.getWaitDurationInOpenState()) + .permittedNumberOfCallsInHalfOpenState( + thresholds.getPermittedNumberOfCallsInHalfOpenState()) + .minimumNumberOfCalls(thresholds.getMinimumNumberOfCalls()) + // Returning false makes Resilience4j count this onError as a success. + .recordException( + throwable -> + throwable instanceof StatusRuntimeException + && failureStatusCodes.contains( + ((StatusRuntimeException) throwable).getStatus().getCode())) + .build(); + } + + private static CircuitBreakerConfig.SlidingWindowType slidingWindowType( + final CircuitBreakerThresholds.SlidingWindowType slidingWindowType) { + switch (slidingWindowType) { + case COUNT_BASED: + return CircuitBreakerConfig.SlidingWindowType.COUNT_BASED; + case TIME_BASED: + default: + return CircuitBreakerConfig.SlidingWindowType.TIME_BASED; + } + } + + private ServiceCircuitBreakerInterceptor( + final Map circuitBreakerConfigs, + final Map> failureStatusCodesByService) { + this.circuitBreakerConfigs = circuitBreakerConfigs; + this.failureStatusCodesByService = failureStatusCodesByService; + this.circuitBreakerRegistry = CircuitBreakerRegistry.ofDefaults(); + circuitBreakerConfigs.forEach(this::registerStateChangeLog); + } + + private void registerStateChangeLog(final String serviceName, final CircuitBreakerConfig config) { + final CircuitBreaker circuitBreaker = + circuitBreakerRegistry.circuitBreaker(serviceName, config); + circuitBreaker + .getEventPublisher() + .onStateTransition( + event -> + log.debug( + "Circuit breaker {} changed from {} to {}", + event.getCircuitBreakerName(), + event.getStateTransition().getFromState(), + event.getStateTransition().getToState())); + } + + @Override + public ClientCall interceptCall( + final MethodDescriptor method, + final CallOptions callOptions, + final Channel next) { + final String serviceName = method.getServiceName(); + final CircuitBreakerConfig circuitBreakerConfig = circuitBreakerConfigs.get(serviceName); + if (circuitBreakerConfig == null) { + return next.newCall(method, callOptions); + } + + final CircuitBreaker circuitBreaker = + circuitBreakerRegistry.circuitBreaker(serviceName, circuitBreakerConfig); + if (!circuitBreaker.tryAcquirePermission()) { + log.debug("Circuit breaker is open for {}; rejecting the call", serviceName); + throw Status.RESOURCE_EXHAUSTED + .withDescription("Circuit breaker is open for " + serviceName) + .asRuntimeException(); + } + + final ClientCall delegate; + try { + delegate = next.newCall(method, callOptions); + } catch (RuntimeException exception) { + circuitBreaker.releasePermission(); + throw exception; + } + return recordingCall(delegate, circuitBreaker, failureStatusCodesByService.get(serviceName)); + } + + private ClientCall recordingCall( + final ClientCall delegate, + final CircuitBreaker circuitBreaker, + final Set failureStatusCodes) { + return new ForwardingClientCall.SimpleForwardingClientCall<>(delegate) { + @Override + public void start(final Listener responseListener, final Metadata headers) { + final long startedAtNanos = System.nanoTime(); + super.start( + new ForwardingClientCallListener.SimpleForwardingClientCallListener<>( + responseListener) { + @Override + public void onClose(final Status status, final Metadata trailers) { + final long elapsedNanos = System.nanoTime() - startedAtNanos; + if (status.isOk() || !failureStatusCodes.contains(status.getCode())) { + log.trace( + "Recording success for {} after {} ns with status {}", + circuitBreaker.getName(), + elapsedNanos, + status.getCode()); + circuitBreaker.onSuccess(elapsedNanos, TimeUnit.NANOSECONDS); + } else { + log.debug( + "Recording failure for {} after {} ns with status {}", + circuitBreaker.getName(), + elapsedNanos, + status.getCode()); + circuitBreaker.onError( + elapsedNanos, TimeUnit.NANOSECONDS, status.asRuntimeException()); + } + super.onClose(status, trailers); + } + }, + headers); + } + }; + } +} diff --git a/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistryTest.java b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistryTest.java new file mode 100644 index 0000000..38e599b --- /dev/null +++ b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerChannelRegistryTest.java @@ -0,0 +1,133 @@ +package org.hypertrace.circuitbreaker.grpcutils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import io.grpc.CallOptions; +import io.grpc.ClientCall; +import io.grpc.ManagedChannel; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import org.hypertrace.core.grpcutils.client.GrpcChannelRegistry; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class CircuitBreakerChannelRegistryTest { + + private static final String INSIGHTS_SERVICE = + "ai.traceable.platform.insights.api.v1.InsightsService"; + private static final String ENABLED_INSIGHTS = + "circuit.breaker.config {\n" + + " enabled = true\n" + + " \"" + + INSIGHTS_SERVICE + + "\" {\n" + + " slowCallDurationThreshold = 1h\n" + + " minimumNumberOfCalls = 1\n" + + " failureRateThreshold = 1\n" + + " slidingWindowSize = 1\n" + + " slidingWindowType = COUNT_BASED\n" + + " waitDurationInOpenState = 60s\n" + + " }\n" + + "}\n"; + + @Test + void nullConfigReturnsTheSameRegistry() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + + assertSame(registry, CircuitBreakerChannelRegistry.wrap(registry, (Config) null)); + } + + @Test + void emptyConfigReturnsTheSameRegistry() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + + assertSame(registry, CircuitBreakerChannelRegistry.wrap(registry, ConfigFactory.empty())); + } + + @Test + void disabledConfigReturnsTheSameRegistry() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + Config config = ConfigFactory.parseString("circuit.breaker.config { enabled = false }"); + + assertSame(registry, CircuitBreakerChannelRegistry.wrap(registry, config)); + } + + @Test + void missingEnabledReturnsTheSameRegistry() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + Config config = ConfigFactory.parseString("circuit.breaker.config { }"); + + assertSame(registry, CircuitBreakerChannelRegistry.wrap(registry, config)); + } + + @Test + void nullRegistryIsReturnedUnchanged() { + assertSame(null, CircuitBreakerChannelRegistry.wrap(null, ConfigFactory.empty())); + } + + @Test + void enabledConfigOpensTheBreakerOnUnavailable() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + ManagedChannel rawChannel = mock(ManagedChannel.class); + when(registry.forPlaintextAddress("host", 80)).thenReturn(rawChannel); + + GrpcChannelRegistry wrapped = + CircuitBreakerChannelRegistry.wrap(registry, ConfigFactory.parseString(ENABLED_INSIGHTS)); + assertNotSame(registry, wrapped); + ManagedChannel channel = wrapped.forPlaintextAddress("host", 80); + + completeCall(channel, rawChannel, Status.UNAVAILABLE); + + StatusRuntimeException rejected = + assertThrows( + StatusRuntimeException.class, () -> channel.newCall(method(), CallOptions.DEFAULT)); + assertEquals(Status.Code.RESOURCE_EXHAUSTED, rejected.getStatus().getCode()); + verify(rawChannel, times(1)).newCall(any(), any()); + } + + @Test + void repeatedLookupReturnsTheSameInterceptedChannel() { + GrpcChannelRegistry registry = mock(GrpcChannelRegistry.class); + ManagedChannel rawChannel = mock(ManagedChannel.class); + when(registry.forPlaintextAddress("host", 80)).thenReturn(rawChannel); + + GrpcChannelRegistry wrapped = + CircuitBreakerChannelRegistry.wrap(registry, ConfigFactory.parseString(ENABLED_INSIGHTS)); + + assertSame(wrapped.forPlaintextAddress("host", 80), wrapped.forPlaintextAddress("host", 80)); + } + + @SuppressWarnings("unchecked") + private static void completeCall(ManagedChannel channel, ManagedChannel raw, Status status) { + ClientCall delegateCall = mock(ClientCall.class); + when(raw.newCall(any(), any())).thenReturn((ClientCall) delegateCall); + ClientCall.Listener listener = mock(ClientCall.Listener.class); + + ClientCall call = channel.newCall(method(), CallOptions.DEFAULT); + call.start(listener, new Metadata()); + + ArgumentCaptor> captor = + ArgumentCaptor.forClass(ClientCall.Listener.class); + verify(delegateCall).start(captor.capture(), any()); + captor.getValue().onClose(status, new Metadata()); + } + + @SuppressWarnings("unchecked") + private static MethodDescriptor method() { + MethodDescriptor method = mock(MethodDescriptor.class); + when(method.getServiceName()).thenReturn(INSIGHTS_SERVICE); + return method; + } +} diff --git a/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParserTest.java b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParserTest.java new file mode 100644 index 0000000..1d7dfcb --- /dev/null +++ b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/CircuitBreakerConfigParserTest.java @@ -0,0 +1,49 @@ +package org.hypertrace.circuitbreaker.grpcutils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.typesafe.config.ConfigException; +import com.typesafe.config.ConfigFactory; +import io.grpc.Status.Code; +import java.time.Duration; +import java.util.Set; +import org.junit.jupiter.api.Test; + +class CircuitBreakerConfigParserTest { + + @Test + void omittedFailureStatusesUseTheLibraryDefault() { + CircuitBreakerThresholds thresholds = + CircuitBreakerConfigParser.parseConfig(ConfigFactory.parseString("enabled = true")) + .build() + .getDefaultThresholds(); + + assertEquals( + Set.of(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE, Code.RESOURCE_EXHAUSTED), + thresholds.getFailureStatusCodes()); + assertEquals(Duration.ofSeconds(2), thresholds.getSlowCallDurationThreshold()); + } + + @Test + void defaultThresholdsReplaceFailureStatuses() { + CircuitBreakerThresholds thresholds = + CircuitBreakerConfigParser.parseConfig( + ConfigFactory.parseString( + "defaultThresholds { failureStatusCodes = [UNAVAILABLE] }")) + .build() + .getDefaultThresholds(); + + assertEquals(Set.of(Code.UNAVAILABLE), thresholds.getFailureStatusCodes()); + } + + @Test + void unknownFailureStatusFailsParsing() { + assertThrows( + ConfigException.BadValue.class, + () -> + CircuitBreakerConfigParser.parseConfig( + ConfigFactory.parseString( + "defaultThresholds { failureStatusCodes = [NOT_A_STATUS] }"))); + } +} diff --git a/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptorTest.java b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptorTest.java new file mode 100644 index 0000000..cb51954 --- /dev/null +++ b/grpc-circuitbreaker-utils/src/test/java/org/hypertrace/circuitbreaker/grpcutils/ServiceCircuitBreakerInterceptorTest.java @@ -0,0 +1,220 @@ +package org.hypertrace.circuitbreaker.grpcutils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigException; +import com.typesafe.config.ConfigFactory; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class ServiceCircuitBreakerInterceptorTest { + + private static final String INSIGHTS_SERVICE = + "ai.traceable.platform.insights.api.v1.InsightsService"; + private static final String OTHER_SERVICE = "ai.traceable.platform.other.OtherService"; + private static final String OPEN_QUICKLY = + "slowCallDurationThreshold = 1h\n" + + "minimumNumberOfCalls = 1\n" + + "failureRateThreshold = 1\n" + + "slidingWindowSize = 1\n" + + "slidingWindowType = COUNT_BASED\n" + + "waitDurationInOpenState = 60s\n"; + + @Test + void absentConfigCallsTheDelegate() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService(ConfigFactory.empty()); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + verify(channel, times(2)).newCall(any(), any()); + } + + @Test + void missingRootEnabledCallsTheDelegate() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker(serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + verify(channel, times(2)).newCall(any(), any()); + } + + @Test + void rootEnabledFalseCallsTheDelegate() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker("enabled = false\n" + serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + verify(channel, times(2)).newCall(any(), any()); + } + + @Test + void enabledServiceCallsTheDelegate() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker( + "enabled = true\n" + + "defaultThresholds { slowCallDurationThreshold = 1h }\n" + + serviceBlock(INSIGHTS_SERVICE, ""))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.OK); + verify(channel).newCall(any(), any()); + } + + @Test + void serviceFailuresDoNotTripAnotherService() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker( + "enabled = true\n" + + serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY) + + serviceBlock(OTHER_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + complete(interceptor, channel, OTHER_SERVICE, Status.OK); + verify(channel, times(2)).newCall(any(), any()); + + StatusRuntimeException rejected = reject(interceptor, channel, INSIGHTS_SERVICE); + assertEquals(Status.Code.RESOURCE_EXHAUSTED, rejected.getStatus().getCode()); + verify(channel, times(2)).newCall(any(), any()); + } + + @Test + void unavailableOpensTheBreakerWithoutCallingTheDelegate() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker("enabled = true\n" + serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + + StatusRuntimeException rejected = reject(interceptor, channel, INSIGHTS_SERVICE); + assertEquals(Status.Code.RESOURCE_EXHAUSTED, rejected.getStatus().getCode()); + assertNotNull(rejected.getStatus().getDescription()); + assertFalse(rejected.getStatus().getDescription().isEmpty()); + verify(channel, times(1)).newCall(any(), any()); + } + + @Test + void internalDoesNotOpenTheBreakerForTheDefaultFailureStatuses() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker("enabled = true\n" + serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.INTERNAL); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.OK); + verify(channel, times(2)).newCall(any(), any()); + } + + @Test + void onlyConfiguredUnavailableCountsAsFailure() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker( + "enabled = true\n" + + "defaultThresholds { failureStatusCodes = [UNAVAILABLE] }\n" + + serviceBlock(INSIGHTS_SERVICE, OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.INTERNAL); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.DEADLINE_EXCEEDED); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + + StatusRuntimeException rejected = reject(interceptor, channel, INSIGHTS_SERVICE); + assertEquals(Status.Code.RESOURCE_EXHAUSTED, rejected.getStatus().getCode()); + verify(channel, times(3)).newCall(any(), any()); + } + + @Test + void unknownFailureStatusFailsConstruction() { + Config config = + circuitBreaker( + "enabled = true\n" + + serviceBlock(INSIGHTS_SERVICE, "failureStatusCodes = [NOT_A_STATUS]")); + + assertThrows( + ConfigException.BadValue.class, () -> CircuitBreakerInterceptorFactory.forService(config)); + } + + @Test + void disabledServicePassesThrough() { + Channel channel = mock(Channel.class); + ClientInterceptor interceptor = + CircuitBreakerInterceptorFactory.forService( + circuitBreaker( + "enabled = true\n" + + serviceBlock(INSIGHTS_SERVICE, "enabled = false\n" + OPEN_QUICKLY))); + + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + complete(interceptor, channel, INSIGHTS_SERVICE, Status.UNAVAILABLE); + verify(channel, times(2)).newCall(any(), any()); + } + + private static Config circuitBreaker(String body) { + return ConfigFactory.parseString("circuit.breaker.config {\n" + body + "\n}"); + } + + private static String serviceBlock(String serviceName, String body) { + return "\"" + serviceName + "\" {\n" + body + "\n}\n"; + } + + @SuppressWarnings("unchecked") + private static void complete( + ClientInterceptor interceptor, Channel channel, String serviceName, Status status) { + ClientCall delegate = mock(ClientCall.class); + when(channel.newCall(any(), any())).thenReturn(delegate); + ClientCall.Listener listener = mock(ClientCall.Listener.class); + Metadata headers = new Metadata(); + ClientCall call = + interceptor.interceptCall(method(serviceName), CallOptions.DEFAULT, channel); + call.start(listener, headers); + + ArgumentCaptor> captor = + ArgumentCaptor.forClass(ClientCall.Listener.class); + verify(delegate).start(captor.capture(), any()); + Metadata trailers = new Metadata(); + captor.getValue().onClose(status, trailers); + verify(listener).onClose(status, trailers); + } + + private static StatusRuntimeException reject( + ClientInterceptor interceptor, Channel channel, String serviceName) { + return assertThrows( + StatusRuntimeException.class, + () -> interceptor.interceptCall(method(serviceName), CallOptions.DEFAULT, channel)); + } + + @SuppressWarnings("unchecked") + private static MethodDescriptor method(String serviceName) { + MethodDescriptor method = mock(MethodDescriptor.class); + when(method.getServiceName()).thenReturn(serviceName); + return method; + } +} diff --git a/grpc-client-utils/build.gradle.kts b/grpc-client-utils/build.gradle.kts index 530a934..c11bcae 100644 --- a/grpc-client-utils/build.gradle.kts +++ b/grpc-client-utils/build.gradle.kts @@ -10,6 +10,7 @@ dependencies { api(commonLibs.grpc.context) api(commonLibs.grpc.api) api(commonLibs.grpc.inprocess) + api(commonLibs.typesafe.config) implementation(projects.grpcContextUtils) implementation(commonLibs.slf4j2.api) diff --git a/grpc-client-utils/gradle.lockfile b/grpc-client-utils/gradle.lockfile index beeec55..ecd50ee 100644 --- a/grpc-client-utils/gradle.lockfile +++ b/grpc-client-utils/gradle.lockfile @@ -21,6 +21,7 @@ com.google.j2objc:j2objc-annotations:2.8=compileClasspath com.google.j2objc:j2objc-annotations:3.0.0=runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.google.protobuf:protobuf-java-util:3.25.8=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath com.google.protobuf:protobuf-java:3.25.8=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath +com.typesafe:config:1.4.2=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-api:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-bom:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath io.grpc:grpc-context:1.75.0=compileClasspath,runtimeClasspath,testCompileClasspath,testRuntimeClasspath diff --git a/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistry.java b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistry.java new file mode 100644 index 0000000..830535e --- /dev/null +++ b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistry.java @@ -0,0 +1,106 @@ +package org.hypertrace.core.grpcutils.client; + +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ClientInterceptors; +import io.grpc.ManagedChannel; +import io.grpc.MethodDescriptor; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Objects; +import lombok.experimental.Delegate; + +public final class InterceptingChannelRegistry extends GrpcChannelRegistry { + + public static GrpcChannelRegistry wrap( + final GrpcChannelRegistry registry, final ClientInterceptor interceptor) { + return new InterceptingChannelRegistry( + Objects.requireNonNull(registry), Objects.requireNonNull(interceptor)); + } + + // Lombok skips a delegate method only when this class overrides it. forAddress is not + // overridden: the superclass implementation calls forPlaintextAddress. Delegating it would + // return the raw channel and skip the interceptor. + @Delegate(excludes = ForAddress.class) + private final GrpcChannelRegistry delegate; + + private final ClientInterceptor interceptor; + private final Map interceptedChannels = + Collections.synchronizedMap(new IdentityHashMap<>()); + + private InterceptingChannelRegistry( + final GrpcChannelRegistry delegate, final ClientInterceptor interceptor) { + this.delegate = delegate; + this.interceptor = interceptor; + } + + @Override + public ManagedChannel forSecureAddress(final String host, final int port) { + return this.intercept(this.delegate.forSecureAddress(host, port)); + } + + @Override + public ManagedChannel forSecureAddress( + final String host, final int port, final GrpcChannelConfig config) { + return this.intercept(this.delegate.forSecureAddress(host, port, config)); + } + + @Override + public ManagedChannel forPlaintextAddress(final String host, final int port) { + return this.intercept(this.delegate.forPlaintextAddress(host, port)); + } + + @Override + public ManagedChannel forPlaintextAddress( + final String host, final int port, final GrpcChannelConfig config) { + return this.intercept(this.delegate.forPlaintextAddress(host, port, config)); + } + + private ManagedChannel intercept(final ManagedChannel channel) { + return this.interceptedChannels.computeIfAbsent( + channel, raw -> new InterceptedManagedChannel(raw, this.interceptor)); + } + + private interface ForAddress { + ManagedChannel forAddress(String host, int port); + } + + private static final class InterceptedManagedChannel extends ManagedChannel { + // Overridden methods are not delegated. shutdown and shutdownNow return this wrapper. + @Delegate private final ManagedChannel delegate; + + private final Channel interceptedChannel; + + private InterceptedManagedChannel( + final ManagedChannel delegate, final ClientInterceptor interceptor) { + this.delegate = delegate; + this.interceptedChannel = ClientInterceptors.intercept(delegate, interceptor); + } + + @Override + public ClientCall newCall( + final MethodDescriptor methodDescriptor, final CallOptions callOptions) { + return this.interceptedChannel.newCall(methodDescriptor, callOptions); + } + + @Override + public String authority() { + return this.interceptedChannel.authority(); + } + + @Override + public ManagedChannel shutdown() { + this.delegate.shutdown(); + return this; + } + + @Override + public ManagedChannel shutdownNow() { + this.delegate.shutdownNow(); + return this; + } + } +} diff --git a/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistry.java b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistry.java new file mode 100644 index 0000000..e862045 --- /dev/null +++ b/grpc-client-utils/src/main/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistry.java @@ -0,0 +1,32 @@ +package org.hypertrace.core.grpcutils.client; + +import com.typesafe.config.Config; +import java.time.Duration; +import lombok.AccessLevel; +import lombok.NoArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +@NoArgsConstructor(access = AccessLevel.PRIVATE) +public final class TimeoutChannelRegistry { + private static final String TIMEOUT_PATH = "grpc.client.timeout"; + private static final String ENABLED_KEY = "enabled"; + private static final String DEFAULT_TIMEOUT_KEY = "defaultTimeout"; + + public static GrpcChannelRegistry wrap(final GrpcChannelRegistry registry, final Config config) { + if (registry == null || config == null || !config.hasPath(TIMEOUT_PATH)) { + log.debug("grpc.client.timeout is absent; leaving the channel registry unchanged"); + return registry; + } + + final Config timeoutConfig = config.getConfig(TIMEOUT_PATH); + if (!timeoutConfig.hasPath(ENABLED_KEY) || !timeoutConfig.getBoolean(ENABLED_KEY)) { + log.debug("grpc.client.timeout is disabled; leaving the channel registry unchanged"); + return registry; + } + + // An enabled block with no timeout must fail at startup rather than silently apply no deadline. + final Duration timeout = timeoutConfig.getDuration(DEFAULT_TIMEOUT_KEY); + return InterceptingChannelRegistry.wrap(registry, new DefaultTimeoutClientInterceptor(timeout)); + } +} diff --git a/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistryTest.java b/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistryTest.java new file mode 100644 index 0000000..cfb2907 --- /dev/null +++ b/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/InterceptingChannelRegistryTest.java @@ -0,0 +1,155 @@ +package org.hypertrace.core.grpcutils.client; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.Deadline; +import io.grpc.ManagedChannel; +import io.grpc.MethodDescriptor; +import io.grpc.MethodDescriptor.Marshaller; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class InterceptingChannelRegistryTest { + + private static final MethodDescriptor METHOD = + MethodDescriptor.newBuilder() + .setType(MethodDescriptor.MethodType.UNARY) + .setFullMethodName("test.Service/Method") + .setRequestMarshaller(new StringMarshaller()) + .setResponseMarshaller(new StringMarshaller()) + .build(); + + private final List> interceptedMethods = new ArrayList<>(); + private final ClientInterceptor recordingInterceptor = + new ClientInterceptor() { + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + interceptedMethods.add(method); + return next.newCall(method, callOptions); + } + }; + + private GrpcChannelRegistry delegate; + private ManagedChannel rawChannel; + private GrpcChannelRegistry wrapper; + + @BeforeEach + void beforeEach() { + this.delegate = mock(GrpcChannelRegistry.class); + this.rawChannel = mock(ManagedChannel.class); + this.wrapper = InterceptingChannelRegistry.wrap(this.delegate, this.recordingInterceptor); + } + + @Test + void wrapReturnsDistinctRegistry() { + assertNotSame(this.delegate, this.wrapper); + assertThrows( + NullPointerException.class, + () -> InterceptingChannelRegistry.wrap(null, this.recordingInterceptor)); + assertThrows( + NullPointerException.class, () -> InterceptingChannelRegistry.wrap(this.delegate, null)); + } + + @Test + void plaintextAndSecureChannelsCarryTheInterceptor() { + GrpcChannelConfig config = GrpcChannelConfig.builder().build(); + ManagedChannel secureRaw = mock(ManagedChannel.class); + when(this.delegate.forPlaintextAddress("host", 1)).thenReturn(this.rawChannel); + when(this.delegate.forPlaintextAddress("host", 2, config)).thenReturn(this.rawChannel); + when(this.delegate.forSecureAddress("host", 3)).thenReturn(secureRaw); + when(this.delegate.forSecureAddress("host", 4, config)).thenReturn(secureRaw); + + List channels = + List.of( + this.wrapper.forPlaintextAddress("host", 1), + this.wrapper.forPlaintextAddress("host", 2, config), + this.wrapper.forSecureAddress("host", 3), + this.wrapper.forSecureAddress("host", 4, config)); + + channels.forEach(channel -> channel.newCall(METHOD, CallOptions.DEFAULT)); + + assertEquals(List.of(METHOD, METHOD, METHOD, METHOD), this.interceptedMethods); + verify(this.rawChannel, times(2)).newCall(METHOD, CallOptions.DEFAULT); + verify(secureRaw, times(2)).newCall(METHOD, CallOptions.DEFAULT); + } + + @Test + void repeatedLookupReturnsTheSameInterceptedChannel() { + when(this.delegate.forPlaintextAddress("host", 1)).thenReturn(this.rawChannel); + + ManagedChannel first = this.wrapper.forPlaintextAddress("host", 1); + + assertNotSame(this.rawChannel, first); + assertSame(first, this.wrapper.forPlaintextAddress("host", 1)); + } + + @Test + void forAddressUsesThePlaintextOverride() { + GrpcChannelConfig config = GrpcChannelConfig.builder().build(); + when(this.delegate.forPlaintextAddress("host", 1, config)).thenReturn(this.rawChannel); + + ManagedChannel channel = this.wrapper.forAddress("host", 1); + channel.newCall(METHOD, CallOptions.DEFAULT); + + verify(this.delegate).forPlaintextAddress("host", 1, config); + verify(this.delegate, never()).forAddress(anyString(), anyInt()); + assertEquals(List.of(METHOD), this.interceptedMethods); + } + + @Test + void shutdownMethodsReturnTheWrapper() { + when(this.delegate.forPlaintextAddress("host", 1)).thenReturn(this.rawChannel); + ManagedChannel channel = this.wrapper.forPlaintextAddress("host", 1); + + assertSame(channel, channel.shutdown()); + assertSame(channel, channel.shutdownNow()); + + verify(this.rawChannel).shutdown(); + verify(this.rawChannel).shutdownNow(); + } + + @Test + void registryShutdownMethodsForwardToTheDelegate() { + Deadline deadline = Deadline.after(1, TimeUnit.SECONDS); + when(this.delegate.isShutdown()).thenReturn(true); + + this.wrapper.shutdown(); + this.wrapper.shutdown(deadline); + + verify(this.delegate).shutdown(); + verify(this.delegate).shutdown(deadline); + assertTrue(this.wrapper.isShutdown()); + } + + private static final class StringMarshaller implements Marshaller { + @Override + public InputStream stream(String value) { + throw new UnsupportedOperationException(); + } + + @Override + public String parse(InputStream stream) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistryTest.java b/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistryTest.java new file mode 100644 index 0000000..a340911 --- /dev/null +++ b/grpc-client-utils/src/test/java/org/hypertrace/core/grpcutils/client/TimeoutChannelRegistryTest.java @@ -0,0 +1,137 @@ +package org.hypertrace.core.grpcutils.client; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigException; +import com.typesafe.config.ConfigFactory; +import io.grpc.CallOptions; +import io.grpc.ClientCall; +import io.grpc.Deadline; +import io.grpc.ManagedChannel; +import io.grpc.MethodDescriptor; +import io.grpc.MethodDescriptor.Marshaller; +import java.io.InputStream; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TimeoutChannelRegistryTest { + + private static final MethodDescriptor METHOD = + MethodDescriptor.newBuilder() + .setType(MethodDescriptor.MethodType.UNARY) + .setFullMethodName("test.Service/Method") + .setRequestMarshaller(new StringMarshaller()) + .setResponseMarshaller(new StringMarshaller()) + .build(); + + private final AtomicReference capturedOptions = new AtomicReference<>(); + + private GrpcChannelRegistry delegate; + private ManagedChannel rawChannel; + + @BeforeEach + @SuppressWarnings("unchecked") + void beforeEach() { + this.delegate = mock(GrpcChannelRegistry.class); + this.rawChannel = mock(ManagedChannel.class); + when(this.delegate.forPlaintextAddress("host", 1)).thenReturn(this.rawChannel); + when(this.rawChannel.newCall(any(), any())) + .thenAnswer( + invocation -> { + this.capturedOptions.set(invocation.getArgument(1)); + return mock(ClientCall.class); + }); + } + + @Test + void missingTimeoutBlockReturnsTheSameRegistry() { + Config config = ConfigFactory.parseString("grpc.client.other = 1"); + + assertSame(this.delegate, TimeoutChannelRegistry.wrap(this.delegate, config)); + } + + @Test + void disabledReturnsTheSameRegistry() { + Config config = + ConfigFactory.parseString( + "grpc.client.timeout { enabled = false, defaultTimeout = 1500ms }"); + + assertSame(this.delegate, TimeoutChannelRegistry.wrap(this.delegate, config)); + } + + @Test + void missingEnabledReturnsTheSameRegistry() { + Config config = ConfigFactory.parseString("grpc.client.timeout { defaultTimeout = 1500ms }"); + + assertSame(this.delegate, TimeoutChannelRegistry.wrap(this.delegate, config)); + } + + @Test + void nullRegistryOrConfigReturnsTheInput() { + Config config = + ConfigFactory.parseString("grpc.client.timeout { enabled = true, defaultTimeout = 1s }"); + + assertSame(this.delegate, TimeoutChannelRegistry.wrap(this.delegate, null)); + assertEquals(null, TimeoutChannelRegistry.wrap(null, config)); + } + + @Test + void enabledWithoutDefaultTimeoutThrows() { + Config config = ConfigFactory.parseString("grpc.client.timeout { enabled = true }"); + + assertThrows(ConfigException.class, () -> TimeoutChannelRegistry.wrap(this.delegate, config)); + } + + @Test + void enabledAppliesTheDefaultDeadlineWhenNoneIsPresent() { + GrpcChannelRegistry wrapped = TimeoutChannelRegistry.wrap(this.delegate, enabledConfig()); + assertNotSame(this.delegate, wrapped); + + wrapped.forPlaintextAddress("host", 1).newCall(METHOD, CallOptions.DEFAULT); + + Deadline deadline = this.capturedOptions.get().getDeadline(); + assertNotNull(deadline); + long remainingMillis = deadline.timeRemaining(TimeUnit.MILLISECONDS); + assertTrue(remainingMillis > 0 && remainingMillis <= 1500, "remaining: " + remainingMillis); + } + + @Test + void existingDeadlineIsKept() { + GrpcChannelRegistry wrapped = TimeoutChannelRegistry.wrap(this.delegate, enabledConfig()); + Deadline existing = Deadline.after(10, TimeUnit.MINUTES); + + wrapped + .forPlaintextAddress("host", 1) + .newCall(METHOD, CallOptions.DEFAULT.withDeadline(existing)); + + assertSame(existing, this.capturedOptions.get().getDeadline()); + } + + private static Config enabledConfig() { + return ConfigFactory.parseString( + "grpc.client.timeout { enabled = true, defaultTimeout = 1500ms }"); + } + + private static final class StringMarshaller implements Marshaller { + @Override + public InputStream stream(String value) { + throw new UnsupportedOperationException(); + } + + @Override + public String parse(InputStream stream) { + throw new UnsupportedOperationException(); + } + } +}