From 062ebb4d77e9906978fadd6c84a6b957f90f369c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Vitor=20Stein?= <joaovstein@gmail.com>
Date: Wed, 3 Jul 2024 12:19:22 -0300
Subject: [PATCH] grpc-core: use retryThrottling from defaultServiceConfig

---
 .../io/grpc/internal/ManagedChannelImpl.java  |  7 +-
 .../grpc/internal/ManagedChannelImplTest.java | 67 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
index b21fc97e64..c5c7b66e15 100644
--- a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
+++ b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
@@ -616,7 +616,7 @@ final class ManagedChannelImpl extends ManagedChannel implements
           parsedDefaultServiceConfig.getError());
       this.defaultServiceConfig =
           (ManagedChannelServiceConfig) parsedDefaultServiceConfig.getConfig();
-      this.lastServiceConfig = this.defaultServiceConfig;
+      this.transportProvider.throttle = this.defaultServiceConfig.getRetryThrottling();
     } else {
       this.defaultServiceConfig = null;
     }
@@ -708,6 +708,11 @@ final class ManagedChannelImpl extends ManagedChannel implements
   InternalConfigSelector getConfigSelector() {
     return realChannel.configSelector.get();
   }
+  
+  @VisibleForTesting
+  boolean hasThrottle() {
+    return this.transportProvider.throttle != null;
+  }
 
   /**
    * Initiates an orderly shutdown in which preexisting calls continue but new calls are immediately
diff --git a/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java b/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
index 2b598db1cf..1d6492f791 100644
--- a/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
+++ b/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
@@ -4058,6 +4058,40 @@ public class ManagedChannelImplTest {
     }
   }
 
+  @Test
+  public void disableServiceConfigLookUp_withDefaultConfig_withRetryThrottle() throws Exception {
+    LoadBalancerRegistry.getDefaultRegistry().register(mockLoadBalancerProvider);
+    try {
+      FakeNameResolverFactory nameResolverFactory =
+              new FakeNameResolverFactory.Builder(expectedUri)
+                      .setServers(ImmutableList.of(addressGroup)).build();
+      channelBuilder.nameResolverFactory(nameResolverFactory);
+      channelBuilder.disableServiceConfigLookUp();
+      channelBuilder.enableRetry();
+      Map<String, Object> defaultServiceConfig =
+              parseConfig("{"
+                      + "\"retryThrottling\":{\"maxTokens\": 1, \"tokenRatio\": 1},"
+                      + "\"methodConfig\":[{"
+                      + "\"name\":[{\"service\":\"SimpleService1\"}],"
+                      + "\"waitForReady\":true"
+                      + "}]}");
+      channelBuilder.defaultServiceConfig(defaultServiceConfig);
+
+      createChannel();
+
+      ArgumentCaptor<ResolvedAddresses> resultCaptor =
+              ArgumentCaptor.forClass(ResolvedAddresses.class);
+      verify(mockLoadBalancer).acceptResolvedAddresses(resultCaptor.capture());
+      assertThat(resultCaptor.getValue().getAddresses()).containsExactly(addressGroup);
+      assertThat(resultCaptor.getValue().getAttributes().get(InternalConfigSelector.KEY)).isNull();
+      verify(mockLoadBalancer, never()).handleNameResolutionError(any(Status.class));
+      assertThat(channel.hasThrottle()).isTrue();
+
+    } finally {
+      LoadBalancerRegistry.getDefaultRegistry().deregister(mockLoadBalancerProvider);
+    }
+  }
+
   @Test
   public void enableServiceConfigLookUp_noDefaultConfig() throws Exception {
     LoadBalancerRegistry.getDefaultRegistry().register(mockLoadBalancerProvider);
@@ -4165,6 +4199,39 @@ public class ManagedChannelImplTest {
     }
   }
 
+
+  @Test
+  public void enableServiceConfigLookUp_usingDefaultConfig_withRetryThrottling() throws Exception {
+    LoadBalancerRegistry.getDefaultRegistry().register(mockLoadBalancerProvider);
+    try {
+      FakeNameResolverFactory nameResolverFactory =
+              new FakeNameResolverFactory.Builder(expectedUri)
+                      .setServers(ImmutableList.of(addressGroup)).build();
+      channelBuilder.nameResolverFactory(nameResolverFactory);
+      channelBuilder.enableRetry();
+      Map<String, Object> defaultServiceConfig =
+              parseConfig("{"
+                      + "\"retryThrottling\":{\"maxTokens\": 1, \"tokenRatio\": 1},"
+                      + "\"methodConfig\":[{"
+                      + "\"name\":[{\"service\":\"SimpleService1\"}],"
+                      + "\"waitForReady\":true"
+                      + "}]}");
+      channelBuilder.defaultServiceConfig(defaultServiceConfig);
+
+      createChannel();
+
+      ArgumentCaptor<ResolvedAddresses> resultCaptor =
+              ArgumentCaptor.forClass(ResolvedAddresses.class);
+      verify(mockLoadBalancer).acceptResolvedAddresses(resultCaptor.capture());
+      assertThat(resultCaptor.getValue().getAddresses()).containsExactly(addressGroup);
+      assertThat(resultCaptor.getValue().getAttributes().get(InternalConfigSelector.KEY)).isNull();
+      verify(mockLoadBalancer, never()).handleNameResolutionError(any(Status.class));
+      assertThat(channel.hasThrottle()).isTrue();
+    } finally {
+      LoadBalancerRegistry.getDefaultRegistry().deregister(mockLoadBalancerProvider);
+    }
+  }
+
   @Test
   public void enableServiceConfigLookUp_resolverReturnsNoConfig_noDefaultConfig() {
     LoadBalancerRegistry.getDefaultRegistry().register(mockLoadBalancerProvider);