api: pass Subchannel state updates to SubchannelStateListener rather than LoadBalancer (take 2) (#5722)

This is a revised version of #5503 (62b03fd), which was rolled back in f8d0868. The newer version passes SubchannelStateListener to Subchannel.start() instead of SubchannelCreationArgs, which allows us to remove the Subchannel argument from the listener, which works as a solution for #5676.

LoadBalancers that call the old createSubchannel() will get start() implicitly called with a listener that passes updates to the deprecated LoadBalancer.handleSubchannelState(). Those who call the new createSubchannel() will have to call start() explicitly.

GRPCLB code is still using the old API, because it's a pain to migrate the SubchannelPool to the new API.  Since CachedSubchannelHelper is on the way, it's easier to switch to it when it's ready. Keeping
GRPCLB with the old API would also confirm the backward compatibility.

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

@@ -16,12 +16,14 @@
 
 package io.grpc;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -329,15 +331,22 @@ public abstract class LoadBalancer {
    * <p>SHUTDOWN can only happen in two cases.  One is that LoadBalancer called {@link
    * Subchannel#shutdown} earlier, thus it should have already discarded this Subchannel.  The other
    * is that Channel is doing a {@link ManagedChannel#shutdownNow forced shutdown} or has already
-   * terminated, thus there won't be further requests to LoadBalancer.  Therefore, SHUTDOWN can be
+   * terminated, thus there won't be further requests to LoadBalancer.  Therefore, the LoadBalancer
-   * safely ignored.
+   * usually don't need to react to a SHUTDOWN state.
    *
    * @param subchannel the involved Subchannel
    * @param stateInfo the new state
    * @since 1.2.0
+   * @deprecated This method will be removed.  Stop overriding it.  Instead, pass {@link
+   *             SubchannelStateListener} to {@link Subchannel#start} to receive Subchannel state
+   *             updates
    */
-  public abstract void handleSubchannelState(
+  @Deprecated
-      Subchannel subchannel, ConnectivityStateInfo stateInfo);
+  public void handleSubchannelState(
+      Subchannel subchannel, ConnectivityStateInfo stateInfo) {
+    // Do nothing.  If the implemetation doesn't implement this, it will get subchannel states from
+    // the new API.  We don't throw because there may be forwarding LoadBalancers still plumb this.
+  }
 
   /**
    * The channel asks the load-balancer to shutdown.  No more callbacks will be called after this
@@ -539,7 +548,7 @@ public abstract class LoadBalancer {
      *       {@code handleSubchannelState}'s javadoc for more details.</li>
      * </ol>
      *
-     * @param subchannel the picked Subchannel
+     * @param subchannel the picked Subchannel.  It must have been {@link Subchannel#start started}
      * @param streamTracerFactory if not null, will be used to trace the activities of the stream
      *                            created as a result of this pick. Note it's possible that no
      *                            stream is created at all in some cases.
@@ -666,6 +675,218 @@ public abstract class LoadBalancer {
     }
   }
 
+  /**
+   * Arguments for creating a {@link Subchannel}.
+   *
+   * @since 1.22.0
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1771")
+  public static final class CreateSubchannelArgs {
+    private final List<EquivalentAddressGroup> addrs;
+    private final Attributes attrs;
+    private final Object[][] customOptions;
+
+    private CreateSubchannelArgs(
+        List<EquivalentAddressGroup> addrs, Attributes attrs, Object[][] customOptions) {
+      this.addrs = checkNotNull(addrs, "addresses are not set");
+      this.attrs = checkNotNull(attrs, "attrs");
+      this.customOptions = checkNotNull(customOptions, "customOptions");
+    }
+
+    /**
+     * Returns the addresses, which is an unmodifiable list.
+     */
+    public List<EquivalentAddressGroup> getAddresses() {
+      return addrs;
+    }
+
+    /**
+     * Returns the attributes.
+     */
+    public Attributes getAttributes() {
+      return attrs;
+    }
+
+    /**
+     * Get the value for a custom option or its inherent default.
+     *
+     * @param key Key identifying option
+     */
+    @SuppressWarnings("unchecked")
+    public <T> T getOption(Key<T> key) {
+      Preconditions.checkNotNull(key, "key");
+      for (int i = 0; i < customOptions.length; i++) {
+        if (key.equals(customOptions[i][0])) {
+          return (T) customOptions[i][1];
+        }
+      }
+      return key.defaultValue;
+    }
+
+    /**
+     * Returns a builder with the same initial values as this object.
+     */
+    public Builder toBuilder() {
+      return newBuilder().setAddresses(addrs).setAttributes(attrs).copyCustomOptions(customOptions);
+    }
+
+    /**
+     * Creates a new builder.
+     */
+    public static Builder newBuilder() {
+      return new Builder();
+    }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper(this)
+          .add("addrs", addrs)
+          .add("attrs", attrs)
+          .add("customOptions", Arrays.deepToString(customOptions))
+          .toString();
+    }
+
+    @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1771")
+    public static final class Builder {
+
+      private List<EquivalentAddressGroup> addrs;
+      private Attributes attrs = Attributes.EMPTY;
+      private Object[][] customOptions = new Object[0][2];
+
+      Builder() {
+      }
+
+      private <T> Builder copyCustomOptions(Object[][] options) {
+        customOptions = new Object[options.length][2];
+        System.arraycopy(options, 0, customOptions, 0, options.length);
+        return this;
+      }
+
+      /**
+       * Add a custom option. Any existing value for the key is overwritten.
+       *
+       * <p>This is an <strong>optional</strong> property.
+       *
+       * @param key the option key
+       * @param value the option value
+       */
+      public <T> Builder addOption(Key<T> key, T value) {
+        Preconditions.checkNotNull(key, "key");
+        Preconditions.checkNotNull(value, "value");
+
+        int existingIdx = -1;
+        for (int i = 0; i < customOptions.length; i++) {
+          if (key.equals(customOptions[i][0])) {
+            existingIdx = i;
+            break;
+          }
+        }
+
+        if (existingIdx == -1) {
+          Object[][] newCustomOptions = new Object[customOptions.length + 1][2];
+          System.arraycopy(customOptions, 0, newCustomOptions, 0, customOptions.length);
+          customOptions = newCustomOptions;
+          existingIdx = customOptions.length - 1;
+        }
+        customOptions[existingIdx] = new Object[]{key, value};
+        return this;
+      }
+
+      /**
+       * The addresses to connect to.  All addresses are considered equivalent and will be tried
+       * in the order they are provided.
+       */
+      public Builder setAddresses(EquivalentAddressGroup addrs) {
+        this.addrs = Collections.singletonList(addrs);
+        return this;
+      }
+
+      /**
+       * The addresses to connect to.  All addresses are considered equivalent and will
+       * be tried in the order they are provided.
+       *
+       * <p>This is a <strong>required</strong> property.
+       *
+       * @throws IllegalArgumentException if {@code addrs} is empty
+       */
+      public Builder setAddresses(List<EquivalentAddressGroup> addrs) {
+        checkArgument(!addrs.isEmpty(), "addrs is empty");
+        this.addrs = Collections.unmodifiableList(new ArrayList<>(addrs));
+        return this;
+      }
+
+      /**
+       * Attributes provided here will be included in {@link Subchannel#getAttributes}.
+       *
+       * <p>This is an <strong>optional</strong> property.  Default is empty if not set.
+       */
+      public Builder setAttributes(Attributes attrs) {
+        this.attrs = checkNotNull(attrs, "attrs");
+        return this;
+      }
+
+      /**
+       * Creates a new args object.
+       */
+      public CreateSubchannelArgs build() {
+        return new CreateSubchannelArgs(addrs, attrs, customOptions);
+      }
+    }
+
+    /**
+     * Key for a key-value pair. Uses reference equality.
+     */
+    @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1771")
+    public static final class Key<T> {
+
+      private final String debugString;
+      private final T defaultValue;
+
+      private Key(String debugString, T defaultValue) {
+        this.debugString = debugString;
+        this.defaultValue = defaultValue;
+      }
+
+      /**
+       * Factory method for creating instances of {@link Key}. The default value of the key is
+       * {@code null}.
+       *
+       * @param debugString a debug string that describes this key.
+       * @param <T> Key type
+       * @return Key object
+       */
+      public static <T> Key<T> create(String debugString) {
+        Preconditions.checkNotNull(debugString, "debugString");
+        return new Key<>(debugString, /*defaultValue=*/ null);
+      }
+
+      /**
+       * Factory method for creating instances of {@link Key}.
+       *
+       * @param debugString a debug string that describes this key.
+       * @param defaultValue default value to return when value for key not set
+       * @param <T> Key type
+       * @return Key object
+       */
+      public static <T> Key<T> createWithDefault(String debugString, T defaultValue) {
+        Preconditions.checkNotNull(debugString, "debugString");
+        return new Key<>(debugString, defaultValue);
+      }
+
+      /**
+       * Returns the user supplied default value for this key.
+       */
+      public T getDefault() {
+        return defaultValue;
+      }
+
+      @Override
+      public String toString() {
+        return debugString;
+      }
+    }
+  }
+
   /**
    * Provides essentials for LoadBalancer implementations.
    *
@@ -679,7 +900,11 @@ public abstract class LoadBalancer {
      * EquivalentAddressGroup}.
      *
      * @since 1.2.0
+     * @deprecated Use {@link #createSubchannel(io.grpc.LoadBalancer.CreateSubchannelArgs)}
+     *             instead. Note the new API must be called from {@link #getSynchronizationContext
+     *             the Synchronization Context}.
      */
+    @Deprecated
     public final Subchannel createSubchannel(EquivalentAddressGroup addrs, Attributes attrs) {
       checkNotNull(addrs, "addrs");
       return createSubchannel(Collections.singletonList(addrs), attrs);
@@ -700,11 +925,32 @@ public abstract class LoadBalancer {
      *
      * @throws IllegalArgumentException if {@code addrs} is empty
      * @since 1.14.0
+     * @deprecated Use {@link #createSubchannel(io.grpc.LoadBalancer.CreateSubchannelArgs)}
+     *             instead. Note the new API must be called from {@link #getSynchronizationContext
+     *             the Synchronization Context}.
      */
+    @Deprecated
     public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
       throw new UnsupportedOperationException();
     }
 
+    /**
+     * Creates a Subchannel, which is a logical connection to the given group of addresses which are
+     * considered equivalent.  The {@code attrs} are custom attributes associated with this
+     * Subchannel, and can be accessed later through {@link Subchannel#getAttributes
+     * Subchannel.getAttributes()}.
+     *
+     * <p>The LoadBalancer is responsible for closing unused Subchannels, and closing all
+     * Subchannels within {@link #shutdown}.
+     *
+     * <p>It must be called from {@link #getSynchronizationContext the Synchronization Context}
+     *
+     * @since 1.22.0
+     */
+    public Subchannel createSubchannel(CreateSubchannelArgs args) {
+      throw new UnsupportedOperationException();
+    }
+
     /**
      * Equivalent to {@link #updateSubchannelAddresses(io.grpc.LoadBalancer.Subchannel, List)} with
      * the given single {@code EquivalentAddressGroup}.
@@ -925,14 +1171,35 @@ public abstract class LoadBalancer {
    * #requestConnection requestConnection()} can be used to ask Subchannel to create a transport if
    * there isn't any.
    *
+   * <p>{@link #start} must be called prior to calling any other methods, with the exception of
+   * {@link #shutdown}, which can be called at any time.
+   *
    * @since 1.2.0
    */
   @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1771")
   public abstract static class Subchannel {
+    /**
+     * Starts the Subchannel.  Can only be called once.
+     *
+     * <p>Must be called prior to any other method on this class, except for {@link #shutdown} which
+     * may be called at any time.
+     *
+     * <p>Must be called from the {@link Helper#getSynchronizationContext Synchronization Context},
+     * otherwise it may throw.  See <a href="https://github.com/grpc/grpc-java/issues/5015">
+     * #5015</a> for more discussions.
+     *
+     * @param listener receives state updates for this Subchannel.
+     */
+    public void start(SubchannelStateListener listener) {
+      throw new UnsupportedOperationException("Not implemented");
+    }
+
     /**
      * Shuts down the Subchannel.  After this method is called, this Subchannel should no longer
      * be returned by the latest {@link SubchannelPicker picker}, and can be safely discarded.
      *
+     * <p>Calling it on an already shut-down Subchannel has no effect.
+     *
      * <p>It should be called from the Synchronization Context.  Currently will log a warning if
      * violated.  It will become an exception eventually.  See <a
      * href="https://github.com/grpc/grpc-java/issues/5015">#5015</a> for the background.
@@ -966,7 +1233,7 @@ public abstract class LoadBalancer {
      */
     public final EquivalentAddressGroup getAddresses() {
       List<EquivalentAddressGroup> groups = getAllAddresses();
-      Preconditions.checkState(groups.size() == 1, "Does not have exactly one group");
+      Preconditions.checkState(groups.size() == 1, "%s does not have exactly one group", groups);
       return groups.get(0);
     }
 
@@ -1031,6 +1298,36 @@ public abstract class LoadBalancer {
     }
   }
 
+  /**
+   * Receives state changes for one {@link Subchannel}. All methods are run under {@link
+   * Helper#getSynchronizationContext}.
+   *
+   * @since 1.22.0
+   */
+  public interface SubchannelStateListener {
+    /**
+     * Handles a state change on a Subchannel.
+     *
+     * <p>The initial state of a Subchannel is IDLE. You won't get a notification for the initial
+     * IDLE state.
+     *
+     * <p>If the new state is not SHUTDOWN, this method should create a new picker and call {@link
+     * Helper#updateBalancingState Helper.updateBalancingState()}.  Failing to do so may result in
+     * unnecessary delays of RPCs. Please refer to {@link PickResult#withSubchannel
+     * PickResult.withSubchannel()}'s javadoc for more information.
+     *
+     * <p>SHUTDOWN can only happen in two cases.  One is that LoadBalancer called {@link
+     * Subchannel#shutdown} earlier, thus it should have already discarded this Subchannel.  The
+     * other is that Channel is doing a {@link ManagedChannel#shutdownNow forced shutdown} or has
+     * already terminated, thus there won't be further requests to LoadBalancer.  Therefore, the
+     * LoadBalancer usually don't need to react to a SHUTDOWN state.
+     * @param newState the new state
+     *
+     * @since 1.22.0
+     */
+    void onSubchannelState(ConnectivityStateInfo newState);
+  }
+
   /**
    * Factory to create {@link LoadBalancer} instance.
    *

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

@@ -97,7 +97,10 @@ public final class ForwardingTestUtil {
       try {
         method.invoke(verify(mockDelegate), args);
       } catch (InvocationTargetException e) {
-        throw new AssertionError(String.format("Method was not forwarded: %s", method));
+        AssertionError ae =
+            new AssertionError(String.format("Method was not forwarded: %s", method));
+        ae.initCause(e);
+        throw ae;
       }
     }
 
@@ -126,7 +129,7 @@ public final class ForwardingTestUtil {
      * method once.  It is recommended that each invocation returns a distinctive object for the
      * same type, in order to verify that arguments are passed by the tested class correctly.
      *
-     * @return a value to be passed as an argument.  If {@code null}, {@link Default#defaultValue}
+     * @return a value to be passed as an argument.  If {@code null}, {@link Defaults#defaultValue}
      *         will be used.
      */
     @Nullable Object get(Method method, int argPos, Class<?> clazz);

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

@@ -17,11 +17,14 @@
 package io.grpc;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import java.net.SocketAddress;
 import java.util.Arrays;
 import java.util.List;
@@ -35,6 +38,8 @@ import org.junit.runners.JUnit4;
 public class LoadBalancerTest {
   private final Subchannel subchannel = mock(Subchannel.class);
   private final Subchannel subchannel2 = mock(Subchannel.class);
+  private final SubchannelStateListener subchannelStateListener =
+      mock(SubchannelStateListener.class);
   private final ClientStreamTracer.Factory tracerFactory = mock(ClientStreamTracer.Factory.class);
   private final Status status = Status.UNAVAILABLE.withDescription("for test");
   private final Status status2 = Status.UNAVAILABLE.withDescription("for test 2");
@@ -120,8 +125,9 @@ public class LoadBalancerTest {
     assertThat(error1).isNotEqualTo(drop1);
   }
 
+  @Deprecated
   @Test
-  public void helper_createSubchannel_delegates() {
+  public void helper_createSubchannel_old_delegates() {
     class OverrideCreateSubchannel extends NoopHelper {
       boolean ran;
 
@@ -140,9 +146,28 @@ public class LoadBalancerTest {
     assertThat(helper.ran).isTrue();
   }
 
-  @Test(expected = UnsupportedOperationException.class)
+  @Test
-  public void helper_createSubchannelList_throws() {
+  @SuppressWarnings("deprecation")
+  public void helper_createSubchannelList_oldApi_throws() {
+    try {
       new NoopHelper().createSubchannel(Arrays.asList(eag), attrs);
+      fail("Should throw");
+    } catch (UnsupportedOperationException e) {
+      // expected
+    }
+  }
+
+  @Test
+  public void helper_createSubchannelList_throws() {
+    try {
+      new NoopHelper().createSubchannel(CreateSubchannelArgs.newBuilder()
+          .setAddresses(eag)
+          .setAttributes(attrs)
+          .build());
+      fail("Should throw");
+    } catch (UnsupportedOperationException e) {
+      // expected
+    }
   }
 
   @Test
@@ -199,6 +224,75 @@ public class LoadBalancerTest {
     }.getAddresses();
   }
 
+  @Test
+  public void createSubchannelArgs_option_keyOps() {
+    CreateSubchannelArgs.Key<String> testKey = CreateSubchannelArgs.Key.create("test-key");
+    String testValue = "test-value";
+    CreateSubchannelArgs.Key<String> testWithDefaultKey = CreateSubchannelArgs.Key
+        .createWithDefault("test-key", testValue);
+    CreateSubchannelArgs args = CreateSubchannelArgs.newBuilder()
+        .setAddresses(eag)
+        .setAttributes(attrs)
+        .build();
+    assertThat(args.getOption(testKey)).isNull();
+    assertThat(args.getOption(testWithDefaultKey)).isSameInstanceAs(testValue);
+  }
+
+  @Test
+  public void createSubchannelArgs_option_addGet() {
+    String testValue = "test-value";
+    CreateSubchannelArgs.Key<String> testKey = CreateSubchannelArgs.Key.create("test-key");
+    CreateSubchannelArgs args = CreateSubchannelArgs.newBuilder()
+        .setAddresses(eag)
+        .setAttributes(attrs)
+        .addOption(testKey, testValue)
+        .build();
+    assertThat(args.getOption(testKey)).isEqualTo(testValue);
+  }
+
+  @Test
+  public void createSubchannelArgs_option_lastOneWins() {
+    String testValue1 = "test-value-1";
+    String testValue2 = "test-value-2";
+    CreateSubchannelArgs.Key<String> testKey = CreateSubchannelArgs.Key.create("test-key");
+    CreateSubchannelArgs args = CreateSubchannelArgs.newBuilder()
+        .setAddresses(eag)
+        .setAttributes(attrs)
+        .addOption(testKey, testValue1)
+        .addOption(testKey, testValue2)
+        .build();
+    assertThat(args.getOption(testKey)).isEqualTo(testValue2);
+  }
+
+  @Test
+  public void createSubchannelArgs_build() {
+    CreateSubchannelArgs.Key<Object> testKey = CreateSubchannelArgs.Key.create("test-key");
+    Object testValue = new Object();
+    CreateSubchannelArgs args = CreateSubchannelArgs.newBuilder()
+        .setAddresses(eag)
+        .setAttributes(attrs)
+        .addOption(testKey, testValue)
+        .build();
+    CreateSubchannelArgs rebuildedArgs = args.toBuilder().build();
+    assertThat(rebuildedArgs.getAddresses()).containsExactly(eag);
+    assertThat(rebuildedArgs.getAttributes()).isSameInstanceAs(attrs);
+    assertThat(rebuildedArgs.getOption(testKey)).isSameInstanceAs(testValue);
+  }
+
+  @Test
+  public void createSubchannelArgs_toString() {
+    CreateSubchannelArgs.Key<String> testKey = CreateSubchannelArgs.Key.create("test-key");
+    CreateSubchannelArgs args = CreateSubchannelArgs.newBuilder()
+        .setAddresses(eag)
+        .setAttributes(attrs)
+        .addOption(testKey, "test-value")
+        .build();
+    String str = args.toString();
+    assertThat(str).contains("addrs=");
+    assertThat(str).contains("attrs=");
+    assertThat(str).contains("customOptions=");
+  }
+
   @Deprecated
   @Test
   public void handleResolvedAddressGroups_delegatesToHandleResolvedAddresses() {

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

@@ -77,9 +77,6 @@ public final class AutoConfiguredLoadBalancerFactory extends LoadBalancer.Factor
     @Override
     public void handleNameResolutionError(Status error) {}
 
-    @Override
-    public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {}
-
     @Override
     public void shutdown() {}
   }
@@ -165,6 +162,7 @@ public final class AutoConfiguredLoadBalancerFactory extends LoadBalancer.Factor
       getDelegate().handleNameResolutionError(error);
     }
 
+    @Deprecated
     @Override
     public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
       getDelegate().handleSubchannelState(subchannel, stateInfo);

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

@@ -53,10 +53,12 @@ import io.grpc.InternalInstrumented;
 import io.grpc.InternalLogId;
 import io.grpc.InternalWithLogId;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.ManagedChannel;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
@@ -228,8 +230,8 @@ final class ManagedChannelImpl extends ManagedChannel implements
   private final AtomicBoolean shutdown = new AtomicBoolean(false);
   // Must only be mutated and read from syncContext
   private boolean shutdownNowed;
-  // Must be mutated from syncContext
+  // Must only be mutated and read from syncContext
-  private volatile boolean terminating;
+  private boolean terminating;
   // Must be mutated from syncContext
   private volatile boolean terminated;
   private final CountDownLatch terminatedLatch = new CountDownLatch(1);
@@ -881,6 +883,13 @@ final class ManagedChannelImpl extends ManagedChannel implements
     }
   }
 
+  // Must be called from syncContext
+  private void handleInternalSubchannelState(ConnectivityStateInfo newState) {
+    if (newState.getState() == TRANSIENT_FAILURE || newState.getState() == IDLE) {
+      refreshAndResetNameResolution();
+    }
+  }
+
   @Override
   @SuppressWarnings("deprecation")
   public ConnectivityState getState(boolean requestConnection) {
@@ -1042,103 +1051,47 @@ final class ManagedChannelImpl extends ManagedChannel implements
   private class LbHelperImpl extends LoadBalancer.Helper {
     LoadBalancer lb;
 
-    // Must be called from syncContext
+    @Deprecated
-    private void handleInternalSubchannelState(ConnectivityStateInfo newState) {
-      if (newState.getState() == TRANSIENT_FAILURE || newState.getState() == IDLE) {
-        refreshAndResetNameResolution();
-      }
-    }
-
     @Override
     public AbstractSubchannel createSubchannel(
         List<EquivalentAddressGroup> addressGroups, Attributes attrs) {
       logWarningIfNotInSyncContext("createSubchannel()");
+      // TODO(ejona): can we be even stricter? Like loadBalancer == null?
       checkNotNull(addressGroups, "addressGroups");
       checkNotNull(attrs, "attrs");
-      // TODO(ejona): can we be even stricter? Like loadBalancer == null?
+      final AbstractSubchannel subchannel = createSubchannelInternal(
-      checkState(!terminated, "Channel is terminated");
+          CreateSubchannelArgs.newBuilder()
-      final SubchannelImpl subchannel = new SubchannelImpl(attrs);
+              .setAddresses(addressGroups)
-      long subchannelCreationTime = timeProvider.currentTimeNanos();
+              .setAttributes(attrs)
-      InternalLogId subchannelLogId = InternalLogId.allocate("Subchannel", /*details=*/ null);
+              .build());
-      ChannelTracer subchannelTracer =
-          new ChannelTracer(
-              subchannelLogId, maxTraceEvents, subchannelCreationTime,
-              "Subchannel for " + addressGroups);
 
-      final class ManagedInternalSubchannelCallback extends InternalSubchannel.Callback {
+      final SubchannelStateListener listener =
-        // All callbacks are run in syncContext
+          new LoadBalancer.SubchannelStateListener() {
             @Override
-        void onTerminated(InternalSubchannel is) {
+            public void onSubchannelState(ConnectivityStateInfo newState) {
-          subchannels.remove(is);
-          channelz.removeSubchannel(is);
-          maybeTerminateChannel();
-        }
-
-        @Override
-        void onStateChange(InternalSubchannel is, ConnectivityStateInfo newState) {
-          handleInternalSubchannelState(newState);
-          // Call LB only if it's not shutdown.  If LB is shutdown, lbHelper won't match.
-          if (LbHelperImpl.this == ManagedChannelImpl.this.lbHelper) {
               lb.handleSubchannelState(subchannel, newState);
             }
-        }
+          };
 
-        @Override
+      syncContext.execute(new Runnable() {
-        void onInUse(InternalSubchannel is) {
-          inUseStateAggregator.updateObjectInUse(is, true);
-        }
-
-        @Override
-        void onNotInUse(InternalSubchannel is) {
-          inUseStateAggregator.updateObjectInUse(is, false);
-        }
-      }
-
-      final InternalSubchannel internalSubchannel = new InternalSubchannel(
-          addressGroups,
-          authority(),
-          userAgent,
-          backoffPolicyProvider,
-          transportFactory,
-          transportFactory.getScheduledExecutorService(),
-          stopwatchSupplier,
-          syncContext,
-          new ManagedInternalSubchannelCallback(),
-          channelz,
-          callTracerFactory.create(),
-          subchannelTracer,
-          subchannelLogId,
-          timeProvider);
-      channelTracer.reportEvent(new ChannelTrace.Event.Builder()
-          .setDescription("Child Subchannel created")
-          .setSeverity(ChannelTrace.Event.Severity.CT_INFO)
-          .setTimestampNanos(subchannelCreationTime)
-          .setSubchannelRef(internalSubchannel)
-          .build());
-      channelz.addSubchannel(internalSubchannel);
-      subchannel.subchannel = internalSubchannel;
-
-      final class AddSubchannel implements Runnable {
           @Override
           public void run() {
-          if (terminating) {
+            subchannel.start(listener);
-            // Because SynchronizationContext doesn't guarantee the runnable has been executed upon
-            // when returning, the subchannel may still be returned to the balancer without being
-            // shutdown even if "terminating" is already true.  The subchannel will not be used in
-            // this case, because delayed transport has terminated when "terminating" becomes true,
-            // and no more requests will be sent to balancer beyond this point.
-            internalSubchannel.shutdown(SHUTDOWN_STATUS);
-          }
-          if (!terminated) {
-            // If channel has not terminated, it will track the subchannel and block termination
-            // for it.
-            subchannels.add(internalSubchannel);
-          }
           }
+        });
+      return subchannel;
     }
 
-      syncContext.execute(new AddSubchannel());
+    @Override
-      return subchannel;
+    public AbstractSubchannel createSubchannel(CreateSubchannelArgs args) {
+      syncContext.throwIfNotInThisSynchronizationContext();
+      return createSubchannelInternal(args);
+    }
+
+    private AbstractSubchannel createSubchannelInternal(CreateSubchannelArgs args) {
+      // TODO(ejona): can we be even stricter? Like loadBalancer == null?
+      checkState(!terminated, "Channel is terminated");
+      return new SubchannelImpl(args, this);
     }
 
     @Override
@@ -1452,46 +1405,143 @@ final class ManagedChannelImpl extends ManagedChannel implements
   }
 
   private final class SubchannelImpl extends AbstractSubchannel {
-    // Set right after SubchannelImpl is created.
+    final CreateSubchannelArgs args;
+    final LbHelperImpl helper;
+    SubchannelStateListener listener;
     InternalSubchannel subchannel;
-    final Object shutdownLock = new Object();
+    boolean started;
-    final Attributes attrs;
+    boolean shutdown;
+    ScheduledHandle delayedShutdownTask;
 
-    @GuardedBy("shutdownLock")
+    SubchannelImpl(CreateSubchannelArgs args, LbHelperImpl helper) {
-    boolean shutdownRequested;
+      this.args = checkNotNull(args, "args");
-    @GuardedBy("shutdownLock")
+      this.helper = checkNotNull(helper, "helper");
-    ScheduledFuture<?> delayedShutdownTask;
+    }
 
-    SubchannelImpl(Attributes attrs) {
+    @Override
-      this.attrs = checkNotNull(attrs, "attrs");
+    public void start(final SubchannelStateListener listener) {
+      syncContext.throwIfNotInThisSynchronizationContext();
+      checkState(!started, "already started");
+      checkState(!shutdown, "already shutdown");
+      started = true;
+      this.listener = listener;
+      if (terminating) {
+        syncContext.execute(new Runnable() {
+            @Override
+            public void run() {
+              listener.onSubchannelState(ConnectivityStateInfo.forNonError(SHUTDOWN));
+            }
+          });
+        return;
+      }
+      final class ManagedInternalSubchannelCallback extends InternalSubchannel.Callback {
+        // All callbacks are run in syncContext
+        @Override
+        void onTerminated(InternalSubchannel is) {
+          subchannels.remove(is);
+          channelz.removeSubchannel(is);
+          maybeTerminateChannel();
+        }
+
+        @Override
+        void onStateChange(InternalSubchannel is, ConnectivityStateInfo newState) {
+          handleInternalSubchannelState(newState);
+          // Call LB only if it's not shutdown.  If LB is shutdown, lbHelper won't match.
+          if (helper == ManagedChannelImpl.this.lbHelper) {
+            checkState(listener != null, "listener is null");
+            listener.onSubchannelState(newState);
+          }
+        }
+
+        @Override
+        void onInUse(InternalSubchannel is) {
+          inUseStateAggregator.updateObjectInUse(is, true);
+        }
+
+        @Override
+        void onNotInUse(InternalSubchannel is) {
+          inUseStateAggregator.updateObjectInUse(is, false);
+        }
+      }
+
+      long subchannelCreationTime = timeProvider.currentTimeNanos();
+      InternalLogId subchannelLogId = InternalLogId.allocate("Subchannel", /*details=*/ null);
+      ChannelTracer subchannelTracer =
+          new ChannelTracer(
+              subchannelLogId, maxTraceEvents, subchannelCreationTime,
+              "Subchannel for " + args.getAddresses());
+      InternalSubchannel internalSubchannel = new InternalSubchannel(
+          args.getAddresses(),
+          authority(),
+          userAgent,
+          backoffPolicyProvider,
+          transportFactory,
+          transportFactory.getScheduledExecutorService(),
+          stopwatchSupplier,
+          syncContext,
+          new ManagedInternalSubchannelCallback(),
+          channelz,
+          callTracerFactory.create(),
+          subchannelTracer,
+          subchannelLogId,
+          timeProvider);
+
+      channelTracer.reportEvent(new ChannelTrace.Event.Builder()
+          .setDescription("Child Subchannel started")
+          .setSeverity(ChannelTrace.Event.Severity.CT_INFO)
+          .setTimestampNanos(subchannelCreationTime)
+          .setSubchannelRef(internalSubchannel)
+          .build());
+
+      channelz.addSubchannel(internalSubchannel);
+      this.subchannel = internalSubchannel;
+      subchannels.add(internalSubchannel);
     }
 
     @Override
     ClientTransport obtainActiveTransport() {
+      checkState(started, "Subchannel is not started");
       return subchannel.obtainActiveTransport();
     }
 
     @Override
     InternalInstrumented<ChannelStats> getInternalSubchannel() {
+      checkState(started, "not started");
       return subchannel;
     }
 
     @Override
     public void shutdown() {
+      // TODO(zhangkun83): replace shutdown() with internalShutdown() to turn the warning into an
+      // exception.
       logWarningIfNotInSyncContext("Subchannel.shutdown()");
-      synchronized (shutdownLock) {
+      syncContext.execute(new Runnable() {
-        if (shutdownRequested) {
+          @Override
+          public void run() {
+            internalShutdown();
+          }
+        });
+    }
+
+    private void internalShutdown() {
+      syncContext.throwIfNotInThisSynchronizationContext();
+      if (subchannel == null) {
+        // start() was not successful
+        shutdown = true;
+        return;
+      }
+      if (shutdown) {
         if (terminating && delayedShutdownTask != null) {
           // shutdown() was previously called when terminating == false, thus a delayed shutdown()
           // was scheduled.  Now since terminating == true, We should expedite the shutdown.
-            delayedShutdownTask.cancel(false);
+          delayedShutdownTask.cancel();
           delayedShutdownTask = null;
           // Will fall through to the subchannel.shutdown() at the end.
         } else {
           return;
         }
       } else {
-          shutdownRequested = true;
+        shutdown = true;
       }
       // Add a delay to shutdown to deal with the race between 1) a transport being picked and
       // newStream() being called on it, and 2) its Subchannel is shut down by LoadBalancer (e.g.,
@@ -1509,12 +1559,12 @@ final class ManagedChannelImpl extends ManagedChannel implements
           }
         }
 
-          delayedShutdownTask = transportFactory.getScheduledExecutorService().schedule(
+        delayedShutdownTask = syncContext.schedule(
-              new LogExceptionRunnable(
+            new LogExceptionRunnable(new ShutdownSubchannel()),
-                  new ShutdownSubchannel()), SUBCHANNEL_SHUTDOWN_DELAY_SECONDS, TimeUnit.SECONDS);
+            SUBCHANNEL_SHUTDOWN_DELAY_SECONDS, TimeUnit.SECONDS,
+            transportFactory.getScheduledExecutorService());
         return;
       }
-      }
       // When terminating == true, no more real streams will be created. It's safe and also
       // desirable to shutdown timely.
       subchannel.shutdown(SHUTDOWN_STATUS);
@@ -1522,18 +1572,20 @@ final class ManagedChannelImpl extends ManagedChannel implements
 
     @Override
     public void requestConnection() {
+      checkState(started, "not started");
       subchannel.obtainActiveTransport();
     }
 
     @Override
     public List<EquivalentAddressGroup> getAllAddresses() {
       logWarningIfNotInSyncContext("Subchannel.getAllAddresses()");
+      checkState(started, "not started");
       return subchannel.getAddressGroups();
     }
 
     @Override
     public Attributes getAttributes() {
-      return attrs;
+      return args.getAttributes();
     }
 
     @Override

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

@@ -21,16 +21,11 @@ import static io.grpc.ConnectivityState.CONNECTING;
 import static io.grpc.ConnectivityState.SHUTDOWN;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 
-import io.grpc.Attributes;
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
-import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancer.SubchannelStateListener;
-import io.grpc.LoadBalancer.PickResult;
-import io.grpc.LoadBalancer.PickSubchannelArgs;
-import io.grpc.LoadBalancer.Subchannel;
-import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.Status;
 import java.util.List;
 
@@ -51,7 +46,17 @@ final class PickFirstLoadBalancer extends LoadBalancer {
   public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
     List<EquivalentAddressGroup> servers = resolvedAddresses.getAddresses();
     if (subchannel == null) {
-      subchannel = helper.createSubchannel(servers, Attributes.EMPTY);
+      final Subchannel subchannel = helper.createSubchannel(
+          CreateSubchannelArgs.newBuilder()
+              .setAddresses(servers)
+              .build());
+      subchannel.start(new SubchannelStateListener() {
+          @Override
+          public void onSubchannelState(ConnectivityStateInfo stateInfo) {
+            processSubchannelState(subchannel, stateInfo);
+          }
+        });
+      this.subchannel = subchannel;
 
       // The channel state does not get updated when doing name resolving today, so for the moment
       // let LB report CONNECTION and call subchannel.requestConnection() immediately.
@@ -73,10 +78,9 @@ final class PickFirstLoadBalancer extends LoadBalancer {
     helper.updateBalancingState(TRANSIENT_FAILURE, new Picker(PickResult.withError(error)));
   }
 
-  @Override
+  private void processSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
-  public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
     ConnectivityState currentState = stateInfo.getState();
-    if (subchannel != this.subchannel || currentState == SHUTDOWN) {
+    if (currentState == SHUTDOWN) {
       return;
     }
 
@@ -99,7 +103,6 @@ final class PickFirstLoadBalancer extends LoadBalancer {
       default:
         throw new IllegalArgumentException("Unsupported state:" + currentState);
     }
-
     helper.updateBalancingState(currentState, picker);
   }
 

core/src/main/java/io/grpc/util/ForwardingLoadBalancer.java

@@ -51,6 +51,7 @@ public abstract class ForwardingLoadBalancer extends LoadBalancer {
     delegate().handleNameResolutionError(error);
   }
 
+  @Deprecated
   @Override
   public void handleSubchannelState(
       Subchannel subchannel, ConnectivityStateInfo stateInfo) {

core/src/main/java/io/grpc/util/ForwardingLoadBalancerHelper.java

@@ -22,6 +22,7 @@ import io.grpc.ChannelLogger;
 import io.grpc.ConnectivityState;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.ExperimentalApi;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.LoadBalancer;
@@ -39,11 +40,17 @@ public abstract class ForwardingLoadBalancerHelper extends LoadBalancer.Helper {
    */
   protected abstract LoadBalancer.Helper delegate();
 
+  @Deprecated
   @Override
   public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
     return delegate().createSubchannel(addrs, attrs);
   }
 
+  @Override
+  public Subchannel createSubchannel(CreateSubchannelArgs args) {
+    return delegate().createSubchannel(args);
+  }
+
   @Override
   public void updateSubchannelAddresses(
       Subchannel subchannel, List<EquivalentAddressGroup> addrs) {

core/src/main/java/io/grpc/util/ForwardingSubchannel.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2019 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.util;
+
+import com.google.common.base.MoreObjects;
+import io.grpc.Attributes;
+import io.grpc.Channel;
+import io.grpc.ChannelLogger;
+import io.grpc.EquivalentAddressGroup;
+import io.grpc.ExperimentalApi;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.Subchannel;
+import io.grpc.LoadBalancer.SubchannelStateListener;
+import java.util.List;
+
+@ExperimentalApi("https://github.com/grpc/grpc-java/issues/1771")
+public abstract class ForwardingSubchannel extends LoadBalancer.Subchannel {
+  /**
+   * Returns the underlying Subchannel.
+   */
+  protected abstract Subchannel delegate();
+
+  @Override
+  public void start(SubchannelStateListener listener) {
+    delegate().start(listener);
+  }
+
+  @Override
+  public void shutdown() {
+    delegate().shutdown();
+  }
+
+  @Override
+  public void requestConnection() {
+    delegate().requestConnection();
+  }
+
+  @Override
+  public List<EquivalentAddressGroup> getAllAddresses() {
+    return delegate().getAllAddresses();
+  }
+
+  @Override
+  public Attributes getAttributes() {
+    return delegate().getAttributes();
+  }
+
+  @Override
+  public Channel asChannel() {
+    return delegate().asChannel();
+  }
+
+  @Override
+  public ChannelLogger getChannelLogger() {
+    return delegate().getChannelLogger();
+  }
+
+  @Override
+  public String toString() {
+    return MoreObjects.toStringHelper(this).add("delegate", delegate()).toString();
+  }
+}

core/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java

@@ -33,10 +33,7 @@ import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
-import io.grpc.LoadBalancer.PickResult;
+import io.grpc.LoadBalancer.SubchannelStateListener;
-import io.grpc.LoadBalancer.PickSubchannelArgs;
-import io.grpc.LoadBalancer.Subchannel;
-import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.Metadata;
 import io.grpc.Metadata.Key;
 import io.grpc.NameResolver;
@@ -129,8 +126,18 @@ final class RoundRobinLoadBalancer extends LoadBalancer {
         subchannelAttrs.set(STICKY_REF, stickyRef = new Ref<>(null));
       }
 
-      Subchannel subchannel = checkNotNull(
+      final Subchannel subchannel = checkNotNull(
-          helper.createSubchannel(addressGroup, subchannelAttrs.build()), "subchannel");
+          helper.createSubchannel(CreateSubchannelArgs.newBuilder()
+              .setAddresses(addressGroup)
+              .setAttributes(subchannelAttrs.build())
+              .build()),
+          "subchannel");
+      subchannel.start(new SubchannelStateListener() {
+          @Override
+          public void onSubchannelState(ConnectivityStateInfo state) {
+            processSubchannelState(subchannel, state);
+          }
+        });
       if (stickyRef != null) {
         stickyRef.value = subchannel;
       }
@@ -160,8 +167,7 @@ final class RoundRobinLoadBalancer extends LoadBalancer {
         currentPicker instanceof ReadyPicker ? currentPicker : new EmptyPicker(error));
   }
 
-  @Override
+  private void processSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
-  public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
     if (subchannels.get(subchannel.getAddresses()) != subchannel) {
       return;
     }

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

@@ -41,10 +41,12 @@ import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.ManagedChannel;
@@ -134,6 +136,7 @@ public class AutoConfiguredLoadBalancerFactoryTest {
     assertThat(lb.getDelegate()).isSameInstanceAs(testLbBalancer);
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void forwardsCalls() {
     AutoConfiguredLoadBalancer lb =
@@ -176,9 +179,9 @@ public class AutoConfiguredLoadBalancerFactoryTest {
         Collections.singletonList(new EquivalentAddressGroup(new SocketAddress(){}));
     Helper helper = new TestHelper() {
       @Override
-      public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
+      public Subchannel createSubchannel(CreateSubchannelArgs args) {
-        assertThat(addrs).isEqualTo(servers);
+        assertThat(args.getAddresses()).isEqualTo(servers);
-        return new TestSubchannel(addrs, attrs);
+        return new TestSubchannel(args);
       }
     };
     AutoConfiguredLoadBalancer lb =
@@ -206,9 +209,9 @@ public class AutoConfiguredLoadBalancerFactoryTest {
         Collections.singletonList(new EquivalentAddressGroup(new SocketAddress(){}));
     Helper helper = new TestHelper() {
       @Override
-      public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
+      public Subchannel createSubchannel(CreateSubchannelArgs args) {
-        assertThat(addrs).isEqualTo(servers);
+        assertThat(args.getAddresses()).isEqualTo(servers);
-        return new TestSubchannel(addrs, attrs);
+        return new TestSubchannel(args);
       }
     };
     AutoConfiguredLoadBalancer lb =
@@ -221,11 +224,6 @@ public class AutoConfiguredLoadBalancerFactoryTest {
         // noop
       }
 
-      @Override
-      public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
-        // noop
-      }
-
       @Override
       public void shutdown() {
         shutdown.set(true);
@@ -704,8 +702,17 @@ public class AutoConfiguredLoadBalancerFactoryTest {
         Collections.singletonList(new EquivalentAddressGroup(new SocketAddress(){}));
     Helper helper = new TestHelper() {
       @Override
+      @Deprecated
       public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
-        return new TestSubchannel(addrs, attrs);
+        return new TestSubchannel(CreateSubchannelArgs.newBuilder()
+            .setAddresses(addrs)
+            .setAttributes(attrs)
+            .build());
+      }
+
+      @Override
+      public Subchannel createSubchannel(CreateSubchannelArgs args) {
+        return new TestSubchannel(args);
       }
 
       @Override
@@ -822,11 +829,6 @@ public class AutoConfiguredLoadBalancerFactoryTest {
       delegate().handleNameResolutionError(error);
     }
 
-    @Override
-    public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
-      delegate().handleSubchannelState(subchannel, stateInfo);
-    }
-
     @Override
     public void shutdown() {
       delegate().shutdown();
@@ -862,14 +864,18 @@ public class AutoConfiguredLoadBalancerFactoryTest {
   }
 
   private static class TestSubchannel extends Subchannel {
-    TestSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
+    TestSubchannel(CreateSubchannelArgs args) {
-      this.addrs = addrs;
+      this.addrs = args.getAddresses();
-      this.attrs = attrs;
+      this.attrs = args.getAttributes();
     }
 
     final List<EquivalentAddressGroup> addrs;
     final Attributes attrs;
 
+    @Override
+    public void start(SubchannelStateListener listener) {
+    }
+
     @Override
     public void shutdown() {
     }

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

@@ -42,12 +42,14 @@ import io.grpc.ConnectivityState;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.IntegerMarshaller;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.ManagedChannel;
@@ -114,6 +116,7 @@ public class ManagedChannelImplIdlenessTest {
 
   @Mock private ClientTransportFactory mockTransportFactory;
   @Mock private LoadBalancer mockLoadBalancer;
+  @Mock private SubchannelStateListener subchannelStateListener;
   private final LoadBalancerProvider mockLoadBalancerProvider =
       mock(LoadBalancerProvider.class, delegatesTo(new LoadBalancerProvider() {
           @Override
@@ -501,14 +504,19 @@ public class ManagedChannelImplIdlenessTest {
   }
 
   // Helper methods to call methods from SynchronizationContext
-  private static Subchannel createSubchannelSafely(
+  private Subchannel createSubchannelSafely(
       final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs) {
     final AtomicReference<Subchannel> resultCapture = new AtomicReference<>();
     helper.getSynchronizationContext().execute(
         new Runnable() {
           @Override
           public void run() {
-            resultCapture.set(helper.createSubchannel(addressGroup, attrs));
+            Subchannel s = helper.createSubchannel(CreateSubchannelArgs.newBuilder()
+                .setAddresses(addressGroup)
+                .setAttributes(attrs)
+                .build());
+            s.start(subchannelStateListener);
+            resultCapture.set(s);
           }
         });
     return resultCapture.get();

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

@@ -78,12 +78,14 @@ import io.grpc.InternalChannelz.ChannelStats;
 import io.grpc.InternalChannelz.ChannelTrace;
 import io.grpc.InternalInstrumented;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.ManagedChannel;
@@ -207,6 +209,8 @@ public class ManagedChannelImplTest {
   private ArgumentCaptor<CallOptions> callOptionsCaptor;
   @Mock
   private LoadBalancer mockLoadBalancer;
+  @Mock
+  private SubchannelStateListener subchannelStateListener;
   private final LoadBalancerProvider mockLoadBalancerProvider =
       mock(LoadBalancerProvider.class, delegatesTo(new LoadBalancerProvider() {
           @Override
@@ -334,8 +338,9 @@ public class ManagedChannelImplTest {
     LoadBalancerRegistry.getDefaultRegistry().deregister(mockLoadBalancerProvider);
   }
 
+  @Deprecated
   @Test
-  public void createSubchannelOutsideSynchronizationContextShouldLogWarning() {
+  public void createSubchannel_old_outsideSynchronizationContextShouldLogWarning() {
     createChannel();
     final AtomicReference<LogRecord> logRef = new AtomicReference<>();
     Handler handler = new Handler() {
@@ -366,6 +371,47 @@ public class ManagedChannelImplTest {
     }
   }
 
+  @Deprecated
+  @Test
+  public void createSubchannel_old_propagateSubchannelStatesToOldApi() {
+    createChannel();
+    final AtomicReference<Subchannel> subchannelCapture = new AtomicReference<>();
+    helper.getSynchronizationContext().execute(new Runnable() {
+        @Override
+        public void run() {
+          subchannelCapture.set(helper.createSubchannel(addressGroup, Attributes.EMPTY));
+        }
+      });
+
+    Subchannel subchannel = subchannelCapture.get();
+    subchannel.requestConnection();
+
+    verify(mockTransportFactory)
+        .newClientTransport(
+            any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
+    verify(mockLoadBalancer).handleSubchannelState(
+        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+
+    MockClientTransportInfo transportInfo = transports.poll();
+    transportInfo.listener.transportReady();
+
+    verify(mockLoadBalancer).handleSubchannelState(
+        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+  }
+
+  @Test
+  public void createSubchannel_outsideSynchronizationContextShouldThrow() {
+    createChannel();
+    try {
+      helper.createSubchannel(CreateSubchannelArgs.newBuilder()
+          .setAddresses(addressGroup)
+          .build());
+      fail("Should throw");
+    } catch (IllegalStateException e) {
+      assertThat(e).hasMessageThat().isEqualTo("Not called from the SynchronizationContext");
+    }
+  }
+
   @Test
   @SuppressWarnings("unchecked")
   public void idleModeDisabled() {
@@ -426,12 +472,12 @@ public class ManagedChannelImplTest {
     assertNotNull(channelz.getRootChannel(channel.getLogId().getId()));
 
     AbstractSubchannel subchannel =
-        (AbstractSubchannel) createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        (AbstractSubchannel) createSubchannelSafely(
+            helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     // subchannels are not root channels
     assertNull(channelz.getRootChannel(subchannel.getInternalSubchannel().getLogId().getId()));
     assertTrue(channelz.containsSubchannel(subchannel.getInternalSubchannel().getLogId()));
-    assertThat(getStats(channel).subchannels)
+    assertThat(getStats(channel).subchannels).containsExactly(subchannel.getInternalSubchannel());
-        .containsExactly(subchannel.getInternalSubchannel());
 
     requestConnectionSafely(helper, subchannel);
     MockClientTransportInfo transportInfo = transports.poll();
@@ -466,8 +512,7 @@ public class ManagedChannelImplTest {
 
     AbstractSubchannel subchannel = (AbstractSubchannel) oob.getSubchannel();
     assertTrue(channelz.containsSubchannel(subchannel.getInternalSubchannel().getLogId()));
-    assertThat(getStats(oob).subchannels)
+    assertThat(getStats(oob).subchannels).containsExactly(subchannel.getInternalSubchannel());
-        .containsExactly(subchannel.getInternalSubchannel());
     assertTrue(channelz.containsSubchannel(subchannel.getInternalSubchannel().getLogId()));
 
     oob.getSubchannel().requestConnection();
@@ -518,7 +563,8 @@ public class ManagedChannelImplTest {
 
     // Configure the picker so that first RPC goes to delayed transport, and second RPC goes to
     // real transport.
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     verify(mockTransportFactory)
         .newClientTransport(
@@ -657,8 +703,12 @@ public class ManagedChannelImplTest {
             .setAttributes(Attributes.EMPTY)
             .build());
 
-    Subchannel subchannel1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    SubchannelStateListener stateListener1 = mock(SubchannelStateListener.class);
-    Subchannel subchannel2 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    SubchannelStateListener stateListener2 = mock(SubchannelStateListener.class);
+    Subchannel subchannel1 =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, stateListener1);
+    Subchannel subchannel2 =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, stateListener2);
     requestConnectionSafely(helper, subchannel1);
     requestConnectionSafely(helper, subchannel2);
     verify(mockTransportFactory, times(2))
@@ -669,13 +719,12 @@ public class ManagedChannelImplTest {
 
     // LoadBalancer receives all sorts of callbacks
     transportInfo1.listener.transportReady();
-    verify(mockLoadBalancer, times(2))
+
-        .handleSubchannelState(same(subchannel1), stateInfoCaptor.capture());
+    verify(stateListener1, times(2)).onSubchannelState(stateInfoCaptor.capture());
     assertSame(CONNECTING, stateInfoCaptor.getAllValues().get(0).getState());
     assertSame(READY, stateInfoCaptor.getAllValues().get(1).getState());
 
-    verify(mockLoadBalancer)
+    verify(stateListener2).onSubchannelState(stateInfoCaptor.capture());
-        .handleSubchannelState(same(subchannel2), stateInfoCaptor.capture());
     assertSame(CONNECTING, stateInfoCaptor.getValue().getState());
 
     resolver.listener.onError(resolutionError);
@@ -724,7 +773,8 @@ public class ManagedChannelImplTest {
     call.start(mockCallListener, headers);
 
     // Make the transport available
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     verify(mockTransportFactory, never())
         .newClientTransport(
             any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
@@ -961,7 +1011,7 @@ public class ManagedChannelImplTest {
           return "badAddress";
         }
       };
-    InOrder inOrder = inOrder(mockLoadBalancer);
+    InOrder inOrder = inOrder(mockLoadBalancer, subchannelStateListener);
 
     List<SocketAddress> resolvedAddrs = Arrays.asList(badAddress, goodAddress);
     FakeNameResolverFactory nameResolverFactory =
@@ -983,12 +1033,12 @@ public class ManagedChannelImplTest {
         ResolvedAddresses.newBuilder()
             .setAddresses(Arrays.asList(addressGroup))
             .build());
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
         .thenReturn(PickResult.withSubchannel(subchannel));
     requestConnectionSafely(helper, subchannel);
-    inOrder.verify(mockLoadBalancer).handleSubchannelState(
+    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
-        same(subchannel), stateInfoCaptor.capture());
     assertEquals(CONNECTING, stateInfoCaptor.getValue().getState());
 
     // The channel will starts with the first address (badAddress)
@@ -1014,8 +1064,7 @@ public class ManagedChannelImplTest {
         .thenReturn(mock(ClientStream.class));
 
     goodTransportInfo.listener.transportReady();
-    inOrder.verify(mockLoadBalancer).handleSubchannelState(
+    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
-        same(subchannel), stateInfoCaptor.capture());
     assertEquals(READY, stateInfoCaptor.getValue().getState());
 
     // A typical LoadBalancer will call this once the subchannel becomes READY
@@ -1105,7 +1154,7 @@ public class ManagedChannelImplTest {
           return "addr2";
         }
       };
-    InOrder inOrder = inOrder(mockLoadBalancer);
+    InOrder inOrder = inOrder(mockLoadBalancer, subchannelStateListener);
 
     List<SocketAddress> resolvedAddrs = Arrays.asList(addr1, addr2);
 
@@ -1133,13 +1182,13 @@ public class ManagedChannelImplTest {
         ResolvedAddresses.newBuilder()
             .setAddresses(Arrays.asList(addressGroup))
             .build());
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
         .thenReturn(PickResult.withSubchannel(subchannel));
     requestConnectionSafely(helper, subchannel);
 
-    inOrder.verify(mockLoadBalancer).handleSubchannelState(
+    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
-        same(subchannel), stateInfoCaptor.capture());
     assertEquals(CONNECTING, stateInfoCaptor.getValue().getState());
 
     // Connecting to server1, which will fail
@@ -1162,8 +1211,7 @@ public class ManagedChannelImplTest {
 
     // ... which makes the subchannel enter TRANSIENT_FAILURE. The last error Status is propagated
     // to LoadBalancer.
-    inOrder.verify(mockLoadBalancer).handleSubchannelState(
+    inOrder.verify(subchannelStateListener).onSubchannelState(stateInfoCaptor.capture());
-        same(subchannel), stateInfoCaptor.capture());
     assertEquals(TRANSIENT_FAILURE, stateInfoCaptor.getValue().getState());
     assertSame(server2Error, stateInfoCaptor.getValue().getStatus());
 
@@ -1192,8 +1240,10 @@ public class ManagedChannelImplTest {
     // createSubchannel() always return a new Subchannel
     Attributes attrs1 = Attributes.newBuilder().set(SUBCHANNEL_ATTR_KEY, "attr1").build();
     Attributes attrs2 = Attributes.newBuilder().set(SUBCHANNEL_ATTR_KEY, "attr2").build();
-    final Subchannel sub1 = createSubchannelSafely(helper, addressGroup, attrs1);
+    SubchannelStateListener listener1 = mock(SubchannelStateListener.class);
-    final Subchannel sub2 = createSubchannelSafely(helper, addressGroup, attrs2);
+    SubchannelStateListener listener2 = mock(SubchannelStateListener.class);
+    final Subchannel sub1 = createSubchannelSafely(helper, addressGroup, attrs1, listener1);
+    final Subchannel sub2 = createSubchannelSafely(helper, addressGroup, attrs2, listener2);
     assertNotSame(sub1, sub2);
     assertNotSame(attrs1, attrs2);
     assertSame(attrs1, sub1.getAttributes());
@@ -1270,8 +1320,10 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelsWhenChannelShutdownNow() {
     createChannel();
-    Subchannel sub1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel sub1 =
-    Subchannel sub2 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
+    Subchannel sub2 =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, sub1);
     requestConnectionSafely(helper, sub2);
 
@@ -1298,8 +1350,10 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelsNoConnectionShutdown() {
     createChannel();
-    Subchannel sub1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel sub1 =
-    Subchannel sub2 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
+    Subchannel sub2 =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
 
     channel.shutdown();
     verify(mockLoadBalancer).shutdown();
@@ -1315,8 +1369,8 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelsNoConnectionShutdownNow() {
     createChannel();
-    createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
-    createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     channel.shutdownNow();
 
     verify(mockLoadBalancer).shutdown();
@@ -1498,7 +1552,8 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelChannel_normalUsage() {
     createChannel();
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     verify(balancerRpcExecutorPool, never()).getObject();
 
     Channel sChannel = subchannel.asChannel();
@@ -1529,7 +1584,8 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelChannel_failWhenNotReady() {
     createChannel();
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     Channel sChannel = subchannel.asChannel();
     Metadata headers = new Metadata();
 
@@ -1557,7 +1613,8 @@ public class ManagedChannelImplTest {
   @Test
   public void subchannelChannel_failWaitForReady() {
     createChannel();
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     Channel sChannel = subchannel.asChannel();
     Metadata headers = new Metadata();
 
@@ -1664,7 +1721,8 @@ public class ManagedChannelImplTest {
       OobChannel oobChannel = (OobChannel) helper.createOobChannel(addressGroup, "oobAuthority");
       oobChannel.getSubchannel().requestConnection();
     } else {
-      Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+      Subchannel subchannel =
+          createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
       requestConnectionSafely(helper, subchannel);
     }
 
@@ -1751,7 +1809,8 @@ public class ManagedChannelImplTest {
 
     // Simulate name resolution results
     EquivalentAddressGroup addressGroup = new EquivalentAddressGroup(socketAddress);
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     verify(mockTransportFactory)
         .newClientTransport(
@@ -1824,7 +1883,8 @@ public class ManagedChannelImplTest {
     ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
     ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
     createChannel();
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     MockClientTransportInfo transportInfo = transports.poll();
     transportInfo.listener.transportReady();
@@ -1862,7 +1922,8 @@ public class ManagedChannelImplTest {
     ClientCall<String, Integer> call = channel.newCall(method, callOptions);
     call.start(mockCallListener, new Metadata());
 
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     MockClientTransportInfo transportInfo = transports.poll();
     transportInfo.listener.transportReady();
@@ -2239,7 +2300,8 @@ public class ManagedChannelImplTest {
     Helper helper2 = helperCaptor.getValue();
 
     // Establish a connection
-    Subchannel subchannel = createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     MockClientTransportInfo transportInfo = transports.poll();
     ConnectionClientTransport mockTransport = transportInfo.transport;
@@ -2307,7 +2369,8 @@ public class ManagedChannelImplTest {
     Helper helper2 = helperCaptor.getValue();
 
     // Establish a connection
-    Subchannel subchannel = createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     ClientStream mockStream = mock(ClientStream.class);
     MockClientTransportInfo transportInfo = transports.poll();
@@ -2336,8 +2399,10 @@ public class ManagedChannelImplTest {
     call.start(mockCallListener, new Metadata());
 
     // Make the transport available with subchannel2
-    Subchannel subchannel1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel1 =
-    Subchannel subchannel2 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
+    Subchannel subchannel2 =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel2);
 
     MockClientTransportInfo transportInfo = transports.poll();
@@ -2476,7 +2541,8 @@ public class ManagedChannelImplTest {
     createChannel();
     assertEquals(TARGET, getStats(channel).target);
 
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     assertEquals(Collections.singletonList(addressGroup).toString(),
         getStats((AbstractSubchannel) subchannel).target);
   }
@@ -2499,9 +2565,10 @@ public class ManagedChannelImplTest {
     createChannel();
     timer.forwardNanos(1234);
     AbstractSubchannel subchannel =
-        (AbstractSubchannel) createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        (AbstractSubchannel) createSubchannelSafely(
+            helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     assertThat(getStats(channel).channelTrace.events).contains(new ChannelTrace.Event.Builder()
-        .setDescription("Child Subchannel created")
+        .setDescription("Child Subchannel started")
         .setSeverity(ChannelTrace.Event.Severity.CT_INFO)
         .setTimestampNanos(timer.getTicker().read())
         .setSubchannelRef(subchannel.getInternalSubchannel())
@@ -2672,7 +2739,8 @@ public class ManagedChannelImplTest {
     channelBuilder.maxTraceEvents(10);
     createChannel();
     AbstractSubchannel subchannel =
-        (AbstractSubchannel) createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        (AbstractSubchannel) createSubchannelSafely(
+            helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     timer.forwardNanos(1234);
     subchannel.obtainActiveTransport();
     assertThat(getStats(subchannel).channelTrace.events).contains(new ChannelTrace.Event.Builder()
@@ -2735,7 +2803,8 @@ public class ManagedChannelImplTest {
     assertEquals(CONNECTING, getStats(channel).state);
 
     AbstractSubchannel subchannel =
-        (AbstractSubchannel) createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        (AbstractSubchannel) createSubchannelSafely(
+            helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
 
     assertEquals(IDLE, getStats(subchannel).state);
     requestConnectionSafely(helper, subchannel);
@@ -2789,7 +2858,8 @@ public class ManagedChannelImplTest {
     ClientStream mockStream = mock(ClientStream.class);
     ClientStreamTracer.Factory factory = mock(ClientStreamTracer.Factory.class);
     AbstractSubchannel subchannel =
-        (AbstractSubchannel) createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+        (AbstractSubchannel) createSubchannelSafely(
+            helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     requestConnectionSafely(helper, subchannel);
     MockClientTransportInfo transportInfo = transports.poll();
     transportInfo.listener.transportReady();
@@ -3026,7 +3096,8 @@ public class ManagedChannelImplTest {
             .build());
 
     // simulating request connection and then transport ready after resolved address
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
         .thenReturn(PickResult.withSubchannel(subchannel));
     requestConnectionSafely(helper, subchannel);
@@ -3125,7 +3196,8 @@ public class ManagedChannelImplTest {
             .build());
 
     // simulating request connection and then transport ready after resolved address
-    Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
+    Subchannel subchannel =
+        createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
     when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
         .thenReturn(PickResult.withSubchannel(subchannel));
     requestConnectionSafely(helper, subchannel);
@@ -3900,13 +3972,19 @@ public class ManagedChannelImplTest {
 
   // Helper methods to call methods from SynchronizationContext
   private static Subchannel createSubchannelSafely(
-      final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs) {
+      final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs,
+      final SubchannelStateListener stateListener) {
     final AtomicReference<Subchannel> resultCapture = new AtomicReference<>();
     helper.getSynchronizationContext().execute(
         new Runnable() {
           @Override
           public void run() {
-            resultCapture.set(helper.createSubchannel(addressGroup, attrs));
+            Subchannel s = helper.createSubchannel(CreateSubchannelArgs.newBuilder()
+                .setAddresses(addressGroup)
+                .setAttributes(attrs)
+                .build());
+            s.start(stateListener);
+            resultCapture.set(s);
           }
         });
     return resultCapture.get();

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

@@ -16,6 +16,7 @@
 
 package io.grpc.internal;
 
+import static com.google.common.truth.Truth.assertThat;
 import static io.grpc.ConnectivityState.CONNECTING;
 import static io.grpc.ConnectivityState.IDLE;
 import static io.grpc.ConnectivityState.READY;
@@ -24,10 +25,8 @@ import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isA;
-import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -38,12 +37,14 @@ import io.grpc.Attributes;
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.Status;
 import java.net.SocketAddress;
 import java.util.List;
@@ -76,7 +77,9 @@ public class PickFirstLoadBalancerTest {
   @Captor
   private ArgumentCaptor<SubchannelPicker> pickerCaptor;
   @Captor
-  private ArgumentCaptor<Attributes> attrsCaptor;
+  private ArgumentCaptor<CreateSubchannelArgs> createArgsCaptor;
+  @Captor
+  private ArgumentCaptor<SubchannelStateListener> stateListenerCaptor;
   @Mock
   private Helper mockHelper;
   @Mock
@@ -93,16 +96,17 @@ public class PickFirstLoadBalancerTest {
     }
 
     when(mockSubchannel.getAllAddresses()).thenThrow(new UnsupportedOperationException());
-    when(mockHelper.createSubchannel(
+    when(mockHelper.createSubchannel(any(CreateSubchannelArgs.class))).thenReturn(mockSubchannel);
-        ArgumentMatchers.<EquivalentAddressGroup>anyList(), any(Attributes.class)))
-        .thenReturn(mockSubchannel);
 
     loadBalancer = new PickFirstLoadBalancer(mockHelper);
   }
 
   @After
+  @SuppressWarnings("deprecation")
   public void tearDown() throws Exception {
     verifyNoMoreInteractions(mockArgs);
+    verify(mockHelper, never()).createSubchannel(
+        ArgumentMatchers.<EquivalentAddressGroup>anyList(), any(Attributes.class));
   }
 
   @Test
@@ -110,7 +114,9 @@ public class PickFirstLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
 
-    verify(mockHelper).createSubchannel(eq(servers), attrsCaptor.capture());
+    verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
     verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
     verify(mockSubchannel).requestConnection();
 
@@ -124,13 +130,14 @@ public class PickFirstLoadBalancerTest {
   public void pickAfterResolvedAndUnchanged() throws Exception {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
+    verify(mockSubchannel).start(any(SubchannelStateListener.class));
     verify(mockSubchannel).requestConnection();
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
     verifyNoMoreInteractions(mockSubchannel);
 
-    verify(mockHelper).createSubchannel(ArgumentMatchers.<EquivalentAddressGroup>anyList(),
+    verify(mockHelper).createSubchannel(createArgsCaptor.capture());
-        any(Attributes.class));
+    assertThat(createArgsCaptor.getValue()).isNotNull();
     verify(mockHelper)
         .updateBalancingState(isA(ConnectivityState.class), isA(SubchannelPicker.class));
     // Updating the subchannel addresses is unnecessary, but doesn't hurt anything
@@ -150,7 +157,10 @@ public class PickFirstLoadBalancerTest {
 
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
-    inOrder.verify(mockHelper).createSubchannel(eq(servers), any(Attributes.class));
+    inOrder.verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    verify(mockSubchannel).start(any(SubchannelStateListener.class));
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
     verify(mockSubchannel).requestConnection();
     assertEquals(mockSubchannel, pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel());
@@ -163,34 +173,30 @@ public class PickFirstLoadBalancerTest {
     verifyNoMoreInteractions(mockHelper);
   }
 
-  @Test
-  public void stateChangeBeforeResolution() throws Exception {
-    loadBalancer.handleSubchannelState(mockSubchannel, ConnectivityStateInfo.forNonError(READY));
-
-    verifyNoMoreInteractions(mockHelper);
-  }
-
   @Test
   public void pickAfterStateChangeAfterResolution() throws Exception {
-    loadBalancer.handleResolvedAddresses(
-        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
-    verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
-    Subchannel subchannel = pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel();
-    reset(mockHelper);
-
     InOrder inOrder = inOrder(mockHelper);
 
+    loadBalancer.handleResolvedAddresses(
+        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
+    inOrder.verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
+    verify(mockSubchannel).start(stateListenerCaptor.capture());
+    SubchannelStateListener stateListener = stateListenerCaptor.getValue();
+    verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
+    Subchannel subchannel = pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel();
+
     Status error = Status.UNAVAILABLE.withDescription("boom!");
-    loadBalancer.handleSubchannelState(subchannel,
+    stateListener.onSubchannelState(ConnectivityStateInfo.forTransientFailure(error));
-        ConnectivityStateInfo.forTransientFailure(error));
     inOrder.verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
     assertEquals(error, pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus());
 
-    loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
+    stateListener.onSubchannelState(ConnectivityStateInfo.forNonError(IDLE));
     inOrder.verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
     assertEquals(Status.OK, pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus());
 
-    loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    stateListener.onSubchannelState(ConnectivityStateInfo.forNonError(READY));
     inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
     assertEquals(subchannel, pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel());
 
@@ -220,8 +226,10 @@ public class PickFirstLoadBalancerTest {
 
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
-
+    inOrder.verify(mockHelper).createSubchannel(createArgsCaptor.capture());
-    inOrder.verify(mockHelper).createSubchannel(eq(servers), eq(Attributes.EMPTY));
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
+    assertThat(args.getAttributes()).isEqualTo(Attributes.EMPTY);
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
     verify(mockSubchannel).requestConnection();
 
@@ -237,9 +245,21 @@ public class PickFirstLoadBalancerTest {
   @Test
   public void nameResolutionErrorWithStateChanges() throws Exception {
     InOrder inOrder = inOrder(mockHelper);
+    loadBalancer.handleResolvedAddresses(
+        ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
+    inOrder.verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    verify(mockSubchannel).start(stateListenerCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
+
+    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class));
+
+    SubchannelStateListener stateListener = stateListenerCaptor.getValue();
+
+    stateListener.onSubchannelState(ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
+    inOrder.verify(mockHelper).updateBalancingState(
+        eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
 
-    loadBalancer.handleSubchannelState(mockSubchannel,
-        ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
     Status error = Status.NOT_FOUND.withDescription("nameResolutionError");
     loadBalancer.handleNameResolutionError(error);
     inOrder.verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
@@ -248,7 +268,6 @@ public class PickFirstLoadBalancerTest {
     assertEquals(null, pickResult.getSubchannel());
     assertEquals(error, pickResult.getStatus());
 
-    loadBalancer.handleSubchannelState(mockSubchannel, ConnectivityStateInfo.forNonError(READY));
     Status error2 = Status.NOT_FOUND.withDescription("nameResolutionError2");
     loadBalancer.handleNameResolutionError(error2);
     inOrder.verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
@@ -263,13 +282,19 @@ public class PickFirstLoadBalancerTest {
   @Test
   public void requestConnection() {
     loadBalancer.requestConnection();
-    verify(mockSubchannel, never()).requestConnection();
 
+    verify(mockSubchannel, never()).requestConnection();
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
     verify(mockSubchannel).requestConnection();
 
-    loadBalancer.handleSubchannelState(mockSubchannel, ConnectivityStateInfo.forNonError(IDLE));
+    verify(mockHelper).createSubchannel(createArgsCaptor.capture());
+    verify(mockSubchannel).start(stateListenerCaptor.capture());
+    CreateSubchannelArgs args = createArgsCaptor.getValue();
+    assertThat(args.getAddresses()).isEqualTo(servers);
+    SubchannelStateListener stateListener = stateListenerCaptor.getValue();
+
+    stateListener.onSubchannelState(ConnectivityStateInfo.forNonError(IDLE));
     verify(mockHelper).updateBalancingState(eq(IDLE), any(SubchannelPicker.class));
 
     verify(mockSubchannel).requestConnection();

core/src/test/java/io/grpc/util/ForwardingSubchannelTest.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2018 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.util;
+
+import static org.mockito.Mockito.mock;
+
+import io.grpc.ForwardingTestUtil;
+import io.grpc.LoadBalancer.Subchannel;
+import java.util.Arrays;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link ForwardingSubchannel}. */
+@RunWith(JUnit4.class)
+public class ForwardingSubchannelTest {
+  private final Subchannel mockDelegate = mock(Subchannel.class);
+
+  private final class TestSubchannel extends ForwardingSubchannel {
+    @Override
+    protected Subchannel delegate() {
+      return mockDelegate;
+    }
+  }
+
+  @Test
+  public void allMethodsForwarded() throws Exception {
+    ForwardingTestUtil.testMethodsForwarded(
+        Subchannel.class,
+        mockDelegate,
+        new TestSubchannel(),
+        Arrays.asList(Subchannel.class.getMethod("getAddresses")));
+  }
+}

core/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java

@@ -51,12 +51,14 @@ import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.Metadata;
 import io.grpc.Metadata.Key;
 import io.grpc.Status;
@@ -79,6 +81,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatchers;
 import org.mockito.Captor;
 import org.mockito.InOrder;
 import org.mockito.Mock;
@@ -94,6 +97,8 @@ public class RoundRobinLoadBalancerTest {
   private RoundRobinLoadBalancer loadBalancer;
   private final List<EquivalentAddressGroup> servers = Lists.newArrayList();
   private final Map<List<EquivalentAddressGroup>, Subchannel> subchannels = Maps.newLinkedHashMap();
+  private final Map<Subchannel, SubchannelStateListener> subchannelStateListeners =
+      Maps.newLinkedHashMap();
   private final Attributes affinity =
       Attributes.newBuilder().set(MAJOR_KEY, "I got the keys").build();
 
@@ -102,7 +107,7 @@ public class RoundRobinLoadBalancerTest {
   @Captor
   private ArgumentCaptor<ConnectivityState> stateCaptor;
   @Captor
-  private ArgumentCaptor<List<EquivalentAddressGroup>> eagListCaptor;
+  private ArgumentCaptor<CreateSubchannelArgs> createArgsCaptor;
   @Mock
   private Helper mockHelper;
 
@@ -119,17 +124,26 @@ public class RoundRobinLoadBalancerTest {
       EquivalentAddressGroup eag = new EquivalentAddressGroup(addr);
       servers.add(eag);
       Subchannel sc = mock(Subchannel.class);
-      when(sc.getAllAddresses()).thenReturn(Arrays.asList(eag));
       subchannels.put(Arrays.asList(eag), sc);
     }
 
-    when(mockHelper.createSubchannel(any(List.class), any(Attributes.class)))
+    when(mockHelper.createSubchannel(any(CreateSubchannelArgs.class)))
         .then(new Answer<Subchannel>() {
           @Override
           public Subchannel answer(InvocationOnMock invocation) throws Throwable {
-            Object[] args = invocation.getArguments();
+            CreateSubchannelArgs args = (CreateSubchannelArgs) invocation.getArguments()[0];
-            Subchannel subchannel = subchannels.get(args[0]);
+            final Subchannel subchannel = subchannels.get(args.getAddresses());
-            when(subchannel.getAttributes()).thenReturn((Attributes) args[1]);
+            when(subchannel.getAllAddresses()).thenReturn(args.getAddresses());
+            when(subchannel.getAttributes()).thenReturn(args.getAttributes());
+            doAnswer(
+                new Answer<Void>() {
+                  @Override
+                  public Void answer(InvocationOnMock invocation) throws Throwable {
+                    subchannelStateListeners.put(
+                        subchannel, (SubchannelStateListener) invocation.getArguments()[0]);
+                    return null;
+                  }
+                }).when(subchannel).start(any(SubchannelStateListener.class));
             return subchannel;
           }
         });
@@ -138,8 +152,11 @@ public class RoundRobinLoadBalancerTest {
   }
 
   @After
+  @SuppressWarnings("deprecation")
   public void tearDown() throws Exception {
     verifyNoMoreInteractions(mockArgs);
+    verify(mockHelper, never()).createSubchannel(
+        ArgumentMatchers.<List<EquivalentAddressGroup>>any(), any(Attributes.class));
   }
 
   @Test
@@ -147,12 +164,15 @@ public class RoundRobinLoadBalancerTest {
     final Subchannel readySubchannel = subchannels.values().iterator().next();
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
-    loadBalancer.handleSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));
 
-    verify(mockHelper, times(3)).createSubchannel(eagListCaptor.capture(),
+    verify(mockHelper, times(3)).createSubchannel(createArgsCaptor.capture());
-        any(Attributes.class));
+    List<List<EquivalentAddressGroup>> capturedAddrs = new ArrayList<>();
+    for (CreateSubchannelArgs arg : createArgsCaptor.getAllValues()) {
+      capturedAddrs.add(arg.getAddresses());
+    }
 
-    assertThat(eagListCaptor.getAllValues()).containsAtLeastElementsIn(subchannels.keySet());
+    assertThat(capturedAddrs).containsAtLeastElementsIn(subchannels.keySet());
     for (Subchannel subchannel : subchannels.values()) {
       verify(subchannel).requestConnection();
       verify(subchannel, never()).shutdown();
@@ -187,36 +207,26 @@ public class RoundRobinLoadBalancerTest {
       Subchannel subchannel = allSubchannels.get(i);
       List<EquivalentAddressGroup> eagList =
           Arrays.asList(new EquivalentAddressGroup(allAddrs.get(i)));
-      when(subchannel.getAttributes()).thenReturn(Attributes.newBuilder().set(STATE_INFO,
+      subchannels.put(eagList, subchannel);
-          new Ref<>(
-              ConnectivityStateInfo.forNonError(READY))).build());
-      when(subchannel.getAllAddresses()).thenReturn(eagList);
     }
 
-    final Map<List<EquivalentAddressGroup>, Subchannel> subchannels2 = Maps.newHashMap();
-    subchannels2.put(Arrays.asList(new EquivalentAddressGroup(removedAddr)), removedSubchannel);
-    subchannels2.put(Arrays.asList(new EquivalentAddressGroup(oldAddr)), oldSubchannel);
-
     List<EquivalentAddressGroup> currentServers =
         Lists.newArrayList(
             new EquivalentAddressGroup(removedAddr),
             new EquivalentAddressGroup(oldAddr));
 
-    doAnswer(new Answer<Subchannel>() {
+    InOrder inOrder = inOrder(mockHelper);
-      @Override
-      public Subchannel answer(InvocationOnMock invocation) throws Throwable {
-        Object[] args = invocation.getArguments();
-        return subchannels2.get(args[0]);
-      }
-    }).when(mockHelper).createSubchannel(any(List.class), any(Attributes.class));
 
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(currentServers).setAttributes(affinity)
             .build());
 
-    InOrder inOrder = inOrder(mockHelper);
+    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
 
-    inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    deliverSubchannelState(removedSubchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(oldSubchannel, ConnectivityStateInfo.forNonError(READY));
+
+    inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
     SubchannelPicker picker = pickerCaptor.getValue();
     assertThat(getList(picker)).containsExactly(removedSubchannel, oldSubchannel);
 
@@ -226,10 +236,6 @@ public class RoundRobinLoadBalancerTest {
     assertThat(loadBalancer.getSubchannels()).containsExactly(removedSubchannel,
         oldSubchannel);
 
-    subchannels2.clear();
-    subchannels2.put(Arrays.asList(new EquivalentAddressGroup(oldAddr)), oldSubchannel);
-    subchannels2.put(Arrays.asList(new EquivalentAddressGroup(newAddr)), newSubchannel);
-
     List<EquivalentAddressGroup> latestServers =
         Lists.newArrayList(
             new EquivalentAddressGroup(oldAddr),
@@ -241,14 +247,14 @@ public class RoundRobinLoadBalancerTest {
     verify(newSubchannel, times(1)).requestConnection();
     verify(removedSubchannel, times(1)).shutdown();
 
-    loadBalancer.handleSubchannelState(removedSubchannel,
+    deliverSubchannelState(removedSubchannel, ConnectivityStateInfo.forNonError(SHUTDOWN));
-            ConnectivityStateInfo.forNonError(SHUTDOWN));
+    deliverSubchannelState(newSubchannel, ConnectivityStateInfo.forNonError(READY));
 
     assertThat(loadBalancer.getSubchannels()).containsExactly(oldSubchannel,
         newSubchannel);
 
-    verify(mockHelper, times(3)).createSubchannel(any(List.class), any(Attributes.class));
+    verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
-    inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
 
     picker = pickerCaptor.getValue();
     assertThat(getList(picker)).containsExactly(oldSubchannel, newSubchannel);
@@ -280,7 +286,7 @@ public class RoundRobinLoadBalancerTest {
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), isA(EmptyPicker.class));
     assertThat(subchannelStateInfo.value).isEqualTo(ConnectivityStateInfo.forNonError(IDLE));
 
-    loadBalancer.handleSubchannelState(subchannel,
+    deliverSubchannelState(subchannel,
         ConnectivityStateInfo.forNonError(READY));
     inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
     assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class);
@@ -288,20 +294,20 @@ public class RoundRobinLoadBalancerTest {
         ConnectivityStateInfo.forNonError(READY));
 
     Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯");
-    loadBalancer.handleSubchannelState(subchannel,
+    deliverSubchannelState(subchannel,
         ConnectivityStateInfo.forTransientFailure(error));
     assertThat(subchannelStateInfo.value).isEqualTo(
         ConnectivityStateInfo.forTransientFailure(error));
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
     assertThat(pickerCaptor.getValue()).isInstanceOf(EmptyPicker.class);
 
-    loadBalancer.handleSubchannelState(subchannel,
+    deliverSubchannelState(subchannel,
         ConnectivityStateInfo.forNonError(IDLE));
     assertThat(subchannelStateInfo.value).isEqualTo(
         ConnectivityStateInfo.forNonError(IDLE));
 
     verify(subchannel, times(2)).requestConnection();
-    verify(mockHelper, times(3)).createSubchannel(any(List.class), any(Attributes.class));
+    verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
     verifyNoMoreInteractions(mockHelper);
   }
 
@@ -353,10 +359,10 @@ public class RoundRobinLoadBalancerTest {
     final Subchannel readySubchannel = subchannels.values().iterator().next();
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity).build());
-    loadBalancer.handleSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));
     loadBalancer.handleNameResolutionError(Status.NOT_FOUND.withDescription("nameResolutionError"));
 
-    verify(mockHelper, times(3)).createSubchannel(any(List.class), any(Attributes.class));
+    verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
     verify(mockHelper, times(3))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
 
@@ -388,12 +394,11 @@ public class RoundRobinLoadBalancerTest {
     verify(sc2, times(1)).requestConnection();
     verify(sc3, times(1)).requestConnection();
 
-    loadBalancer.handleSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
-    loadBalancer.handleSubchannelState(sc2, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(sc2, ConnectivityStateInfo.forNonError(READY));
-    loadBalancer.handleSubchannelState(sc3, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(sc3, ConnectivityStateInfo.forNonError(READY));
-    loadBalancer.handleSubchannelState(sc2, ConnectivityStateInfo.forNonError(IDLE));
+    deliverSubchannelState(sc2, ConnectivityStateInfo.forNonError(IDLE));
-    loadBalancer
+    deliverSubchannelState(sc3, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
-        .handleSubchannelState(sc3, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
 
     verify(mockHelper, times(6))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -426,7 +431,7 @@ public class RoundRobinLoadBalancerTest {
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY)
             .build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(any(ConnectivityState.class), pickerCaptor.capture());
@@ -460,7 +465,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -493,7 +498,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -524,7 +529,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -570,7 +575,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -585,8 +590,7 @@ public class RoundRobinLoadBalancerTest {
     Subchannel sc1 = picker.pickSubchannel(mockArgs).getSubchannel();
 
     // go to transient failure
-    loadBalancer
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
-        .handleSubchannelState(sc1, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
 
     verify(mockHelper, times(5))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -596,7 +600,7 @@ public class RoundRobinLoadBalancerTest {
     Subchannel sc2 = picker.pickSubchannel(mockArgs).getSubchannel();
 
     // go back to ready
-    loadBalancer.handleSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
 
     verify(mockHelper, times(6))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -619,7 +623,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -634,8 +638,7 @@ public class RoundRobinLoadBalancerTest {
     Subchannel sc1 = picker.pickSubchannel(mockArgs).getSubchannel();
 
     // go to transient failure
-    loadBalancer
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
-        .handleSubchannelState(sc1, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
 
     Metadata headerWithStickinessValue2 = new Metadata();
     headerWithStickinessValue2.put(stickinessKey, "my-sticky-value2");
@@ -649,7 +652,7 @@ public class RoundRobinLoadBalancerTest {
     Subchannel sc2 = picker.pickSubchannel(mockArgs).getSubchannel();
 
     // go back to ready
-    loadBalancer.handleSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forNonError(READY));
 
     doReturn(headerWithStickinessValue1).when(mockArgs).getHeaders();
     verify(mockHelper, times(6))
@@ -674,7 +677,7 @@ public class RoundRobinLoadBalancerTest {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(attributes).build());
     for (Subchannel subchannel : subchannels.values()) {
-      loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     }
     verify(mockHelper, times(4))
         .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());
@@ -690,8 +693,7 @@ public class RoundRobinLoadBalancerTest {
     Subchannel sc1 = picker.pickSubchannel(mockArgs).getSubchannel();
 
     // shutdown channel directly
-    loadBalancer
+    deliverSubchannelState(sc1, ConnectivityStateInfo.forNonError(ConnectivityState.SHUTDOWN));
-        .handleSubchannelState(sc1, ConnectivityStateInfo.forNonError(ConnectivityState.SHUTDOWN));
 
     assertNull(loadBalancer.getStickinessMapForTest().get("my-sticky-value").value);
 
@@ -713,7 +715,7 @@ public class RoundRobinLoadBalancerTest {
 
     verify(sc2, times(1)).shutdown();
 
-    loadBalancer.handleSubchannelState(sc2, ConnectivityStateInfo.forNonError(SHUTDOWN));
+    deliverSubchannelState(sc2, ConnectivityStateInfo.forNonError(SHUTDOWN));
 
     assertNull(loadBalancer.getStickinessMapForTest().get("my-sticky-value").value);
 
@@ -799,6 +801,10 @@ public class RoundRobinLoadBalancerTest {
         Collections.<Subchannel>emptyList();
   }
 
+  private void deliverSubchannelState(Subchannel subchannel, ConnectivityStateInfo newState) {
+    subchannelStateListeners.get(subchannel).onSubchannelState(newState);
+  }
+
   private static class FakeSocketAddress extends SocketAddress {
     final String name;
 

grpclb/src/main/java/io/grpc/grpclb/CachedSubchannelPool.java

@@ -51,11 +51,14 @@ final class CachedSubchannelPool implements SubchannelPool {
   }
 
   @Override
+  @SuppressWarnings("deprecation")
   public Subchannel takeOrCreateSubchannel(
       EquivalentAddressGroup eag, Attributes defaultAttributes) {
     final CacheEntry entry = cache.remove(eag);
     final Subchannel subchannel;
     if (entry == null) {
+      // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to the
+      // new createSubchannel().
       subchannel = helper.createSubchannel(eag, defaultAttributes);
     } else {
       subchannel = entry.subchannel;

grpclb/src/main/java/io/grpc/grpclb/GrpclbLoadBalancer.java

@@ -80,6 +80,7 @@ class GrpclbLoadBalancer extends LoadBalancer {
     checkNotNull(grpclbState, "grpclbState");
   }
 
+  @Deprecated
   @Override
   public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo newState) {
     // grpclbState should never be null here since handleSubchannelState cannot be called while the

grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java

@@ -371,6 +371,7 @@ final class GrpclbState {
   /**
    * Populate the round-robin lists with the given values.
    */
+  @SuppressWarnings("deprecation")
   private void useRoundRobinLists(
       List<DropEntry> newDropList, List<BackendAddressGroup> newBackendAddrList,
       @Nullable GrpclbClientLoadRecorder loadRecorder) {
@@ -430,6 +431,8 @@ final class GrpclbState {
         }
         Subchannel subchannel;
         if (subchannels.isEmpty()) {
+          // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+          // the new createSubchannel().
           subchannel = helper.createSubchannel(eagList, createSubchannelAttrs());
         } else {
           checkState(subchannels.size() == 1, "Unexpected Subchannel count: %s", subchannels);

grpclb/src/test/java/io/grpc/grpclb/CachedSubchannelPoolTest.java

@@ -93,7 +93,7 @@ public class CachedSubchannelPoolTest {
   private final ArrayList<Subchannel> mockSubchannels = new ArrayList<>();
 
   @Before
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "deprecation"})
   public void setUp() {
     doAnswer(new Answer<Subchannel>() {
         @Override
@@ -107,6 +107,8 @@ public class CachedSubchannelPoolTest {
           mockSubchannels.add(subchannel);
           return subchannel;
         }
+        // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+        // the new createSubchannel().
       }).when(helper).createSubchannel(any(List.class), any(Attributes.class));
     doAnswer(new Answer<Void>() {
         @Override
@@ -122,20 +124,26 @@ public class CachedSubchannelPoolTest {
   }
 
   @After
+  @SuppressWarnings("deprecation")
   public void wrapUp() {
     // Sanity checks
     for (Subchannel subchannel : mockSubchannels) {
       verify(subchannel, atMost(1)).shutdown();
     }
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+    // the new API.
     verify(balancer, atLeast(0))
         .handleSubchannelState(any(Subchannel.class), any(ConnectivityStateInfo.class));
     verifyNoMoreInteractions(balancer);
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void subchannelExpireAfterReturned() {
     Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
     assertThat(subchannel1).isNotNull();
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to the new
+    // createSubchannel().
     verify(helper).createSubchannel(eq(Arrays.asList(EAG1)), same(ATTRS1));
 
     Subchannel subchannel2 = pool.takeOrCreateSubchannel(EAG2, ATTRS2);
@@ -162,10 +170,13 @@ public class CachedSubchannelPoolTest {
     assertThat(clock.numPendingTasks()).isEqualTo(0);
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void subchannelReused() {
     Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
     assertThat(subchannel1).isNotNull();
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to the new
+    // createSubchannel().
     verify(helper).createSubchannel(eq(Arrays.asList(EAG1)), same(ATTRS1));
 
     Subchannel subchannel2 = pool.takeOrCreateSubchannel(EAG2, ATTRS2);
@@ -205,6 +216,7 @@ public class CachedSubchannelPoolTest {
     assertThat(clock.numPendingTasks()).isEqualTo(0);
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void updateStateWhileInPool() {
     Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
@@ -217,6 +229,8 @@ public class CachedSubchannelPoolTest {
 
     pool.handleSubchannelState(subchannel1, anotherFailureState);
 
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to the new
+    // createSubchannel().
     verify(balancer, never())
         .handleSubchannelState(any(Subchannel.class), any(ConnectivityStateInfo.class));
 
@@ -229,11 +243,14 @@ public class CachedSubchannelPoolTest {
     verifyNoMoreInteractions(balancer);
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void updateStateWhileInPool_notSameObject() {
     Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
     pool.returnSubchannel(subchannel1, READY_STATE);
 
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to the new
+    // createSubchannel().
     Subchannel subchannel2 = helper.createSubchannel(EAG1, ATTRS1);
     Subchannel subchannel3 = helper.createSubchannel(EAG2, ATTRS2);
 

grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java

@@ -193,7 +193,7 @@ public class GrpclbLoadBalancerTest {
   private BackoffPolicy backoffPolicy2;
   private GrpclbLoadBalancer balancer;
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "deprecation"})
   @Before
   public void setUp() throws Exception {
     MockitoAnnotations.initMocks(this);
@@ -262,6 +262,8 @@ public class GrpclbLoadBalancerTest {
           unpooledSubchannelTracker.add(subchannel);
           return subchannel;
         }
+        // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+        // the new createSubchannel().
       }).when(helper).createSubchannel(any(List.class), any(Attributes.class));
     when(helper.getSynchronizationContext()).thenReturn(syncContext);
     when(helper.getScheduledExecutorService()).thenReturn(fakeClock.getScheduledExecutorService());
@@ -1680,7 +1682,7 @@ public class GrpclbLoadBalancerTest {
     verify(helper, times(4)).refreshNameResolution();
   }
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "deprecation"})
   @Test
   public void grpclbWorking_pickFirstMode() throws Exception {
     InOrder inOrder = inOrder(helper);
@@ -1711,6 +1713,8 @@ public class GrpclbLoadBalancerTest {
     lbResponseObserver.onNext(buildInitialResponse());
     lbResponseObserver.onNext(buildLbResponse(backends1));
 
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+    // the new createSubchannel().
     inOrder.verify(helper).createSubchannel(
         eq(Arrays.asList(
                 new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
@@ -1804,7 +1808,7 @@ public class GrpclbLoadBalancerTest {
         .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class));
   }
 
-  @SuppressWarnings("unchecked")
+  @SuppressWarnings({"unchecked", "deprecation"})
   @Test
   public void pickFirstMode_fallback() throws Exception {
     InOrder inOrder = inOrder(helper);
@@ -1828,6 +1832,8 @@ public class GrpclbLoadBalancerTest {
     fakeClock.forwardTime(GrpclbState.FALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
 
     // Entering fallback mode
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+    // the new createSubchannel().
     inOrder.verify(helper).createSubchannel(
         eq(Arrays.asList(grpclbResolutionList.get(0), grpclbResolutionList.get(2))),
         any(Attributes.class));
@@ -1883,6 +1889,7 @@ public class GrpclbLoadBalancerTest {
         .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class));
   }
 
+  @SuppressWarnings("deprecation")
   @Test
   public void switchMode() throws Exception {
     InOrder inOrder = inOrder(helper);
@@ -1960,6 +1967,8 @@ public class GrpclbLoadBalancerTest {
     lbResponseObserver.onNext(buildLbResponse(backends1));
 
     // PICK_FIRST Subchannel
+    // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+    // the new createSubchannel().
     inOrder.verify(helper).createSubchannel(
         eq(Arrays.asList(
                 new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
@@ -2044,11 +2053,14 @@ public class GrpclbLoadBalancerTest {
     assertThat(mode).isEqualTo(Mode.ROUND_ROBIN);
   }
 
+  @SuppressWarnings("deprecation")
   private void deliverSubchannelState(
       final Subchannel subchannel, final ConnectivityStateInfo newState) {
     syncContext.execute(new Runnable() {
         @Override
         public void run() {
+          // TODO(zhangkun83): remove the deprecation suppression on this method once migrated to
+          // the new API.
           balancer.handleSubchannelState(subchannel, newState);
         }
       });

services/src/main/java/io/grpc/services/HealthCheckingLoadBalancerFactory.java

@@ -23,21 +23,22 @@ import static io.grpc.ConnectivityState.IDLE;
 import static io.grpc.ConnectivityState.READY;
 import static io.grpc.ConnectivityState.SHUTDOWN;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Objects;
 import com.google.common.base.Stopwatch;
 import com.google.common.base.Supplier;
-import io.grpc.Attributes;
 import io.grpc.CallOptions;
 import io.grpc.ChannelLogger;
 import io.grpc.ChannelLogger.ChannelLogLevel;
 import io.grpc.ClientCall;
 import io.grpc.ConnectivityStateInfo;
-import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Factory;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.Subchannel;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.Status.Code;
@@ -52,8 +53,8 @@ import io.grpc.internal.GrpcAttributes;
 import io.grpc.internal.ServiceConfigUtil;
 import io.grpc.util.ForwardingLoadBalancer;
 import io.grpc.util.ForwardingLoadBalancerHelper;
+import io.grpc.util.ForwardingSubchannel;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -70,8 +71,6 @@ import javax.annotation.Nullable;
  * SynchronizationContext, or it will throw.
  */
 final class HealthCheckingLoadBalancerFactory extends Factory {
-  private static final Attributes.Key<HealthCheckState> KEY_HEALTH_CHECK_STATE =
-      Attributes.Key.create("io.grpc.services.HealthCheckingLoadBalancerFactory.healthCheckState");
   private static final Logger logger =
       Logger.getLogger(HealthCheckingLoadBalancerFactory.class.getName());
 
@@ -91,7 +90,6 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
   public LoadBalancer newLoadBalancer(Helper helper) {
     HelperImpl wrappedHelper = new HelperImpl(helper);
     LoadBalancer delegateBalancer = delegateFactory.newLoadBalancer(wrappedHelper);
-    wrappedHelper.init(delegateBalancer);
     return new HealthCheckingLoadBalancer(wrappedHelper, delegateBalancer);
   }
 
@@ -99,7 +97,6 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
     private final Helper delegate;
     private final SynchronizationContext syncContext;
 
-    private LoadBalancer delegateBalancer;
     @Nullable String healthCheckedService;
     private boolean balancerShutdown;
 
@@ -110,27 +107,21 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
       this.syncContext = checkNotNull(delegate.getSynchronizationContext(), "syncContext");
     }
 
-    void init(LoadBalancer delegateBalancer) {
-      checkState(this.delegateBalancer == null, "init() already called");
-      this.delegateBalancer = checkNotNull(delegateBalancer, "delegateBalancer");
-    }
-
     @Override
     protected Helper delegate() {
       return delegate;
     }
 
     @Override
-    public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
+    public Subchannel createSubchannel(CreateSubchannelArgs args) {
       // HealthCheckState is not thread-safe, we are requiring the original LoadBalancer calls
       // createSubchannel() from the SynchronizationContext.
       syncContext.throwIfNotInThisSynchronizationContext();
+      Subchannel originalSubchannel = super.createSubchannel(args);
       HealthCheckState hcState = new HealthCheckState(
-          this, delegateBalancer, syncContext, delegate.getScheduledExecutorService());
+          this, originalSubchannel, syncContext, delegate.getScheduledExecutorService());
       hcStates.add(hcState);
-      Subchannel subchannel = super.createSubchannel(
+      Subchannel subchannel = new SubchannelImpl(originalSubchannel, hcState);
-          addrs, attrs.toBuilder().set(KEY_HEALTH_CHECK_STATE, hcState).build());
-      hcState.init(subchannel);
       if (healthCheckedService != null) {
         hcState.setServiceName(healthCheckedService);
       }
@@ -150,6 +141,28 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
     }
   }
 
+  @VisibleForTesting
+  static final class SubchannelImpl extends ForwardingSubchannel {
+    final Subchannel delegate;
+    final HealthCheckState hcState;
+
+    SubchannelImpl(Subchannel delegate, HealthCheckState hcState) {
+      this.delegate = checkNotNull(delegate, "delegate");
+      this.hcState = checkNotNull(hcState, "hcState");
+    }
+
+    @Override
+    protected Subchannel delegate() {
+      return delegate;
+    }
+
+    @Override
+    public void start(final SubchannelStateListener listener) {
+      hcState.init(listener);
+      delegate().start(hcState);
+    }
+  }
+
   private static final class HealthCheckingLoadBalancer extends ForwardingLoadBalancer {
     final LoadBalancer delegate;
     final HelperImpl helper;
@@ -177,27 +190,15 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
       super.handleResolvedAddresses(resolvedAddresses);
     }
 
-    @Override
-    public void handleSubchannelState(
-        Subchannel subchannel, ConnectivityStateInfo stateInfo) {
-      HealthCheckState hcState =
-          checkNotNull(subchannel.getAttributes().get(KEY_HEALTH_CHECK_STATE), "hcState");
-      hcState.updateRawState(stateInfo);
-
-      if (Objects.equal(stateInfo.getState(), SHUTDOWN)) {
-        helper.hcStates.remove(hcState);
-      }
-    }
-
     @Override
     public void shutdown() {
       super.shutdown();
       helper.balancerShutdown = true;
       for (HealthCheckState hcState : helper.hcStates) {
-        // ManagedChannel will stop calling handleSubchannelState() after shutdown() is called,
+        // ManagedChannel will stop calling onSubchannelState() after shutdown() is called,
         // which is required by LoadBalancer API semantics. We need to deliver the final SHUTDOWN
         // signal to health checkers so that they can cancel the streams.
-        hcState.updateRawState(ConnectivityStateInfo.forNonError(SHUTDOWN));
+        hcState.onSubchannelState(ConnectivityStateInfo.forNonError(SHUTDOWN));
       }
       helper.hcStates.clear();
     }
@@ -210,7 +211,7 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
 
   
   // All methods are run from syncContext
-  private final class HealthCheckState {
+  private final class HealthCheckState implements SubchannelStateListener {
     private final Runnable retryTask = new Runnable() {
         @Override
         public void run() {
@@ -218,13 +219,12 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
         }
       };
 
-    private final LoadBalancer delegate;
     private final SynchronizationContext syncContext;
     private final ScheduledExecutorService timerService;
     private final HelperImpl helperImpl;
-
+    private final Subchannel subchannel;
-    private Subchannel subchannel;
+    private final ChannelLogger subchannelLogger;
-    private ChannelLogger subchannelLogger;
+    private SubchannelStateListener stateListener;
 
     // Set when RPC started. Cleared when the RPC has closed or abandoned.
     @Nullable
@@ -246,18 +246,18 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
 
     HealthCheckState(
         HelperImpl helperImpl,
-        LoadBalancer delegate, SynchronizationContext syncContext,
+        Subchannel subchannel, SynchronizationContext syncContext,
         ScheduledExecutorService timerService) {
       this.helperImpl = checkNotNull(helperImpl, "helperImpl");
-      this.delegate = checkNotNull(delegate, "delegate");
+      this.subchannel = checkNotNull(subchannel, "subchannel");
+      this.subchannelLogger = checkNotNull(subchannel.getChannelLogger(), "subchannelLogger");
       this.syncContext = checkNotNull(syncContext, "syncContext");
       this.timerService = checkNotNull(timerService, "timerService");
     }
 
-    void init(Subchannel subchannel) {
+    void init(SubchannelStateListener listener) {
-      checkState(this.subchannel == null, "init() already called");
+      checkState(this.stateListener == null, "init() already called");
-      this.subchannel = checkNotNull(subchannel, "subchannel");
+      this.stateListener = checkNotNull(listener, "listener");
-      this.subchannelLogger = checkNotNull(subchannel.getChannelLogger(), "subchannelLogger");
     }
 
     void setServiceName(@Nullable String newServiceName) {
@@ -274,13 +274,17 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
       adjustHealthCheck();
     }
 
-    void updateRawState(ConnectivityStateInfo rawState) {
+    @Override
+    public void onSubchannelState(ConnectivityStateInfo rawState) {
       if (Objects.equal(this.rawState.getState(), READY)
           && !Objects.equal(rawState.getState(), READY)) {
         // A connection was lost.  We will reset disabled flag because health check
         // may be available on the new connection.
         disabled = false;
       }
+      if (Objects.equal(rawState.getState(), SHUTDOWN)) {
+        helperImpl.hcStates.remove(this);
+      }
       this.rawState = rawState;
       adjustHealthCheck();
     }
@@ -339,7 +343,7 @@ final class HealthCheckingLoadBalancerFactory extends Factory {
       checkState(subchannel != null, "init() not called");
       if (!helperImpl.balancerShutdown && !Objects.equal(concludedState, newState)) {
         concludedState = newState;
-        delegate.handleSubchannelState(subchannel, concludedState);
+        stateListener.onSubchannelState(concludedState);
       }
     }
 

services/src/test/java/io/grpc/services/HealthCheckingLoadBalancerFactoryTest.java

@@ -24,7 +24,6 @@ import static io.grpc.ConnectivityState.IDLE;
 import static io.grpc.ConnectivityState.READY;
 import static io.grpc.ConnectivityState.SHUTDOWN;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
-import static org.junit.Assert.fail;
 import static org.mockito.AdditionalAnswers.delegatesTo;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -32,6 +31,7 @@ import static org.mockito.ArgumentMatchers.same;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.verifyZeroInteractions;
@@ -47,11 +47,13 @@ import io.grpc.Context;
 import io.grpc.Context.CancellationListener;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Factory;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.ManagedChannel;
 import io.grpc.NameResolver;
 import io.grpc.Server;
@@ -68,13 +70,13 @@ import io.grpc.internal.BackoffPolicy;
 import io.grpc.internal.FakeClock;
 import io.grpc.internal.GrpcAttributes;
 import io.grpc.internal.ServiceConfigUtil;
+import io.grpc.services.HealthCheckingLoadBalancerFactory.SubchannelImpl;
 import io.grpc.stub.StreamObserver;
 import java.net.SocketAddress;
 import java.text.MessageFormat;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Queue;
@@ -108,6 +110,8 @@ public class HealthCheckingLoadBalancerFactoryTest {
   private final EquivalentAddressGroup[] eags = new EquivalentAddressGroup[NUM_SUBCHANNELS];
   @SuppressWarnings({"rawtypes", "unchecked"})
   private final List<EquivalentAddressGroup>[] eagLists = new List[NUM_SUBCHANNELS];
+  private final SubchannelStateListener[] mockStateListeners =
+      new SubchannelStateListener[NUM_SUBCHANNELS];
   private List<EquivalentAddressGroup> resolvedAddressList;
   private final FakeSubchannel[] subchannels = new FakeSubchannel[NUM_SUBCHANNELS];
   private final ManagedChannel[] channels = new ManagedChannel[NUM_SUBCHANNELS];
@@ -139,7 +143,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
   private LoadBalancer origLb;
   private LoadBalancer hcLb;
   @Captor
-  ArgumentCaptor<Attributes> attrsCaptor;
+  ArgumentCaptor<CreateSubchannelArgs> createArgsCaptor;
   @Mock
   private BackoffPolicy.Provider backoffPolicyProvider;
   @Mock
@@ -171,6 +175,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
       eags[i] = eag;
       List<EquivalentAddressGroup> eagList = Arrays.asList(eag);
       eagLists[i] = eagList;
+      mockStateListeners[i] = mock(SubchannelStateListener.class);
     }
     resolvedAddressList = Arrays.asList(eags);
 
@@ -199,19 +204,6 @@ public class HealthCheckingLoadBalancerFactoryTest {
             });
         }
 
-        @Override
-        public void handleSubchannelState(
-            final Subchannel subchannel, final ConnectivityStateInfo stateInfo) {
-          syncContext.execute(new Runnable() {
-              @Override
-              public void run() {
-                if (!shutdown) {
-                  hcLb.handleSubchannelState(subchannel, stateInfo);
-                }
-              }
-            });
-        }
-
         @Override
         public void handleNameResolutionError(Status error) {
           throw new AssertionError("Not supposed to be called");
@@ -237,7 +229,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
   public void teardown() throws Exception {
     // All scheduled tasks have been accounted for
     assertThat(clock.getPendingTasks()).isEmpty();
-    // Health-check streams are usually not closed in the tests because handleSubchannelState() is
+    // Health-check streams are usually not closed in the tests because onSubchannelState() is
     // faked.  Force closing for clean up.
     for (Server server : servers) {
       server.shutdownNow();
@@ -252,16 +244,6 @@ public class HealthCheckingLoadBalancerFactoryTest {
     }
   }
 
-  @Test
-  public void createSubchannelThrowsIfCalledOutsideSynchronizationContext() {
-    try {
-      wrappedHelper.createSubchannel(eagLists[0], Attributes.EMPTY);
-      fail("Should throw");
-    } catch (IllegalStateException e) {
-      assertThat(e.getMessage()).isEqualTo("Not called from the SynchronizationContext");
-    }
-  }
-
   @Test
   public void typicalWorkflow() {
     Attributes resolutionAttrs = attrsWithHealthCheckService("FooService");
@@ -284,42 +266,42 @@ public class HealthCheckingLoadBalancerFactoryTest {
       Attributes attrs = Attributes.newBuilder()
           .set(SUBCHANNEL_ATTR_KEY, subchannelAttrValue).build();
       // We don't wrap Subchannels, thus origLb gets the original Subchannels.
-      assertThat(createSubchannel(i, attrs)).isSameInstanceAs(subchannels[i]);
+      assertThat(unwrap(createSubchannel(i, attrs))).isSameInstanceAs(subchannels[i]);
-      verify(origHelper).createSubchannel(same(eagLists[i]), attrsCaptor.capture());
+      verify(origHelper, times(i + 1)).createSubchannel(createArgsCaptor.capture());
-      assertThat(attrsCaptor.getValue().get(SUBCHANNEL_ATTR_KEY)).isEqualTo(subchannelAttrValue);
+      assertThat(createArgsCaptor.getValue().getAddresses()).isEqualTo(eagLists[i]);
+      assertThat(createArgsCaptor.getValue().getAttributes().get(SUBCHANNEL_ATTR_KEY))
+          .isEqualTo(subchannelAttrValue);
     }
 
     for (int i = NUM_SUBCHANNELS - 1; i >= 0; i--) {
       // Not starting health check until underlying Subchannel is READY
       FakeSubchannel subchannel = subchannels[i];
       HealthImpl healthImpl = healthImpls[i];
-      InOrder inOrder = inOrder(origLb);
+      SubchannelStateListener mockStateListener = mockStateListeners[i];
-      hcLbEventDelivery.handleSubchannelState(
+      InOrder inOrder = inOrder(mockStateListener);
-          subchannel, ConnectivityStateInfo.forNonError(CONNECTING));
+      deliverSubchannelState(i, ConnectivityStateInfo.forNonError(CONNECTING));
-      hcLbEventDelivery.handleSubchannelState(
+      deliverSubchannelState(i, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
-          subchannel, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE));
+      deliverSubchannelState(i, ConnectivityStateInfo.forNonError(IDLE));
-      hcLbEventDelivery.handleSubchannelState(
-          subchannel, ConnectivityStateInfo.forNonError(IDLE));
 
-      inOrder.verify(origLb).handleSubchannelState(
+      inOrder.verify(mockStateListener).onSubchannelState(
-          same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+          eq(ConnectivityStateInfo.forNonError(CONNECTING)));
-      inOrder.verify(origLb).handleSubchannelState(
+      inOrder.verify(mockStateListener).onSubchannelState(
-          same(subchannel), eq(ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)));
+          eq(ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)));
-      inOrder.verify(origLb).handleSubchannelState(
+      inOrder.verify(mockStateListener).onSubchannelState(
-          same(subchannel), eq(ConnectivityStateInfo.forNonError(IDLE)));
+          eq(ConnectivityStateInfo.forNonError(IDLE)));
-      verifyNoMoreInteractions(origLb);
+      verifyNoMoreInteractions(mockStateListener);
 
       assertThat(subchannel.logs).isEmpty();
       assertThat(healthImpl.calls).isEmpty();
-      hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      deliverSubchannelState(i, ConnectivityStateInfo.forNonError(READY));
       assertThat(healthImpl.calls).hasSize(1);
       ServerSideCall serverCall = healthImpl.calls.peek();
       assertThat(serverCall.request).isEqualTo(makeRequest("FooService"));
 
       // Starting the health check will make the Subchannel appear CONNECTING to the origLb.
-      inOrder.verify(origLb).handleSubchannelState(
+      inOrder.verify(mockStateListener).onSubchannelState(
-          same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+          eq(ConnectivityStateInfo.forNonError(CONNECTING)));
-      verifyNoMoreInteractions(origLb);
+      verifyNoMoreInteractions(mockStateListener);
 
       assertThat(subchannel.logs).containsExactly(
           "INFO: CONNECTING: Starting health-check for \"FooService\"");
@@ -333,36 +315,37 @@ public class HealthCheckingLoadBalancerFactoryTest {
         serverCall.responseObserver.onNext(makeResponse(servingStatus));
         // SERVING is mapped to READY, while other statuses are mapped to TRANSIENT_FAILURE
         if (servingStatus == ServingStatus.SERVING) {
-          inOrder.verify(origLb).handleSubchannelState(
+          inOrder.verify(mockStateListener).onSubchannelState(
-              same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+              eq(ConnectivityStateInfo.forNonError(READY)));
           assertThat(subchannel.logs).containsExactly(
               "INFO: READY: health-check responded SERVING");
         } else {
-          inOrder.verify(origLb).handleSubchannelState(
+          inOrder.verify(mockStateListener).onSubchannelState(
-              same(subchannel),unavailableStateWithMsg(
+              unavailableStateWithMsg(
                   "Health-check service responded " + servingStatus + " for 'FooService'"));
           assertThat(subchannel.logs).containsExactly(
               "INFO: TRANSIENT_FAILURE: health-check responded " + servingStatus);
         }
         subchannel.logs.clear();
-        verifyNoMoreInteractions(origLb);
+        verifyNoMoreInteractions(mockStateListener);
       }
     }
 
     // origLb shuts down Subchannels
     for (int i = 0; i < NUM_SUBCHANNELS; i++) {
       FakeSubchannel subchannel = subchannels[i];
+      SubchannelStateListener mockStateListener = mockStateListeners[i];
 
       ServerSideCall serverCall = healthImpls[i].calls.peek();
       assertThat(serverCall.cancelled).isFalse();
-      verifyNoMoreInteractions(origLb);
+      verifyNoMoreInteractions(mockStateListener);
 
       // Subchannel enters SHUTDOWN state as a response to shutdown(), and that will cancel the
       // health check RPC
       subchannel.shutdown();
       assertThat(serverCall.cancelled).isTrue();
-      verify(origLb).handleSubchannelState(
+      verify(mockStateListener).onSubchannelState(
-          same(subchannel), eq(ConnectivityStateInfo.forNonError(SHUTDOWN)));
+          eq(ConnectivityStateInfo.forNonError(SHUTDOWN)));
       assertThat(subchannel.logs).isEmpty();
     }
 
@@ -390,14 +373,13 @@ public class HealthCheckingLoadBalancerFactoryTest {
       createSubchannel(i, Attributes.EMPTY);
     }
 
-    InOrder inOrder = inOrder(origLb);
+    InOrder inOrder = inOrder(mockStateListeners[0], mockStateListeners[1]);
 
     for (int i = 0; i < 2; i++) {
-      hcLbEventDelivery.handleSubchannelState(
+      deliverSubchannelState(i, ConnectivityStateInfo.forNonError(READY));
-          subchannels[i], ConnectivityStateInfo.forNonError(READY));
       assertThat(healthImpls[i].calls).hasSize(1);
-      inOrder.verify(origLb).handleSubchannelState(
+      inOrder.verify(mockStateListeners[i]).onSubchannelState(
-          same(subchannels[i]), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+          eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     }
 
     ServerSideCall serverCall0 = healthImpls[0].calls.poll();
@@ -409,40 +391,37 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // In reality UNIMPLEMENTED is generated by GRPC server library, but the client can't tell
     // whether it's the server library or the service implementation that returned this status.
     serverCall0.responseObserver.onError(Status.UNIMPLEMENTED.asException());
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannels[0]), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
     assertThat(subchannels[0].logs).containsExactly(
         "ERROR: Health-check disabled: " + Status.UNIMPLEMENTED,
         "INFO: READY (no health-check)").inOrder();
 
     // subchannels[1] has normal health checking
     serverCall1.responseObserver.onNext(makeResponse(ServingStatus.NOT_SERVING));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[1]).onSubchannelState(
-        same(subchannels[1]),
         unavailableStateWithMsg("Health-check service responded NOT_SERVING for 'BarService'"));
 
-    // Without health checking, states from underlying Subchannel are delivered directly to origLb
+    // Without health checking, states from underlying Subchannel are delivered directly to the mock
-    hcLbEventDelivery.handleSubchannelState(
+    // listeners.
-        subchannels[0], ConnectivityStateInfo.forNonError(IDLE));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(IDLE));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannels[0]), eq(ConnectivityStateInfo.forNonError(IDLE)));
+        eq(ConnectivityStateInfo.forNonError(IDLE)));
 
     // Re-connecting on a Subchannel will reset the "disabled" flag.
     assertThat(healthImpls[0].calls).hasSize(0);
-    hcLbEventDelivery.handleSubchannelState(
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-        subchannels[0], ConnectivityStateInfo.forNonError(READY));
     assertThat(healthImpls[0].calls).hasSize(1);
     serverCall0 = healthImpls[0].calls.poll();
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannels[0]), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     // Health check now works as normal
     serverCall0.responseObserver.onNext(makeResponse(ServingStatus.SERVICE_UNKNOWN));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannels[0]),
         unavailableStateWithMsg("Health-check service responded SERVICE_UNKNOWN for 'BarService'"));
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockStateListeners[0], mockStateListeners[1]);
     verifyZeroInteractions(backoffPolicyProvider);
   }
 
@@ -458,13 +437,14 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verify(origLb).handleResolvedAddresses(result);
     verifyNoMoreInteractions(origLb);
 
-    FakeSubchannel subchannel = (FakeSubchannel) createSubchannel(0, Attributes.EMPTY);
+    FakeSubchannel subchannel = unwrap(createSubchannel(0, Attributes.EMPTY));
     assertThat(subchannel).isSameInstanceAs(subchannels[0]);
-    InOrder inOrder = inOrder(origLb, backoffPolicyProvider, backoffPolicy1, backoffPolicy2);
+    SubchannelStateListener mockListener = mockStateListeners[0];
+    InOrder inOrder = inOrder(mockListener, backoffPolicyProvider, backoffPolicy1, backoffPolicy2);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
     assertThat(clock.getPendingTasks()).isEmpty();
@@ -474,8 +454,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     healthImpl.calls.poll().responseObserver.onCompleted();
 
     // which results in TRANSIENT_FAILURE
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with " + Status.OK + " for 'TeeService'"));
     assertThat(subchannel.logs).containsExactly(
@@ -487,7 +466,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     inOrder.verify(backoffPolicy1).nextBackoffNanos();
     assertThat(clock.getPendingTasks()).hasSize(1);
 
-    verifyRetryAfterNanos(inOrder, subchannel, healthImpl, 11);
+    verifyRetryAfterNanos(inOrder, mockListener, subchannel, healthImpl, 11);
     assertThat(clock.getPendingTasks()).isEmpty();
 
     subchannel.logs.clear();
@@ -495,8 +474,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     healthImpl.calls.poll().responseObserver.onError(Status.CANCELLED.asException());
 
     // which also results in TRANSIENT_FAILURE, with a different description
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with "
             + Status.CANCELLED + " for 'TeeService'"));
@@ -507,15 +485,15 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // Retry with backoff
     inOrder.verify(backoffPolicy1).nextBackoffNanos();
 
-    verifyRetryAfterNanos(inOrder, subchannel, healthImpl, 21);
+    verifyRetryAfterNanos(inOrder, mockListener, subchannel, healthImpl, 21);
 
     // Server responds this time
     healthImpl.calls.poll().responseObserver.onNext(makeResponse(ServingStatus.SERVING));
 
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
-    verifyNoMoreInteractions(origLb, backoffPolicyProvider, backoffPolicy1);
+    verifyNoMoreInteractions(origLb, mockListener, backoffPolicyProvider, backoffPolicy1);
   }
 
   @Test
@@ -530,13 +508,15 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verify(origLb).handleResolvedAddresses(result);
     verifyNoMoreInteractions(origLb);
 
+    SubchannelStateListener mockStateListener = mockStateListeners[0];
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
-    InOrder inOrder = inOrder(origLb, backoffPolicyProvider, backoffPolicy1, backoffPolicy2);
+    InOrder inOrder =
+        inOrder(mockStateListener, backoffPolicyProvider, backoffPolicy1, backoffPolicy2);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
     assertThat(clock.getPendingTasks()).isEmpty();
@@ -545,8 +525,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     healthImpl.calls.poll().responseObserver.onError(Status.CANCELLED.asException());
 
     // which results in TRANSIENT_FAILURE
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with "
             + Status.CANCELLED + " for 'TeeService'"));
@@ -556,20 +535,19 @@ public class HealthCheckingLoadBalancerFactoryTest {
     inOrder.verify(backoffPolicy1).nextBackoffNanos();
     assertThat(clock.getPendingTasks()).hasSize(1);
 
-    verifyRetryAfterNanos(inOrder, subchannel, healthImpl, 11);
+    verifyRetryAfterNanos(inOrder, mockStateListener, subchannel, healthImpl, 11);
     assertThat(clock.getPendingTasks()).isEmpty();
 
     // Server responds
     healthImpl.calls.peek().responseObserver.onNext(makeResponse(ServingStatus.SERVING));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(mockStateListener);
 
     // then closes the stream
     healthImpl.calls.poll().responseObserver.onError(Status.UNAVAILABLE.asException());
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with "
             + Status.UNAVAILABLE + " for 'TeeService'"));
@@ -577,16 +555,15 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // Because server has responded, the first retry is not subject to backoff.
     // But the backoff policy has been reset.  A new backoff policy will be used for
     // the next backed-off retry.
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     assertThat(healthImpl.calls).hasSize(1);
     assertThat(clock.getPendingTasks()).isEmpty();
     inOrder.verifyNoMoreInteractions();
 
     // then closes the stream for this retry
     healthImpl.calls.poll().responseObserver.onError(Status.UNAVAILABLE.asException());
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with "
             + Status.UNAVAILABLE + " for 'TeeService'"));
@@ -596,19 +573,19 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // Retry with a new backoff policy
     inOrder.verify(backoffPolicy2).nextBackoffNanos();
 
-    verifyRetryAfterNanos(inOrder, subchannel, healthImpl, 12);
+    verifyRetryAfterNanos(inOrder, mockStateListener, subchannel, healthImpl, 12);
   }
 
   private void verifyRetryAfterNanos(
-      InOrder inOrder, Subchannel subchannel, HealthImpl impl, long nanos) {
+      InOrder inOrder, SubchannelStateListener listener, Subchannel subchannel, HealthImpl impl,
+      long nanos) {
     assertThat(impl.calls).isEmpty();
     clock.forwardNanos(nanos - 1);
     assertThat(impl.calls).isEmpty();
     inOrder.verifyNoMoreInteractions();
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(listener);
     clock.forwardNanos(1);
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(listener).onSubchannelState(eq(ConnectivityStateInfo.forNonError(CONNECTING)));
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     assertThat(impl.calls).hasSize(1);
   }
 
@@ -628,13 +605,11 @@ public class HealthCheckingLoadBalancerFactoryTest {
     createSubchannel(0, Attributes.EMPTY);
 
     // No health check activity.  Underlying Subchannel states are directly propagated
-    hcLbEventDelivery.handleSubchannelState(
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-        subchannels[0], ConnectivityStateInfo.forNonError(READY));
     assertThat(healthImpls[0].calls).isEmpty();
-    verify(origLb).handleSubchannelState(
+    verify(mockStateListeners[0]).onSubchannelState(eq(ConnectivityStateInfo.forNonError(READY)));
-        same(subchannels[0]), eq(ConnectivityStateInfo.forNonError(READY)));
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(mockStateListeners[0]);
 
     // Service config enables health check
     Attributes resolutionAttrs = attrsWithHealthCheckService("FooService");
@@ -649,13 +624,12 @@ public class HealthCheckingLoadBalancerFactoryTest {
     assertThat(healthImpls[0].calls).hasSize(1);
 
     // State stays in READY, instead of switching to CONNECTING.
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(mockStateListeners[0]);
 
     // Start Subchannel 1, which will have health check
     createSubchannel(1, Attributes.EMPTY);
     assertThat(healthImpls[1].calls).isEmpty();
-    hcLbEventDelivery.handleSubchannelState(
+    deliverSubchannelState(1, ConnectivityStateInfo.forNonError(READY));
-        subchannels[1], ConnectivityStateInfo.forNonError(READY));
     assertThat(healthImpls[1].calls).hasSize(1);
   }
 
@@ -672,12 +646,12 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
-    InOrder inOrder = inOrder(origLb);
+    InOrder inOrder = inOrder(origLb, mockStateListeners[0]);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     inOrder.verifyNoMoreInteractions();
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
@@ -694,12 +668,12 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // Health check RPC cancelled.
     assertThat(serverCall.cancelled).isTrue();
     // Subchannel uses original state
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
     inOrder.verify(origLb).handleResolvedAddresses(result2);
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockStateListeners[0]);
     assertThat(healthImpl.calls).isEmpty();
   }
 
@@ -716,12 +690,12 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
-    InOrder inOrder = inOrder(origLb);
+    InOrder inOrder = inOrder(origLb, mockStateListeners[0]);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
     inOrder.verifyNoMoreInteractions();
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
@@ -730,8 +704,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     assertThat(clock.getPendingTasks()).isEmpty();
     healthImpl.calls.poll().responseObserver.onCompleted();
     assertThat(clock.getPendingTasks()).hasSize(1);
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with " + Status.OK + " for 'TeeService'"));
 
@@ -749,12 +722,12 @@ public class HealthCheckingLoadBalancerFactoryTest {
     assertThat(healthImpl.calls).isEmpty();
 
     // Subchannel uses original state
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
     inOrder.verify(origLb).handleResolvedAddresses(result2);
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockStateListeners[0]);
   }
 
   @Test
@@ -770,15 +743,15 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
-    InOrder inOrder = inOrder(origLb);
+    InOrder inOrder = inOrder(origLb, mockStateListeners[0]);
 
     // Underlying subchannel is not READY initially
     ConnectivityStateInfo underlyingErrorState =
         ConnectivityStateInfo.forTransientFailure(
             Status.UNAVAILABLE.withDescription("connection refused"));
-    hcLbEventDelivery.handleSubchannelState(subchannel, underlyingErrorState);
+    deliverSubchannelState(0, underlyingErrorState);
-    inOrder.verify(origLb).handleSubchannelState(same(subchannel), same(underlyingErrorState));
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(same(underlyingErrorState));
     inOrder.verifyNoMoreInteractions();
 
     // NameResolver gives an update without service config, thus health check will be disabled
@@ -791,16 +764,16 @@ public class HealthCheckingLoadBalancerFactoryTest {
     inOrder.verify(origLb).handleResolvedAddresses(result2);
 
     // Underlying subchannel is now ready
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
 
     // Since health check is disabled, READY state is propagated directly.
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockStateListeners[0]).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
     // and there is no health check activity.
     assertThat(healthImpls[0].calls).isEmpty();
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockStateListeners[0]);
   }
 
   @Test
@@ -816,12 +789,13 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    SubchannelStateListener mockListener = mockStateListeners[0];
-    InOrder inOrder = inOrder(origLb);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
+    InOrder inOrder = inOrder(origLb, mockListener);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
@@ -831,14 +805,14 @@ public class HealthCheckingLoadBalancerFactoryTest {
 
     // Health check responded
     serverCall.responseObserver.onNext(makeResponse(ServingStatus.SERVING));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(READY)));
+        eq(ConnectivityStateInfo.forNonError(READY)));
 
     // Service config returns with the same health check name.
     hcLbEventDelivery.handleResolvedAddresses(result1);
     // It's delivered to origLb, but nothing else happens
     inOrder.verify(origLb).handleResolvedAddresses(result1);
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
 
     // Service config returns a different health check name.
     resolutionAttrs = attrsWithHealthCheckService("FooService");
@@ -859,7 +833,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     assertThat(serverCall.request).isEqualTo(makeRequest("FooService"));
 
     // State stays in READY, instead of switching to CONNECTING.
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
   }
 
   @Test
@@ -875,12 +849,13 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    SubchannelStateListener mockListener = mockStateListeners[0];
-    InOrder inOrder = inOrder(origLb);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
+    InOrder inOrder = inOrder(origLb, mockListener);
 
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
@@ -893,8 +868,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     serverCall.responseObserver.onCompleted();
     assertThat(clock.getPendingTasks()).hasSize(1);
     assertThat(healthImpl.calls).isEmpty();
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel),
         unavailableStateWithMsg(
             "Health-check stream unexpectedly closed with " + Status.OK + " for 'TeeService'"));
 
@@ -903,7 +877,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     hcLbEventDelivery.handleResolvedAddresses(result1);
     // It's delivered to origLb, but nothing else happens
     inOrder.verify(origLb).handleResolvedAddresses(result1);
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
     assertThat(clock.getPendingTasks()).hasSize(1);
     assertThat(healthImpl.calls).isEmpty();
 
@@ -916,8 +890,8 @@ public class HealthCheckingLoadBalancerFactoryTest {
     hcLbEventDelivery.handleResolvedAddresses(result2);
 
     // Concluded CONNECTING state
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     inOrder.verify(origLb).handleResolvedAddresses(result2);
 
@@ -930,7 +904,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     // with the new service name
     assertThat(serverCall.request).isEqualTo(makeRequest("FooService"));
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
   }
 
   @Test
@@ -946,16 +920,17 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    SubchannelStateListener mockListener = mockStateListeners[0];
-    InOrder inOrder = inOrder(origLb);
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
+    InOrder inOrder = inOrder(origLb, mockListener);
     HealthImpl healthImpl = healthImpls[0];
 
     // Underlying subchannel is not READY initially
     ConnectivityStateInfo underlyingErrorState =
         ConnectivityStateInfo.forTransientFailure(
             Status.UNAVAILABLE.withDescription("connection refused"));
-    hcLbEventDelivery.handleSubchannelState(subchannel, underlyingErrorState);
+    deliverSubchannelState(0, underlyingErrorState);
-    inOrder.verify(origLb).handleSubchannelState(same(subchannel), same(underlyingErrorState));
+    inOrder.verify(mockListener).onSubchannelState(same(underlyingErrorState));
     inOrder.verifyNoMoreInteractions();
 
     // Service config returns with the same health check name.
@@ -976,18 +951,18 @@ public class HealthCheckingLoadBalancerFactoryTest {
     inOrder.verify(origLb).handleResolvedAddresses(result2);
 
     // Underlying subchannel is now ready
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
 
     // Concluded CONNECTING state
-    inOrder.verify(origLb).handleSubchannelState(
+    inOrder.verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     // Health check RPC is started
     assertThat(healthImpl.calls).hasSize(1);
     // with the new service name
     assertThat(healthImpl.calls.poll().request).isEqualTo(makeRequest("FooService"));
 
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
   }
 
   @Test
@@ -1031,18 +1006,19 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verifyNoMoreInteractions(origLb);
 
     Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
-    assertThat(subchannel).isSameInstanceAs(subchannels[0]);
+    SubchannelStateListener mockListener = mockStateListeners[0];
+    assertThat(unwrap(subchannel)).isSameInstanceAs(subchannels[0]);
 
     // Trigger the health check
-    hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
 
     HealthImpl healthImpl = healthImpls[0];
     assertThat(healthImpl.calls).hasSize(1);
     ServerSideCall serverCall = healthImpl.calls.poll();
     assertThat(serverCall.cancelled).isFalse();
 
-    verify(origLb).handleSubchannelState(
+    verify(mockListener).onSubchannelState(
-        same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));
+        eq(ConnectivityStateInfo.forNonError(CONNECTING)));
 
     // Shut down the balancer
     hcLbEventDelivery.shutdown();
@@ -1052,7 +1028,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     assertThat(serverCall.cancelled).isTrue();
 
     // LoadBalancer API requires no more callbacks on LoadBalancer after shutdown() is called.
-    verifyNoMoreInteractions(origLb);
+    verifyNoMoreInteractions(origLb, mockListener);
 
     // No more health check call is made or scheduled
     assertThat(healthImpl.calls).isEmpty();
@@ -1085,8 +1061,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
     verify(origLb).handleResolvedAddresses(result);
     createSubchannel(0, Attributes.EMPTY);
     assertThat(healthImpls[0].calls).isEmpty();
-    hcLbEventDelivery.handleSubchannelState(
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
-        subchannels[0], ConnectivityStateInfo.forNonError(READY));
     assertThat(healthImpls[0].calls).hasSize(1);
   }
 
@@ -1179,6 +1154,8 @@ public class HealthCheckingLoadBalancerFactoryTest {
     final Attributes attrs;
     final Channel channel;
     final ArrayList<String> logs = new ArrayList<>();
+    final int index;
+    SubchannelStateListener listener;
     private final ChannelLogger logger = new ChannelLogger() {
         @Override
         public void log(ChannelLogLevel level, String msg) {
@@ -1191,15 +1168,22 @@ public class HealthCheckingLoadBalancerFactoryTest {
         }
       };
 
-    FakeSubchannel(List<EquivalentAddressGroup> eagList, Attributes attrs, Channel channel) {
+    FakeSubchannel(int index, CreateSubchannelArgs args, Channel channel) {
-      this.eagList = Collections.unmodifiableList(eagList);
+      this.index = index;
-      this.attrs = checkNotNull(attrs);
+      this.eagList = args.getAddresses();
+      this.attrs = args.getAttributes();
       this.channel = checkNotNull(channel);
     }
 
+    @Override
+    public void start(SubchannelStateListener listener) {
+      checkState(this.listener == null);
+      this.listener = listener;
+    }
+
     @Override
     public void shutdown() {
-      hcLbEventDelivery.handleSubchannelState(this, ConnectivityStateInfo.forNonError(SHUTDOWN));
+      deliverSubchannelState(index, ConnectivityStateInfo.forNonError(SHUTDOWN));
     }
 
     @Override
@@ -1230,16 +1214,16 @@ public class HealthCheckingLoadBalancerFactoryTest {
 
   private class FakeHelper extends Helper {
     @Override
-    public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
+    public Subchannel createSubchannel(CreateSubchannelArgs args) {
       int index = -1;
       for (int i = 0; i < NUM_SUBCHANNELS; i++) {
-        if (eagLists[i] == addrs) {
+        if (eagLists[i].equals(args.getAddresses())) {
           index = i;
           break;
         }
       }
-      checkState(index >= 0, "addrs " + addrs + " not found");
+      checkState(index >= 0, "addrs " + args.getAddresses() + " not found");
-      FakeSubchannel subchannel = new FakeSubchannel(addrs, attrs, channels[index]);
+      FakeSubchannel subchannel = new FakeSubchannel(index, args, channels[index]);
       checkState(subchannels[index] == null, "subchannels[" + index + "] already created");
       subchannels[index] = subchannel;
       return subchannel;
@@ -1297,9 +1281,27 @@ public class HealthCheckingLoadBalancerFactoryTest {
     syncContext.execute(new Runnable() {
         @Override
         public void run() {
-          returnedSubchannel.set(wrappedHelper.createSubchannel(eagLists[index], attrs));
+          Subchannel s = wrappedHelper.createSubchannel(CreateSubchannelArgs.newBuilder()
+              .setAddresses(eagLists[index])
+              .setAttributes(attrs)
+              .build());
+          s.start(mockStateListeners[index]);
+          returnedSubchannel.set(s);
         }
       });
     return returnedSubchannel.get();
   }
+
+  private void deliverSubchannelState(final int index, final ConnectivityStateInfo newState) {
+    syncContext.execute(new Runnable() {
+        @Override
+        public void run() {
+          subchannels[index].listener.onSubchannelState(newState);
+        }
+      });
+  }
+
+  private static FakeSubchannel unwrap(Subchannel s) {
+    return (FakeSubchannel) ((SubchannelImpl) s).delegate();
+  }
 }

xds/src/test/java/io/grpc/xds/LocalityStoreTest.java

@@ -65,6 +65,8 @@ import org.mockito.junit.MockitoRule;
  * Tests for {@link LocalityStore}.
  */
 @RunWith(JUnit4.class)
+// TODO(dapengzhang0): remove this after switching to Subchannel.start().
+@SuppressWarnings("deprecation")
 public class LocalityStoreTest {
   @Rule
   public final MockitoRule mockitoRule = MockitoJUnit.rule();