Merge branch 'grpc:master' into Issue_fixed_12142

Author: Sangamesh1997

api/src/main/java/io/grpc/NameResolver.java

@@ -303,6 +303,7 @@ public static final class Args {
     @Nullable private final Executor executor;
     @Nullable private final String overrideAuthority;
     @Nullable private final MetricRecorder metricRecorder;
+    @Nullable private final NameResolverRegistry nameResolverRegistry;
     @Nullable private final IdentityHashMap<Key<?>, Object> customArgs;
 
     private Args(Builder builder) {
@@ -316,6 +317,7 @@ private Args(Builder builder) {
       this.executor = builder.executor;
       this.overrideAuthority = builder.overrideAuthority;
       this.metricRecorder = builder.metricRecorder;
+      this.nameResolverRegistry = builder.nameResolverRegistry;
       this.customArgs = cloneCustomArgs(builder.customArgs);
     }
 
@@ -447,6 +449,18 @@ public MetricRecorder getMetricRecorder() {
       return metricRecorder;
     }
 
+    /**
+     * Returns the {@link NameResolverRegistry} that the Channel uses to look for {@link
+     * NameResolver}s.
+     *
+     * @since 1.74.0
+     */
+    public NameResolverRegistry getNameResolverRegistry() {
+      if (nameResolverRegistry == null) {
+        throw new IllegalStateException("NameResolverRegistry is not set in Builder");
+      }
+      return nameResolverRegistry;
+    }
 
     @Override
     public String toString() {
@@ -461,6 +475,7 @@ public String toString() {
           .add("executor", executor)
           .add("overrideAuthority", overrideAuthority)
           .add("metricRecorder", metricRecorder)
+          .add("nameResolverRegistry", nameResolverRegistry)
           .toString();
     }
 
@@ -480,6 +495,7 @@ public Builder toBuilder() {
       builder.setOffloadExecutor(executor);
       builder.setOverrideAuthority(overrideAuthority);
       builder.setMetricRecorder(metricRecorder);
+      builder.setNameResolverRegistry(nameResolverRegistry);
       builder.customArgs = cloneCustomArgs(customArgs);
       return builder;
     }
@@ -508,6 +524,7 @@ public static final class Builder {
       private Executor executor;
       private String overrideAuthority;
       private MetricRecorder metricRecorder;
+      private NameResolverRegistry nameResolverRegistry;
       private IdentityHashMap<Key<?>, Object> customArgs;
 
       Builder() {
@@ -614,6 +631,16 @@ public Builder setMetricRecorder(MetricRecorder metricRecorder) {
         return this;
       }
 
+      /**
+       * See {@link Args#getNameResolverRegistry}.  This is an optional field.
+       *
+       * @since 1.74.0
+       */
+      public Builder setNameResolverRegistry(NameResolverRegistry registry) {
+        this.nameResolverRegistry = registry;
+        return this;
+      }
+
       /**
        * Builds an {@link Args}.
        *

core/src/main/java/io/grpc/internal/ManagedChannelImpl.java

@@ -94,7 +94,6 @@
 import io.grpc.internal.ManagedChannelServiceConfig.ServiceConfigConvertedSelector;
 import io.grpc.internal.RetriableStream.ChannelBufferMeter;
 import io.grpc.internal.RetriableStream.Throttle;
-import io.grpc.internal.RetryingNameResolver.ResolutionResultListener;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -598,7 +597,8 @@ ClientStream newSubstream(
             .setChannelLogger(channelLogger)
             .setOffloadExecutor(this.offloadExecutorHolder)
             .setOverrideAuthority(this.authorityOverride)
-            .setMetricRecorder(this.metricRecorder);
+            .setMetricRecorder(this.metricRecorder)
+            .setNameResolverRegistry(builder.nameResolverRegistry);
     builder.copyAllNameResolverCustomArgsTo(nameResolverArgsBuilder);
     this.nameResolverArgs = nameResolverArgsBuilder.build();
     this.nameResolver = getNameResolver(
@@ -686,11 +686,7 @@ static NameResolver getNameResolver(
     // We wrap the name resolver in a RetryingNameResolver to give it the ability to retry failures.
     // TODO: After a transition period, all NameResolver implementations that need retry should use
     //       RetryingNameResolver directly and this step can be removed.
-    NameResolver usedNameResolver = new RetryingNameResolver(resolver,
-          new BackoffPolicyRetryScheduler(new ExponentialBackoffPolicy.Provider(),
-              nameResolverArgs.getScheduledExecutorService(),
-              nameResolverArgs.getSynchronizationContext()),
-          nameResolverArgs.getSynchronizationContext());
+    NameResolver usedNameResolver = RetryingNameResolver.wrap(resolver, nameResolverArgs);
 
     if (overrideAuthority == null) {
       return usedNameResolver;
@@ -1653,18 +1649,7 @@ final class NameResolverListener extends NameResolver.Listener2 {
 
     @Override
     public void onResult(final ResolutionResult resolutionResult) {
-      final class NamesResolved implements Runnable {
-
-        @Override
-        public void run() {
-          Status status = onResult2(resolutionResult);
-          ResolutionResultListener resolutionResultListener = resolutionResult.getAttributes()
-              .get(RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY);
-          resolutionResultListener.resolutionAttempted(status);
-        }
-      }
-
-      syncContext.execute(new NamesResolved());
+      syncContext.execute(() -> onResult2(resolutionResult));
     }
 
     @SuppressWarnings("ReferenceEquality")

core/src/main/java/io/grpc/internal/RetryingNameResolver.java

@@ -17,7 +17,6 @@
 package io.grpc.internal;
 
 import com.google.common.annotations.VisibleForTesting;
-import io.grpc.Attributes;
 import io.grpc.NameResolver;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext;
@@ -28,16 +27,22 @@
  *
  * <p>The {@link NameResolver} used with this
  */
-final class RetryingNameResolver extends ForwardingNameResolver {
+public final class RetryingNameResolver extends ForwardingNameResolver {
+  public static NameResolver wrap(NameResolver retriedNameResolver, Args args) {
+    // For migration, this might become conditional
+    return new RetryingNameResolver(
+        retriedNameResolver,
+        new BackoffPolicyRetryScheduler(
+            new ExponentialBackoffPolicy.Provider(),
+            args.getScheduledExecutorService(),
+            args.getSynchronizationContext()),
+        args.getSynchronizationContext());
+  }
 
   private final NameResolver retriedNameResolver;
   private final RetryScheduler retryScheduler;
   private final SynchronizationContext syncContext;
 
-  static final Attributes.Key<ResolutionResultListener> RESOLUTION_RESULT_LISTENER_KEY
-      = Attributes.Key.create(
-          "io.grpc.internal.RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY");
-
   /**
    * Creates a new {@link RetryingNameResolver}.
    *
@@ -88,18 +93,7 @@ private class RetryingListener extends Listener2 {
 
     @Override
     public void onResult(ResolutionResult resolutionResult) {
-      // If the resolution result listener is already an attribute it indicates that a name resolver
-      // has already been wrapped with this class. This indicates a misconfiguration.
-      if (resolutionResult.getAttributes().get(RESOLUTION_RESULT_LISTENER_KEY) != null) {
-        throw new IllegalStateException(
-            "RetryingNameResolver can only be used once to wrap a NameResolver");
-      }
-
-      // To have retry behavior for name resolvers that haven't migrated to onResult2.
-      delegateListener.onResult(resolutionResult.toBuilder().setAttributes(
-              resolutionResult.getAttributes().toBuilder()
-                  .set(RESOLUTION_RESULT_LISTENER_KEY, new ResolutionResultListener()).build())
-          .build());
+      syncContext.execute(() -> onResult2(resolutionResult));
     }
 
     @Override
@@ -119,19 +113,4 @@ public void onError(Status error) {
       syncContext.execute(() -> retryScheduler.schedule(new DelayedNameResolverRefresh()));
     }
   }
-
-  /**
-   * Simple callback class to store in {@link ResolutionResult} attributes so that
-   * ManagedChannel can indicate if the resolved addresses were accepted. Temporary until
-   * the Listener2.onResult() API can be changed to return a boolean for this purpose.
-   */
-  class ResolutionResultListener {
-    public void resolutionAttempted(Status successStatus) {
-      if (successStatus.isOk()) {
-        retryScheduler.reset();
-      } else {
-        retryScheduler.schedule(new DelayedNameResolverRefresh());
-      }
-    }
-  }
 }

core/src/test/java/io/grpc/internal/DnsNameResolverTest.java

@@ -207,14 +207,7 @@ private RetryingNameResolver newResolver(
 
     // In practice the DNS name resolver provider always wraps the resolver in a
     // RetryingNameResolver which adds retry capabilities to it. We use the same setup here.
-    return new RetryingNameResolver(
-        dnsResolver,
-        new BackoffPolicyRetryScheduler(
-            new ExponentialBackoffPolicy.Provider(),
-            fakeExecutor.getScheduledExecutorService(),
-            syncContext
-        ),
-        syncContext);
+    return (RetryingNameResolver) RetryingNameResolver.wrap(dnsResolver, args);
   }
 
   @Before

core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java

@@ -3844,8 +3844,6 @@ public double nextDouble() {
     verify(mockLoadBalancer).acceptResolvedAddresses(resolvedAddressCaptor.capture());
     ResolvedAddresses resolvedAddresses = resolvedAddressCaptor.getValue();
     assertThat(resolvedAddresses.getAddresses()).isEqualTo(nameResolverFactory.servers);
-    assertThat(resolvedAddresses.getAttributes()
-        .get(RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY)).isNotNull();
 
     // simulating request connection and then transport ready after resolved address
     Subchannel subchannel =
@@ -3951,8 +3949,6 @@ public void hedgingScheduledThenChannelShutdown_hedgeShouldStillHappen_newCallSh
     verify(mockLoadBalancer).acceptResolvedAddresses(resolvedAddressCaptor.capture());
     ResolvedAddresses resolvedAddresses = resolvedAddressCaptor.getValue();
     assertThat(resolvedAddresses.getAddresses()).isEqualTo(nameResolverFactory.servers);
-    assertThat(resolvedAddresses.getAttributes()
-        .get(RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY)).isNotNull();
 
     // simulating request connection and then transport ready after resolved address
     Subchannel subchannel =

core/src/test/java/io/grpc/internal/RetryingNameResolverTest.java

@@ -17,7 +17,6 @@
 package io.grpc.internal;
 
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.isA;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -28,7 +27,6 @@
 import io.grpc.NameResolver.ResolutionResult;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext;
-import io.grpc.internal.RetryingNameResolver.ResolutionResultListener;
 import java.lang.Thread.UncaughtExceptionHandler;
 import org.junit.Before;
 import org.junit.Rule;
@@ -58,8 +56,6 @@ public class RetryingNameResolverTest {
   private RetryScheduler mockRetryScheduler;
   @Captor
   private ArgumentCaptor<Listener2> listenerCaptor;
-  @Captor
-  private ArgumentCaptor<ResolutionResult> onResultCaptor;
   private final SynchronizationContext syncContext = new SynchronizationContext(
       mock(UncaughtExceptionHandler.class));
 
@@ -77,21 +73,14 @@ public void startAndShutdown() {
     retryingNameResolver.shutdown();
   }
 
-  // Make sure the ResolutionResultListener callback is added to the ResolutionResult attributes,
-  // and the retry scheduler is reset since the name resolution was successful.
   @Test
   public void onResult_success() {
+    when(mockListener.onResult2(isA(ResolutionResult.class))).thenReturn(Status.OK);
     retryingNameResolver.start(mockListener);
     verify(mockNameResolver).start(listenerCaptor.capture());
 
     listenerCaptor.getValue().onResult(ResolutionResult.newBuilder().build());
-    verify(mockListener).onResult(onResultCaptor.capture());
-    ResolutionResultListener resolutionResultListener = onResultCaptor.getValue()
-        .getAttributes()
-        .get(RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY);
-    assertThat(resolutionResultListener).isNotNull();
 
-    resolutionResultListener.resolutionAttempted(Status.OK);
     verify(mockRetryScheduler).reset();
   }
 
@@ -107,21 +96,15 @@ public void onResult2_sucesss() {
     verify(mockRetryScheduler).reset();
   }
 
-  // Make sure the ResolutionResultListener callback is added to the ResolutionResult attributes,
-  // and that a retry gets scheduled when the resolution results are rejected.
+  // Make sure that a retry gets scheduled when the resolution results are rejected.
   @Test
   public void onResult_failure() {
+    when(mockListener.onResult2(isA(ResolutionResult.class))).thenReturn(Status.UNAVAILABLE);
     retryingNameResolver.start(mockListener);
     verify(mockNameResolver).start(listenerCaptor.capture());
 
     listenerCaptor.getValue().onResult(ResolutionResult.newBuilder().build());
-    verify(mockListener).onResult(onResultCaptor.capture());
-    ResolutionResultListener resolutionResultListener = onResultCaptor.getValue()
-        .getAttributes()
-        .get(RetryingNameResolver.RESOLUTION_RESULT_LISTENER_KEY);
-    assertThat(resolutionResultListener).isNotNull();
 
-    resolutionResultListener.resolutionAttempted(Status.UNAVAILABLE);
     verify(mockRetryScheduler).schedule(isA(Runnable.class));
   }
 
@@ -138,24 +121,6 @@ public void onResult2_failure() {
     verify(mockRetryScheduler).schedule(isA(Runnable.class));
   }
 
-  // Wrapping a NameResolver more than once is a misconfiguration.
-  @Test
-  public void onResult_failure_doubleWrapped() {
-    NameResolver doubleWrappedResolver = new RetryingNameResolver(retryingNameResolver,
-        mockRetryScheduler, syncContext);
-
-    doubleWrappedResolver.start(mockListener);
-    verify(mockNameResolver).start(listenerCaptor.capture());
-
-    try {
-      listenerCaptor.getValue().onResult(ResolutionResult.newBuilder().build());
-    } catch (IllegalStateException e) {
-      assertThat(e).hasMessageThat().contains("can only be used once");
-      return;
-    }
-    fail("An exception should have been thrown for a double wrapped NAmeResolver");
-  }
-
   // A retry should get scheduled when name resolution fails.
   @Test
   public void onError() {
@@ -165,4 +130,4 @@ public void onError() {
     verify(mockListener).onError(Status.DEADLINE_EXCEEDED);
     verify(mockRetryScheduler).schedule(isA(Runnable.class));
   }
-}
+}