Trigger R8's ServiceLoader optimization

This simplifies R8 Full Mode's configuration when paired with R8
optimizations (which is made more difficult to avoid in AGP 9), as when
the optimization is triggered Full Mode will automatically keep the
constructor for the relevant classes.

android-interop-test doesn't currently enable R8 optimizations, so it
doesn't actually demonstrate the benefit, but I have manually confirmed
that enabling proguard-android-optimize.txt does cause the ServiceLoader
optimization in android-interop-test. Note that full mode is a separate
configuration and not necessary to get the ServiceLoader optimization.

Author: ejona86

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

@@ -17,7 +17,9 @@
 package io.grpc;
 
 import com.google.common.annotations.VisibleForTesting;
+import java.util.Iterator;
 import java.util.List;
+import java.util.ServiceLoader;
 
 @Internal
 public final class InternalServiceProviders {
@@ -27,12 +29,28 @@ private InternalServiceProviders() {
   /**
    * Accessor for method.
    */
+  @Deprecated
   public static <T> List<T> loadAll(
       Class<T> klass,
       Iterable<Class<?>> hardCodedClasses,
       ClassLoader classLoader,
       PriorityAccessor<T> priorityAccessor) {
-    return ServiceProviders.loadAll(klass, hardCodedClasses, classLoader, priorityAccessor);
+    return loadAll(
+        klass,
+        ServiceLoader.load(klass, classLoader).iterator(),
+        () -> hardCodedClasses,
+        priorityAccessor);
+  }
+
+  /**
+   * Accessor for method.
+   */
+  public static <T> List<T> loadAll(
+      Class<T> klass,
+      Iterator<T> serviceLoader,
+      Supplier<Iterable<Class<?>>> hardCodedClasses,
+      PriorityAccessor<T> priorityAccessor) {
+    return ServiceProviders.loadAll(klass, serviceLoader, hardCodedClasses::get, priorityAccessor);
   }
 
   /**
@@ -60,4 +78,8 @@ public static boolean isAndroid(ClassLoader cl) {
   }
 
   public interface PriorityAccessor<T> extends ServiceProviders.PriorityAccessor<T> {}
+
+  public interface Supplier<T> {
+    T get();
+  }
 }

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

@@ -26,6 +26,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -42,7 +43,6 @@
 public final class LoadBalancerRegistry {
   private static final Logger logger = Logger.getLogger(LoadBalancerRegistry.class.getName());
   private static LoadBalancerRegistry instance;
-  private static final Iterable<Class<?>> HARDCODED_CLASSES = getHardCodedClasses();
 
   private final LinkedHashSet<LoadBalancerProvider> allProviders =
       new LinkedHashSet<>();
@@ -101,8 +101,10 @@ public static synchronized LoadBalancerRegistry getDefaultRegistry() {
     if (instance == null) {
       List<LoadBalancerProvider> providerList = ServiceProviders.loadAll(
           LoadBalancerProvider.class,
-          HARDCODED_CLASSES,
-          LoadBalancerProvider.class.getClassLoader(),
+          ServiceLoader
+            .load(LoadBalancerProvider.class, LoadBalancerProvider.class.getClassLoader())
+            .iterator(),
+          LoadBalancerRegistry::getHardCodedClasses,
           new LoadBalancerPriorityAccessor());
       instance = new LoadBalancerRegistry();
       for (LoadBalancerProvider provider : providerList) {

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

@@ -29,6 +29,7 @@
 import java.util.Comparator;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.ServiceLoader;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.concurrent.ThreadSafe;
@@ -100,8 +101,10 @@ public static synchronized ManagedChannelRegistry getDefaultRegistry() {
     if (instance == null) {
       List<ManagedChannelProvider> providerList = ServiceProviders.loadAll(
           ManagedChannelProvider.class,
-          getHardCodedClasses(),
-          ManagedChannelProvider.class.getClassLoader(),
+          ServiceLoader
+            .load(ManagedChannelProvider.class, ManagedChannelProvider.class.getClassLoader())
+            .iterator(),
+          ManagedChannelRegistry::getHardCodedClasses,
           new ManagedChannelPriorityAccessor());
       instance = new ManagedChannelRegistry();
       for (ManagedChannelProvider provider : providerList) {

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

@@ -29,6 +29,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.ServiceLoader;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -125,8 +126,10 @@ public static synchronized NameResolverRegistry getDefaultRegistry() {
     if (instance == null) {
       List<NameResolverProvider> providerList = ServiceProviders.loadAll(
           NameResolverProvider.class,
-          getHardCodedClasses(),
-          NameResolverProvider.class.getClassLoader(),
+          ServiceLoader
+            .load(NameResolverProvider.class, NameResolverProvider.class.getClassLoader())
+            .iterator(),
+          NameResolverRegistry::getHardCodedClasses,
           new NameResolverPriorityAccessor());
       if (providerList.isEmpty()) {
         logger.warning("No NameResolverProviders found via ServiceLoader, including for DNS. This "

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

@@ -24,6 +24,7 @@
 import java.util.Comparator;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.ServiceLoader;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.concurrent.ThreadSafe;
@@ -93,8 +94,9 @@ public static synchronized ServerRegistry getDefaultRegistry() {
     if (instance == null) {
       List<ServerProvider> providerList = ServiceProviders.loadAll(
           ServerProvider.class,
-          getHardCodedClasses(),
-          ServerProvider.class.getClassLoader(),
+          ServiceLoader.load(ServerProvider.class, ServerProvider.class.getClassLoader())
+            .iterator(),
+          ServerRegistry::getHardCodedClasses,
           new ServerPriorityAccessor());
       instance = new ServerRegistry();
       for (ServerProvider provider : providerList) {

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

@@ -17,10 +17,13 @@
 package io.grpc;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Supplier;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.Iterator;
 import java.util.List;
+import java.util.ListIterator;
 import java.util.ServiceConfigurationError;
 import java.util.ServiceLoader;
 
@@ -34,20 +37,39 @@ private ServiceProviders() {
    * {@link ServiceLoader}.
    * If this is Android, returns all available implementations in {@code hardcoded}.
    * The list is sorted in descending priority order.
+   *
+   * <p>{@code serviceLoader} should be created with {@code ServiceLoader.load(MyClass.class,
+   * MyClass.class.getClassLoader()).iterator()} in order to be detected by R8 so that R8 full mode
+   * will keep the constructors for the provider classes.
    */
   public static <T> List<T> loadAll(
       Class<T> klass,
-      Iterable<Class<?>> hardcoded,
-      ClassLoader cl,
+      Iterator<T> serviceLoader,
+      Supplier<Iterable<Class<?>>> hardcoded,
       final PriorityAccessor<T> priorityAccessor) {
-    Iterable<T> candidates;
-    if (isAndroid(cl)) {
-      candidates = getCandidatesViaHardCoded(klass, hardcoded);
+    Iterator<T> candidates;
+    if (serviceLoader instanceof ListIterator) {
+      // A rewriting tool has replaced the ServiceLoader with a List of some sort (R8 uses
+      // ArrayList, AppReduce uses singletonList). We prefer to use such iterators on Android as
+      // they won't need reflection like the hard-coded list does. In addition, the provider
+      // instances will have already been created, so it seems we should use them.
+      //
+      // R8: https://r8.googlesource.com/r8/+/490bc53d9310d4cc2a5084c05df4aadaec8c885d/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
+      // AppReduce: service_loader_pass.cc
+      candidates = serviceLoader;
+    } else if (isAndroid(klass.getClassLoader())) {
+      // Avoid getResource() on Android, which must read from a zip which uses a lot of memory
+      candidates = getCandidatesViaHardCoded(klass, hardcoded.get()).iterator();
+    } else if (!serviceLoader.hasNext()) {
+      // Attempt to load using the context class loader and ServiceLoader.
+      // This allows frameworks like http://aries.apache.org/modules/spi-fly.html to plug in.
+      candidates = ServiceLoader.load(klass).iterator();
     } else {
-      candidates = getCandidatesViaServiceLoader(klass, cl);
+      candidates = serviceLoader;
     }
     List<T> list = new ArrayList<>();
-    for (T current: candidates) {
+    while (candidates.hasNext()) {
+      T current = candidates.next();
       if (!priorityAccessor.isAvailable(current)) {
         continue;
       }
@@ -84,15 +106,14 @@ static boolean isAndroid(ClassLoader cl) {
   }
 
   /**
-   * Loads service providers for the {@code klass} service using {@link ServiceLoader}.
+   * For testing only: Loads service providers for the {@code klass} service using {@link
+   * ServiceLoader}. Does not support spi-fly and related tricks.
    */
   @VisibleForTesting
   public static <T> Iterable<T> getCandidatesViaServiceLoader(Class<T> klass, ClassLoader cl) {
     Iterable<T> i = ServiceLoader.load(klass, cl);
-    // Attempt to load using the context class loader and ServiceLoader.
-    // This allows frameworks like http://aries.apache.org/modules/spi-fly.html to plug in.
     if (!i.iterator().hasNext()) {
-      i = ServiceLoader.load(klass);
+      return null;
     }
     return i;
   }

api/src/test/java/io/grpc/ServiceProvidersTest.java

@@ -16,27 +16,30 @@
 
 package io.grpc;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
 import io.grpc.InternalServiceProviders.PriorityAccessor;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.ServiceConfigurationError;
+import java.util.ServiceLoader;
+import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 /** Unit tests for {@link ServiceProviders}. */
 @RunWith(JUnit4.class)
 public class ServiceProvidersTest {
-  private static final List<Class<?>> NO_HARDCODED = Collections.emptyList();
   private static final PriorityAccessor<ServiceProvidersTestAbstractProvider> ACCESSOR =
       new PriorityAccessor<ServiceProvidersTestAbstractProvider>() {
         @Override
@@ -51,6 +54,19 @@ public int getPriority(ServiceProvidersTestAbstractProvider provider) {
       };
   private final String serviceFile =
       "META-INF/services/io.grpc.ServiceProvidersTestAbstractProvider";
+  private boolean failingHardCodedAccessed;
+  private final Supplier<Iterable<Class<?>>> failingHardCoded = new Supplier<Iterable<Class<?>>>() {
+    @Override
+    public Iterable<Class<?>> get() {
+      failingHardCodedAccessed = true;
+      throw new AssertionError();
+    }
+  };
+
+  @After
+  public void tearDown() {
+    assertThat(failingHardCodedAccessed).isFalse();
+  }
 
   @Test
   public void contextClassLoaderProvider() {
@@ -69,7 +85,8 @@ public void contextClassLoaderProvider() {
       Thread.currentThread().setContextClassLoader(rcll);
       assertEquals(
           Available7Provider.class,
-          load(ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR).getClass());
+          load(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR)
+            .getClass());
     } finally {
       Thread.currentThread().setContextClassLoader(ccl);
     }
@@ -84,7 +101,7 @@ public void noProvider() {
           serviceFile,
           "io/grpc/ServiceProvidersTestAbstractProvider-doesNotExist.txt");
       Thread.currentThread().setContextClassLoader(cl);
-      assertNull(load(ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR));
+      assertNull(load(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR));
     } finally {
       Thread.currentThread().setContextClassLoader(ccl);
     }
@@ -96,10 +113,11 @@ public void multipleProvider() throws Exception {
         "io/grpc/ServiceProvidersTestAbstractProvider-multipleProvider.txt");
     assertSame(
         Available7Provider.class,
-        load(ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR).getClass());
+        load(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR)
+          .getClass());
 
-    List<ServiceProvidersTestAbstractProvider> providers = ServiceProviders.loadAll(
-        ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR);
+    List<ServiceProvidersTestAbstractProvider> providers = loadAll(
+        ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR);
     assertEquals(3, providers.size());
     assertEquals(Available7Provider.class, providers.get(0).getClass());
     assertEquals(Available5Provider.class, providers.get(1).getClass());
@@ -113,16 +131,16 @@ public void unavailableProvider() {
         "io/grpc/ServiceProvidersTestAbstractProvider-unavailableProvider.txt");
     assertEquals(
         Available7Provider.class,
-        load(ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR).getClass());
+        load(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR)
+          .getClass());
   }
 
   @Test
   public void unknownClassProvider() {
     ClassLoader cl = new ReplacingClassLoader(getClass().getClassLoader(), serviceFile,
         "io/grpc/ServiceProvidersTestAbstractProvider-unknownClassProvider.txt");
     try {
-      ServiceProviders.loadAll(
-          ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR);
+      loadAll(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR);
       fail("Exception expected");
     } catch (ServiceConfigurationError e) {
       // noop
@@ -136,8 +154,7 @@ public void exceptionSurfacedToCaller_failAtInit() {
     try {
       // Even though there is a working provider, if any providers fail then we should fail
       // completely to avoid returning something unexpected.
-      ServiceProviders.loadAll(
-          ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR);
+      loadAll(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR);
       fail("Expected exception");
     } catch (ServiceConfigurationError expected) {
       // noop
@@ -150,8 +167,7 @@ public void exceptionSurfacedToCaller_failAtPriority() {
         "io/grpc/ServiceProvidersTestAbstractProvider-failAtPriorityProvider.txt");
     try {
       // The exception should be surfaced to the caller
-      ServiceProviders.loadAll(
-          ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR);
+      loadAll(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR);
       fail("Expected exception");
     } catch (FailAtPriorityProvider.PriorityException expected) {
       // noop
@@ -164,8 +180,7 @@ public void exceptionSurfacedToCaller_failAtAvailable() {
         "io/grpc/ServiceProvidersTestAbstractProvider-failAtAvailableProvider.txt");
     try {
       // The exception should be surfaced to the caller
-      ServiceProviders.loadAll(
-          ServiceProvidersTestAbstractProvider.class, NO_HARDCODED, cl, ACCESSOR);
+      loadAll(ServiceProvidersTestAbstractProvider.class, failingHardCoded, cl, ACCESSOR);
       fail("Expected exception");
     } catch (FailAtAvailableProvider.AvailableException expected) {
       // noop
@@ -242,16 +257,28 @@ class RandomClass {}
 
   private static <T> T load(
       Class<T> klass,
-      Iterable<Class<?>> hardcoded,
+      Supplier<Iterable<Class<?>>> hardCoded,
       ClassLoader cl,
       PriorityAccessor<T> priorityAccessor) {
-    List<T> candidates = ServiceProviders.loadAll(klass, hardcoded, cl, priorityAccessor);
+    List<T> candidates = loadAll(klass, hardCoded, cl, priorityAccessor);
     if (candidates.isEmpty()) {
       return null;
     }
     return candidates.get(0);
   }
 
+  private static <T> List<T> loadAll(
+      Class<T> klass,
+      Supplier<Iterable<Class<?>>> hardCoded,
+      ClassLoader classLoader,
+      PriorityAccessor<T> priorityAccessor) {
+    return ServiceProviders.loadAll(
+        klass,
+        ServiceLoader.load(klass, classLoader).iterator(),
+        hardCoded,
+        priorityAccessor);
+  }
+
   private static class BaseProvider extends ServiceProvidersTestAbstractProvider {
     private final boolean isAvailable;
     private final int priority;