grpclb: support "pick_first" child policy (#5438)

The PICK_FIRST mode puts all backend addresses in a single Subchannel. There are a few points where it's different from the default ROUND_ROBIN mode:

1. PICK_FIRST doesn't eagerly connect to backends like ROUND_ROBIN does. Instead, it requests for connections when the Subchannel is picked.

2. PICK_FIRST adds tokens to the headers via a different code path (`TokenAttachingTracerFactory`) than ROUND_ROBIN

3. For simple implementation, when the mode is changed by service config when the LoadBalancer is working, we will shut down `GrpclbState` and starts a new one with the new mode. All connections will be closed during the transition. We don't expect this to happen in practice given the specific use case of PICK_FIRST.

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

@@ -16,6 +16,8 @@
 
 package io.grpc.grpclb;
 
+import io.grpc.Attributes;
+import io.grpc.EquivalentAddressGroup;
 import io.grpc.ExperimentalApi;
 import io.grpc.Metadata;
 
@@ -32,5 +34,12 @@ public final class GrpclbConstants {
   public static final Metadata.Key<String> TOKEN_METADATA_KEY =
       Metadata.Key.of("lb-token", Metadata.ASCII_STRING_MARSHALLER);
 
+  /**
+   * For passing LB tokens via the EAG attributes.
+   */
+  @EquivalentAddressGroup.Attr
+  static final Attributes.Key<String> TOKEN_ATTRIBUTE_KEY =
+      Attributes.Key.create("lb-token");
+
   private GrpclbConstants() { }
 }

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

@@ -17,6 +17,7 @@
 package io.grpc.grpclb;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.annotations.VisibleForTesting;
 import io.grpc.Attributes;
@@ -51,7 +52,11 @@ class GrpclbLoadBalancer extends LoadBalancer {
   private static final Logger logger = Logger.getLogger(GrpclbLoadBalancer.class.getName());
 
   private final Helper helper;
+  private final TimeProvider time;
   private final SubchannelPool subchannelPool;
+  private final BackoffPolicy.Provider backoffPolicyProvider;
+
+  private Mode mode = Mode.ROUND_ROBIN;
 
   // All mutable states in this class are mutated ONLY from Channel Executor
   @Nullable
@@ -63,12 +68,12 @@ class GrpclbLoadBalancer extends LoadBalancer {
       TimeProvider time,
       BackoffPolicy.Provider backoffPolicyProvider) {
     this.helper = checkNotNull(helper, "helper");
-    checkNotNull(time, "time provider");
+    this.time = checkNotNull(time, "time provider");
-    checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
+    this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
     this.subchannelPool = checkNotNull(subchannelPool, "subchannelPool");
     this.subchannelPool.init(helper);
-    grpclbState =
+    recreateStates();
-        new GrpclbState(helper, subchannelPool, time, backoffPolicyProvider);
+    checkNotNull(grpclbState, "grpclbState");
   }
 
   @Override
@@ -97,7 +102,12 @@ class GrpclbLoadBalancer extends LoadBalancer {
     newBackendServers = Collections.unmodifiableList(newBackendServers);
     Map<String, Object> rawLbConfigValue = attributes.get(ATTR_LOAD_BALANCING_CONFIG);
     Mode newMode = retrieveModeFromLbConfig(rawLbConfigValue, helper.getChannelLogger());
-    grpclbState.handleAddresses(newLbAddressGroups, newBackendServers, newMode);
+    if (!mode.equals(newMode)) {
+      mode = newMode;
+      helper.getChannelLogger().log(ChannelLogLevel.INFO, "Mode: " + newMode);
+      recreateStates();
+    }
+    grpclbState.handleAddresses(newLbAddressGroups, newBackendServers);
   }
 
   @VisibleForTesting
@@ -141,6 +151,12 @@ class GrpclbLoadBalancer extends LoadBalancer {
     }
   }
 
+  private void recreateStates() {
+    resetStates();
+    checkState(grpclbState == null, "Should've been cleared");
+    grpclbState = new GrpclbState(mode, helper, subchannelPool, time, backoffPolicyProvider);
+  }
+
   @Override
   public void shutdown() {
     resetStates();

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

@@ -138,8 +138,8 @@ final class GrpclbState {
 
   @Nullable
   private LbStream lbStream;
-  private Map<EquivalentAddressGroup, Subchannel> subchannels = Collections.emptyMap();
+  private Map<List<EquivalentAddressGroup>, Subchannel> subchannels = Collections.emptyMap();
-  private Mode mode;
+  private final Mode mode;
 
   // Has the same size as the round-robin list from the balancer.
   // A drop entry from the round-robin list becomes a DropEntry here.
@@ -151,10 +151,12 @@ final class GrpclbState {
       new RoundRobinPicker(Collections.<DropEntry>emptyList(), Arrays.asList(BUFFER_ENTRY));
 
   GrpclbState(
+      Mode mode,
       Helper helper,
       SubchannelPool subchannelPool,
       TimeProvider time,
       BackoffPolicy.Provider backoffPolicyProvider) {
+    this.mode = checkNotNull(mode, "mode");
     this.helper = checkNotNull(helper, "helper");
     this.syncContext = checkNotNull(helper.getSynchronizationContext(), "syncContext");
     this.subchannelPool = checkNotNull(subchannelPool, "subchannelPool");
@@ -169,7 +171,7 @@ final class GrpclbState {
     if (newState.getState() == SHUTDOWN || !subchannels.values().contains(subchannel)) {
       return;
     }
-    if (newState.getState() == IDLE) {
+    if (mode == Mode.ROUND_ROBIN && newState.getState() == IDLE) {
       subchannel.requestConnection();
     }
     subchannel.getAttributes().get(STATE_INFO).set(newState);
@@ -182,8 +184,7 @@ final class GrpclbState {
    * not yet connected.
    */
   void handleAddresses(
-      List<LbAddressGroup> newLbAddressGroups, List<EquivalentAddressGroup> newBackendServers,
+      List<LbAddressGroup> newLbAddressGroups, List<EquivalentAddressGroup> newBackendServers) {
-      Mode mode) {
     if (newLbAddressGroups.isEmpty()) {
       propagateError(Status.UNAVAILABLE.withDescription(
               "NameResolver returned no LB address while asking for GRPCLB"));
@@ -305,11 +306,21 @@ final class GrpclbState {
 
   void shutdown() {
     shutdownLbComm();
+    switch (mode) {
+      case ROUND_ROBIN:
         // We close the subchannels through subchannelPool instead of helper just for convenience of
         // testing.
         for (Subchannel subchannel : subchannels.values()) {
           subchannelPool.returnSubchannel(subchannel);
         }
+        break;
+      case PICK_FIRST:
+        checkState(subchannels.size() == 1, "Excessive Subchannels: %s", subchannels);
+        subchannels.values().iterator().next().shutdown();
+        break;
+      default:
+        throw new AssertionError("Missing case for " + mode);
+    }
     subchannels = Collections.emptyMap();
     subchannelPool.clear();
     cancelFallbackTimer();
@@ -341,25 +352,23 @@ final class GrpclbState {
       @Nullable GrpclbClientLoadRecorder loadRecorder) {
     logger.log(
         ChannelLogLevel.INFO, "Using RR list={0}, drop={1}", newBackendAddrList, newDropList);
-    HashMap<EquivalentAddressGroup, Subchannel> newSubchannelMap =
+    HashMap<List<EquivalentAddressGroup>, Subchannel> newSubchannelMap =
         new HashMap<>();
     List<BackendEntry> newBackendList = new ArrayList<>();
 
+    switch (mode) {
+      case ROUND_ROBIN:
         for (BackendAddressGroup backendAddr : newBackendAddrList) {
           EquivalentAddressGroup eag = backendAddr.getAddresses();
-      Subchannel subchannel = newSubchannelMap.get(eag);
+          List<EquivalentAddressGroup> eagAsList = Collections.singletonList(eag);
+          Subchannel subchannel = newSubchannelMap.get(eagAsList);
           if (subchannel == null) {
-        subchannel = subchannels.get(eag);
+            subchannel = subchannels.get(eagAsList);
             if (subchannel == null) {
-          Attributes subchannelAttrs = Attributes.newBuilder()
+              subchannel = subchannelPool.takeOrCreateSubchannel(eag, createSubchannelAttrs());
-              .set(STATE_INFO,
-                  new AtomicReference<>(
-                      ConnectivityStateInfo.forNonError(IDLE)))
-              .build();
-          subchannel = subchannelPool.takeOrCreateSubchannel(eag, subchannelAttrs);
               subchannel.requestConnection();
             }
-        newSubchannelMap.put(eag, subchannel);
+            newSubchannelMap.put(eagAsList, subchannel);
           }
           BackendEntry entry;
           // Only picks with tokens are reported to LoadRecorder
@@ -370,16 +379,47 @@ final class GrpclbState {
           }
           newBackendList.add(entry);
         }
-
         // Close Subchannels whose addresses have been delisted
-    for (Entry<EquivalentAddressGroup, Subchannel> entry : subchannels.entrySet()) {
+        for (Entry<List<EquivalentAddressGroup>, Subchannel> entry : subchannels.entrySet()) {
-      EquivalentAddressGroup eag = entry.getKey();
+          List<EquivalentAddressGroup> eagList = entry.getKey();
-      if (!newSubchannelMap.containsKey(eag)) {
+          if (!newSubchannelMap.containsKey(eagList)) {
             subchannelPool.returnSubchannel(entry.getValue());
           }
         }
-
         subchannels = Collections.unmodifiableMap(newSubchannelMap);
+        break;
+      case PICK_FIRST:
+        List<EquivalentAddressGroup> eagList = new ArrayList<>();
+        // Because for PICK_FIRST, we create a single Subchannel for all addresses, we have to
+        // attach the tokens to the EAG attributes and use TokenAttachingLoadRecorder to put them on
+        // headers.
+        //
+        // The PICK_FIRST code path doesn't cache Subchannels.
+        for (BackendAddressGroup bag : newBackendAddrList) {
+          EquivalentAddressGroup origEag = bag.getAddresses();
+          Attributes eagAttrs = origEag.getAttributes();
+          if (bag.getToken() != null) {
+            eagAttrs = eagAttrs.toBuilder()
+                .set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, bag.getToken()).build();
+          }
+          eagList.add(new EquivalentAddressGroup(origEag.getAddresses(), eagAttrs));
+        }
+        Subchannel subchannel;
+        if (subchannels.isEmpty()) {
+          subchannel = helper.createSubchannel(eagList, createSubchannelAttrs());
+        } else {
+          checkState(subchannels.size() == 1, "Unexpected Subchannel count: %s", subchannels);
+          subchannel = subchannels.values().iterator().next();
+          helper.updateSubchannelAddresses(subchannel, eagList);
+        }
+        subchannels = Collections.singletonMap(eagList, subchannel);
+        newBackendList.add(
+            new BackendEntry(subchannel, new TokenAttachingTracerFactory(loadRecorder)));
+        break;
+      default:
+        throw new AssertionError("Missing case for " + mode);
+    }
+
     dropList = Collections.unmodifiableList(newDropList);
     backendList = Collections.unmodifiableList(newBackendList);
   }
@@ -619,11 +659,15 @@ final class GrpclbState {
    * changed since the last picker created.
    */
   private void maybeUpdatePicker() {
-    List<RoundRobinEntry> pickList = new ArrayList<>(backendList.size());
+    List<RoundRobinEntry> pickList;
+    ConnectivityState state;
+    switch (mode) {
+      case ROUND_ROBIN:
+        pickList = new ArrayList<>(backendList.size());
         Status error = null;
         boolean hasIdle = false;
         for (BackendEntry entry : backendList) {
-      Subchannel subchannel = entry.result.getSubchannel();
+          Subchannel subchannel = entry.subchannel;
           Attributes attrs = subchannel.getAttributes();
           ConnectivityStateInfo stateInfo = attrs.get(STATE_INFO).get();
           if (stateInfo.getState() == READY) {
@@ -634,7 +678,6 @@ final class GrpclbState {
             hasIdle = true;
           }
         }
-    ConnectivityState state;
         if (pickList.isEmpty()) {
           if (error != null && !hasIdle) {
             pickList.add(new ErrorEntry(error));
@@ -646,6 +689,38 @@ final class GrpclbState {
         } else {
           state = READY;
         }
+        break;
+      case PICK_FIRST:
+        if (backendList.isEmpty()) {
+          pickList = Collections.singletonList(BUFFER_ENTRY);
+          // Have not received server addresses
+          state = CONNECTING;
+        } else {
+          checkState(backendList.size() == 1, "Excessive backend entries: %s", backendList);
+          BackendEntry onlyEntry = backendList.get(0);
+          ConnectivityStateInfo stateInfo =
+              onlyEntry.subchannel.getAttributes().get(STATE_INFO).get();
+          state = stateInfo.getState();
+          switch (state) {
+            case READY:
+              pickList = Collections.<RoundRobinEntry>singletonList(onlyEntry);
+              break;
+            case TRANSIENT_FAILURE:
+              pickList =
+                  Collections.<RoundRobinEntry>singletonList(new ErrorEntry(stateInfo.getStatus()));
+              break;
+            case CONNECTING:
+              pickList = Collections.singletonList(BUFFER_ENTRY);
+              break;
+            default:
+              pickList = Collections.<RoundRobinEntry>singletonList(
+                  new IdleSubchannelEntry(onlyEntry.subchannel));
+          }
+        }
+        break;
+      default:
+        throw new AssertionError("Missing case for " + mode);
+    }
     maybeUpdatePicker(state, new RoundRobinPicker(dropList, pickList));
   }
 
@@ -704,6 +779,14 @@ final class GrpclbState {
     return new EquivalentAddressGroup(addrs, attrs);
   }
 
+  private static Attributes createSubchannelAttrs() {
+    return Attributes.newBuilder()
+        .set(STATE_INFO,
+            new AtomicReference<>(
+                ConnectivityStateInfo.forNonError(IDLE)))
+        .build();
+  }
+
   @VisibleForTesting
   static final class DropEntry {
     private final GrpclbClientLoadRecorder loadRecorder;
@@ -740,34 +823,45 @@ final class GrpclbState {
     }
   }
 
-  private interface RoundRobinEntry {
+  @VisibleForTesting
+  interface RoundRobinEntry {
     PickResult picked(Metadata headers);
   }
 
   @VisibleForTesting
   static final class BackendEntry implements RoundRobinEntry {
+    final Subchannel subchannel;
     @VisibleForTesting
     final PickResult result;
     @Nullable
-    private final GrpclbClientLoadRecorder loadRecorder;
-    @Nullable
     private final String token;
 
     /**
-     * Creates a BackendEntry whose usage will be reported to load recorder.
+     * For ROUND_ROBIN: creates a BackendEntry whose usage will be reported to load recorder.
      */
     BackendEntry(Subchannel subchannel, GrpclbClientLoadRecorder loadRecorder, String token) {
-      this.result = PickResult.withSubchannel(subchannel, loadRecorder);
+      this.subchannel = checkNotNull(subchannel, "subchannel");
-      this.loadRecorder = checkNotNull(loadRecorder, "loadRecorder");
+      this.result =
+          PickResult.withSubchannel(subchannel, checkNotNull(loadRecorder, "loadRecorder"));
       this.token = checkNotNull(token, "token");
     }
 
     /**
-     * Creates a BackendEntry whose usage will not be reported.
+     * For ROUND_ROBIN/PICK_FIRST: creates a BackendEntry whose usage will not be reported.
      */
     BackendEntry(Subchannel subchannel) {
+      this.subchannel = checkNotNull(subchannel, "subchannel");
       this.result = PickResult.withSubchannel(subchannel);
-      this.loadRecorder = null;
+      this.token = null;
+    }
+
+    /**
+     * For PICK_FIRST: creates a BackendEntry that includes all addresses.
+     */
+    BackendEntry(Subchannel subchannel, TokenAttachingTracerFactory tracerFactory) {
+      this.subchannel = checkNotNull(subchannel, "subchannel");
+      this.result =
+          PickResult.withSubchannel(subchannel, checkNotNull(tracerFactory, "tracerFactory"));
       this.token = null;
     }
 
@@ -783,12 +877,12 @@ final class GrpclbState {
     @Override
     public String toString() {
       // This is printed in logs.  Only give out useful information.
-      return "[" + result.getSubchannel().getAllAddresses().toString() + "(" + token + ")]";
+      return "[" + subchannel.getAllAddresses().toString() + "(" + token + ")]";
     }
 
     @Override
     public int hashCode() {
-      return Objects.hashCode(loadRecorder, result, token);
+      return Objects.hashCode(result, token);
     }
 
     @Override
@@ -797,8 +891,42 @@ final class GrpclbState {
         return false;
       }
       BackendEntry that = (BackendEntry) other;
-      return Objects.equal(result, that.result) && Objects.equal(token, that.token)
+      return Objects.equal(result, that.result) && Objects.equal(token, that.token);
-          && Objects.equal(loadRecorder, that.loadRecorder);
+    }
+  }
+
+  @VisibleForTesting
+  static final class IdleSubchannelEntry implements RoundRobinEntry {
+    private final Subchannel subchannel;
+
+    IdleSubchannelEntry(Subchannel subchannel) {
+      this.subchannel = checkNotNull(subchannel, "subchannel");
+    }
+
+    @Override
+    public PickResult picked(Metadata headers) {
+      subchannel.requestConnection();
+      return PickResult.withNoResult();
+    }
+
+    @Override
+    public String toString() {
+      // This is printed in logs.  Only give out useful information.
+      return "(idle)[" + subchannel.getAllAddresses().toString() + "]";
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(subchannel);
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      if (!(other instanceof IdleSubchannelEntry)) {
+        return false;
+      }
+      IdleSubchannelEntry that = (IdleSubchannelEntry) other;
+      return Objects.equal(subchannel, that.subchannel);
     }
   }
 
@@ -860,8 +988,7 @@ final class GrpclbState {
         // First round-robin on dropList. If a drop entry is selected, request will be dropped.  If
         // a non-drop entry is selected, then round-robin on pickList.  This makes sure requests are
         // dropped at the same proportion as the drop entries appear on the round-robin list from
-        // the balancer, while only READY backends (that make up pickList) are selected for the
+        // the balancer, while only backends from pickList are selected for the non-drop cases.
-        // non-drop cases.
         if (!dropList.isEmpty()) {
           DropEntry drop = dropList.get(dropIndex);
           dropIndex++;
@@ -881,5 +1008,14 @@ final class GrpclbState {
         return pick.picked(args.getHeaders());
       }
     }
+
+    @Override
+    public void requestConnection() {
+      for (RoundRobinEntry entry : pickList) {
+        if (entry instanceof IdleSubchannelEntry) {
+          ((IdleSubchannelEntry) entry).subchannel.requestConnection();
+        }
+      }
+    }
   }
 }

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

@@ -0,0 +1,72 @@
+/*
+ * 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.grpclb;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.base.Objects;
+import io.grpc.Attributes;
+import io.grpc.ClientStreamTracer;
+import io.grpc.Metadata;
+import io.grpc.internal.GrpcAttributes;
+import javax.annotation.Nullable;
+
+/**
+ * Wraps a {@link ClientStreamTracer.Factory}, retrieves tokens from transport attributes and
+ * attaches them to headers.  This is only used in the PICK_FIRST mode.
+ */
+final class TokenAttachingTracerFactory extends ClientStreamTracer.Factory {
+  private static final ClientStreamTracer NOOP_TRACER = new ClientStreamTracer() {};
+
+  @Nullable
+  private final ClientStreamTracer.Factory delegate;
+
+  TokenAttachingTracerFactory(@Nullable ClientStreamTracer.Factory delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public ClientStreamTracer newClientStreamTracer(
+      ClientStreamTracer.StreamInfo info, Metadata headers) {
+    Attributes transportAttrs = checkNotNull(info.getTransportAttrs(), "transportAttrs");
+    Attributes eagAttrs =
+        checkNotNull(transportAttrs.get(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS), "eagAttrs");
+    String token = eagAttrs.get(GrpclbConstants.TOKEN_ATTRIBUTE_KEY);
+    headers.discardAll(GrpclbConstants.TOKEN_METADATA_KEY);
+    if (token != null) {
+      headers.put(GrpclbConstants.TOKEN_METADATA_KEY, token);
+    }
+    if (delegate != null) {
+      return delegate.newClientStreamTracer(info, headers);
+    } else {
+      return NOOP_TRACER;
+    }
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(delegate);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (!(other instanceof TokenAttachingTracerFactory)) {
+      return false;
+    }
+    return Objects.equal(delegate, ((TokenAttachingTracerFactory) other).delegate);
+  }
+}

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

@@ -56,6 +56,7 @@ import io.grpc.ClientStreamTracer;
 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.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
@@ -68,7 +69,9 @@ import io.grpc.SynchronizationContext;
 import io.grpc.grpclb.GrpclbState.BackendEntry;
 import io.grpc.grpclb.GrpclbState.DropEntry;
 import io.grpc.grpclb.GrpclbState.ErrorEntry;
+import io.grpc.grpclb.GrpclbState.IdleSubchannelEntry;
 import io.grpc.grpclb.GrpclbState.Mode;
+import io.grpc.grpclb.GrpclbState.RoundRobinEntry;
 import io.grpc.grpclb.GrpclbState.RoundRobinPicker;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
@@ -166,7 +169,8 @@ public class GrpclbLoadBalancerTest {
       new LinkedList<>();
   private final LinkedList<Subchannel> mockSubchannels = new LinkedList<>();
   private final LinkedList<ManagedChannel> fakeOobChannels = new LinkedList<>();
-  private final ArrayList<Subchannel> subchannelTracker = new ArrayList<>();
+  private final ArrayList<Subchannel> pooledSubchannelTracker = new ArrayList<>();
+  private final ArrayList<Subchannel> unpooledSubchannelTracker = new ArrayList<>();
   private final ArrayList<ManagedChannel> oobChannelTracker = new ArrayList<>();
   private final ArrayList<String> failingLbAuthorities = new ArrayList<>();
   private final SynchronizationContext syncContext = new SynchronizationContext(
@@ -251,11 +255,25 @@ public class GrpclbLoadBalancerTest {
           when(subchannel.getAllAddresses()).thenReturn(Arrays.asList(eag));
           when(subchannel.getAttributes()).thenReturn(attrs);
           mockSubchannels.add(subchannel);
-          subchannelTracker.add(subchannel);
+          pooledSubchannelTracker.add(subchannel);
           return subchannel;
         }
       }).when(subchannelPool).takeOrCreateSubchannel(
           any(EquivalentAddressGroup.class), any(Attributes.class));
+    doAnswer(new Answer<Subchannel>() {
+        @Override
+        public Subchannel answer(InvocationOnMock invocation) throws Throwable {
+          Subchannel subchannel = mock(Subchannel.class);
+          List<EquivalentAddressGroup> eagList =
+              (List<EquivalentAddressGroup>) invocation.getArguments()[0];
+          Attributes attrs = (Attributes) invocation.getArguments()[1];
+          when(subchannel.getAllAddresses()).thenReturn(eagList);
+          when(subchannel.getAttributes()).thenReturn(attrs);
+          mockSubchannels.add(subchannel);
+          unpooledSubchannelTracker.add(subchannel);
+          return subchannel;
+        }
+      }).when(helper).createSubchannel(any(List.class), any(Attributes.class));
     when(helper.getSynchronizationContext()).thenReturn(syncContext);
     when(helper.getScheduledExecutorService()).thenReturn(fakeClock.getScheduledExecutorService());
     when(helper.getChannelLogger()).thenReturn(channelLogger);
@@ -294,14 +312,15 @@ public class GrpclbLoadBalancerTest {
         assertTrue(channel + " is terminated", channel.isTerminated());
       }
       // GRPCLB manages subchannels only through subchannelPool
-      for (Subchannel subchannel: subchannelTracker) {
+      for (Subchannel subchannel : pooledSubchannelTracker) {
         verify(subchannelPool).returnSubchannel(same(subchannel));
         // Our mock subchannelPool never calls Subchannel.shutdown(), thus we can tell if
         // LoadBalancer has called it expectedly.
         verify(subchannel, never()).shutdown();
       }
-      verify(helper, never())
+      for (Subchannel subchannel : unpooledSubchannelTracker) {
-          .createSubchannel(any(List.class), any(Attributes.class));
+        verify(subchannel).shutdown();
+      }
       // No timer should linger after shutdown
       assertThat(fakeClock.getPendingTasks()).isEmpty();
     } finally {
@@ -406,6 +425,65 @@ public class GrpclbLoadBalancerTest {
     verify(subchannel, never()).getAttributes();
   }
 
+  @Test
+  public void roundRobinPickerWithIdleEntry_noDrop() {
+    Subchannel subchannel = mock(Subchannel.class);
+    IdleSubchannelEntry entry = new IdleSubchannelEntry(subchannel);
+
+    RoundRobinPicker picker =
+        new RoundRobinPicker(Collections.<DropEntry>emptyList(), Collections.singletonList(entry));
+    PickSubchannelArgs args = mock(PickSubchannelArgs.class);
+
+    verify(subchannel, never()).requestConnection();
+    assertThat(picker.pickSubchannel(args)).isSameAs(PickResult.withNoResult());
+    verify(subchannel).requestConnection();
+  }
+
+  @Test
+  public void roundRobinPickerWithIdleEntry_andDrop() {
+    GrpclbClientLoadRecorder loadRecorder =
+        new GrpclbClientLoadRecorder(fakeClock.getTimeProvider());
+    // 1 out of 2 requests are to be dropped
+    DropEntry d = new DropEntry(loadRecorder, "LBTOKEN0003");
+    List<DropEntry> dropList = Arrays.asList(null, d);
+
+    Subchannel subchannel = mock(Subchannel.class);
+    IdleSubchannelEntry entry = new IdleSubchannelEntry(subchannel);
+
+    RoundRobinPicker picker = new RoundRobinPicker(dropList, Collections.singletonList(entry));
+    PickSubchannelArgs args = mock(PickSubchannelArgs.class);
+
+    verify(subchannel, never()).requestConnection();
+    assertThat(picker.pickSubchannel(args)).isSameAs(PickResult.withNoResult());
+    verify(subchannel).requestConnection();
+
+    assertThat(picker.pickSubchannel(args)).isSameAs(DROP_PICK_RESULT);
+
+    verify(subchannel).requestConnection();
+    assertThat(picker.pickSubchannel(args)).isSameAs(PickResult.withNoResult());
+    verify(subchannel, times(2)).requestConnection();
+  }
+
+  @Test
+  public void roundRobinPicker_requestConnection() {
+    // requestConnection() on RoundRobinPicker is only passed to IdleSubchannelEntry
+
+    Subchannel subchannel1 = mock(Subchannel.class);
+    Subchannel subchannel2 = mock(Subchannel.class);
+
+    RoundRobinPicker picker = new RoundRobinPicker(
+        Collections.<DropEntry>emptyList(),
+        Arrays.<RoundRobinEntry>asList(
+            new BackendEntry(subchannel1), new IdleSubchannelEntry(subchannel2),
+            new ErrorEntry(Status.UNAVAILABLE)));
+
+    verify(subchannel2, never()).requestConnection();
+
+    picker.requestConnection();
+    verify(subchannel2).requestConnection();
+    verify(subchannel1, never()).requestConnection();
+  }
+
   @Test
   public void loadReporting() {
     Metadata headers = new Metadata();
@@ -1591,6 +1669,297 @@ public class GrpclbLoadBalancerTest {
     verify(helper, times(4)).refreshNameResolution();
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void grpclbWorking_pickFirstMode() throws Exception {
+    InOrder inOrder = inOrder(helper);
+
+    String lbConfig = "{\"childPolicy\" : [ {\"pick_first\" : {}} ]}";
+    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+
+    assertEquals(1, fakeOobChannels.size());
+    ManagedChannel oobChannel = fakeOobChannels.poll();
+    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
+    StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
+    assertEquals(1, lbRequestObservers.size());
+    StreamObserver<LoadBalanceRequest> lbRequestObserver = lbRequestObservers.poll();
+    verify(lbRequestObserver).onNext(
+        eq(LoadBalanceRequest.newBuilder().setInitialRequest(
+                InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build())
+            .build()));
+
+    // Simulate receiving LB response
+    List<ServerEntry> backends1 = Arrays.asList(
+        new ServerEntry("127.0.0.1", 2000, "token0001"),
+        new ServerEntry("127.0.0.1", 2010, "token0002"));
+    inOrder.verify(helper, never())
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
+    lbResponseObserver.onNext(buildInitialResponse());
+    lbResponseObserver.onNext(buildLbResponse(backends1));
+
+    inOrder.verify(helper).createSubchannel(
+        eq(Arrays.asList(
+                new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+                new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")))),
+        any(Attributes.class));
+
+    // Initially IDLE
+    inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue();
+
+    // Only one subchannel is created
+    assertThat(mockSubchannels).hasSize(1);
+    Subchannel subchannel = mockSubchannels.poll();
+    assertThat(picker0.dropList).containsExactly(null, null);
+    assertThat(picker0.pickList).containsExactly(new IdleSubchannelEntry(subchannel));
+
+    // PICK_FIRST doesn't eagerly connect
+    verify(subchannel, never()).requestConnection();
+
+    // CONNECTING
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(CONNECTING));
+
+    inOrder.verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
+    RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker1.dropList).containsExactly(null, null);
+    assertThat(picker1.pickList).containsExactly(BUFFER_ENTRY);
+
+    // TRANSIENT_FAILURE
+    Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));
+    inOrder.verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
+    RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker2.dropList).containsExactly(null, null);
+    assertThat(picker2.pickList).containsExactly(new ErrorEntry(error));
+
+    // READY
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker3.dropList).containsExactly(null, null);
+    assertThat(picker3.pickList).containsExactly(
+        new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder())));
+
+
+    // New server list with drops
+    List<ServerEntry> backends2 = Arrays.asList(
+        new ServerEntry("127.0.0.1", 2000, "token0001"),
+        new ServerEntry("token0003"),  // drop
+        new ServerEntry("127.0.0.1", 2020, "token0004"));
+    inOrder.verify(helper, never())
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
+    lbResponseObserver.onNext(buildLbResponse(backends2));
+
+    // new addresses will be updated to the existing subchannel
+    // createSubchannel() has ever been called only once
+    verify(helper, times(1)).createSubchannel(any(List.class), any(Attributes.class));
+    assertThat(mockSubchannels).isEmpty();
+    inOrder.verify(helper).updateSubchannelAddresses(
+        same(subchannel),
+        eq(Arrays.asList(
+                new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
+                new EquivalentAddressGroup(backends2.get(2).addr,
+                    eagAttrsWithToken("token0004")))));
+    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    RoundRobinPicker picker4 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker4.dropList).containsExactly(
+        null, new DropEntry(getLoadRecorder(), "token0003"), null);
+    assertThat(picker4.pickList).containsExactly(
+        new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder())));
+
+    // Subchannel goes IDLE, but PICK_FIRST will not try to reconnect
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
+    inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    RoundRobinPicker picker5 = (RoundRobinPicker) pickerCaptor.getValue();
+    verify(subchannel, never()).requestConnection();
+
+    // ... until it's selected
+    PickSubchannelArgs args = mock(PickSubchannelArgs.class);
+    PickResult pick = picker5.pickSubchannel(args);
+    assertThat(pick).isSameAs(PickResult.withNoResult());
+    verify(subchannel).requestConnection();
+
+    // ... or requested by application
+    picker5.requestConnection();
+    verify(subchannel, times(2)).requestConnection();
+
+    // PICK_FIRST doesn't use subchannelPool
+    verify(subchannelPool, never())
+        .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class));
+    verify(subchannelPool, never()).returnSubchannel(any(Subchannel.class));
+  }
+
+  @SuppressWarnings("unchecked")
+  @Test
+  public void pickFirstMode_fallback() throws Exception {
+    InOrder inOrder = inOrder(helper);
+
+    String lbConfig = "{\"childPolicy\" : [ {\"pick_first\" : {}} ]}";
+
+    // Name resolver returns a mix of balancer and backend addresses
+    List<EquivalentAddressGroup> grpclbResolutionList =
+        createResolvedServerAddresses(false, true, false);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+
+    // Attempted to connect to balancer
+    assertEquals(1, fakeOobChannels.size());
+    ManagedChannel oobChannel = fakeOobChannels.poll();
+    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
+    StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
+    assertEquals(1, lbRequestObservers.size());
+
+    // Fallback timer expires with no response
+    fakeClock.forwardTime(GrpclbState.FALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+
+    // Entering fallback mode
+    inOrder.verify(helper).createSubchannel(
+        eq(Arrays.asList(grpclbResolutionList.get(0), grpclbResolutionList.get(2))),
+        any(Attributes.class));
+
+    assertThat(mockSubchannels).hasSize(1);
+    Subchannel subchannel = mockSubchannels.poll();
+
+    // Initially IDLE
+    inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue();
+
+    // READY
+    deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker1.dropList).containsExactly(null, null);
+    assertThat(picker1.pickList).containsExactly(
+        new BackendEntry(subchannel, new TokenAttachingTracerFactory(null)));
+
+    assertThat(picker0.dropList).containsExactly(null, null);
+    assertThat(picker0.pickList).containsExactly(new IdleSubchannelEntry(subchannel));
+
+
+    // Finally, an LB response, which brings us out of fallback
+    List<ServerEntry> backends1 = Arrays.asList(
+        new ServerEntry("127.0.0.1", 2000, "token0001"),
+        new ServerEntry("127.0.0.1", 2010, "token0002"));
+    inOrder.verify(helper, never())
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
+    lbResponseObserver.onNext(buildInitialResponse());
+    lbResponseObserver.onNext(buildLbResponse(backends1));
+
+    // new addresses will be updated to the existing subchannel
+    // createSubchannel() has ever been called only once
+    verify(helper, times(1)).createSubchannel(any(List.class), any(Attributes.class));
+    assertThat(mockSubchannels).isEmpty();
+    inOrder.verify(helper).updateSubchannelAddresses(
+        same(subchannel),
+        eq(Arrays.asList(
+                new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+                new EquivalentAddressGroup(backends1.get(1).addr,
+                    eagAttrsWithToken("token0002")))));
+    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();
+    assertThat(picker2.dropList).containsExactly(null, null);
+    assertThat(picker2.pickList).containsExactly(
+        new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder())));
+
+    // PICK_FIRST doesn't use subchannelPool
+    verify(subchannelPool, never())
+        .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class));
+    verify(subchannelPool, never()).returnSubchannel(any(Subchannel.class));
+  }
+
+  @Test
+  public void switchMode() throws Exception {
+    InOrder inOrder = inOrder(helper);
+
+    String lbConfig = "{\"childPolicy\" : [ {\"round_robin\" : {}} ]}";
+    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+
+    assertEquals(1, fakeOobChannels.size());
+    ManagedChannel oobChannel = fakeOobChannels.poll();
+    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
+    StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
+    assertEquals(1, lbRequestObservers.size());
+    StreamObserver<LoadBalanceRequest> lbRequestObserver = lbRequestObservers.poll();
+    verify(lbRequestObserver).onNext(
+        eq(LoadBalanceRequest.newBuilder().setInitialRequest(
+                InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build())
+            .build()));
+
+    // Simulate receiving LB response
+    List<ServerEntry> backends1 = Arrays.asList(
+        new ServerEntry("127.0.0.1", 2000, "token0001"),
+        new ServerEntry("127.0.0.1", 2010, "token0002"));
+    inOrder.verify(helper, never())
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
+    lbResponseObserver.onNext(buildInitialResponse());
+    lbResponseObserver.onNext(buildLbResponse(backends1));
+
+    // ROUND_ROBIN: create one subchannel per server
+    verify(subchannelPool).takeOrCreateSubchannel(
+        eq(new EquivalentAddressGroup(backends1.get(0).addr, LB_BACKEND_ATTRS)),
+        any(Attributes.class));
+    verify(subchannelPool).takeOrCreateSubchannel(
+        eq(new EquivalentAddressGroup(backends1.get(1).addr, LB_BACKEND_ATTRS)),
+        any(Attributes.class));
+    inOrder.verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class));
+    assertEquals(2, mockSubchannels.size());
+    Subchannel subchannel1 = mockSubchannels.poll();
+    Subchannel subchannel2 = mockSubchannels.poll();
+    verify(subchannelPool, never()).returnSubchannel(any(Subchannel.class));
+
+    // Switch to PICK_FIRST
+    lbConfig = "{\"childPolicy\" : [ {\"pick_first\" : {}} ]}";
+    grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+
+
+    // GrpclbState will be shutdown, and a new one will be created
+    assertThat(oobChannel.isShutdown()).isTrue();
+    verify(subchannelPool).returnSubchannel(same(subchannel1));
+    verify(subchannelPool).returnSubchannel(same(subchannel2));
+
+    // A new LB stream is created
+    assertEquals(1, fakeOobChannels.size());
+    oobChannel = fakeOobChannels.poll();
+    verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture());
+    lbResponseObserver = lbResponseObserverCaptor.getValue();
+    assertEquals(1, lbRequestObservers.size());
+    lbRequestObserver = lbRequestObservers.poll();
+    verify(lbRequestObserver).onNext(
+        eq(LoadBalanceRequest.newBuilder().setInitialRequest(
+                InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build())
+            .build()));
+
+    // Simulate receiving LB response
+    inOrder.verify(helper, never())
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
+    lbResponseObserver.onNext(buildInitialResponse());
+    lbResponseObserver.onNext(buildLbResponse(backends1));
+
+    // PICK_FIRST Subchannel
+    inOrder.verify(helper).createSubchannel(
+        eq(Arrays.asList(
+                new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+                new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")))),
+        any(Attributes.class));
+    
+    inOrder.verify(helper).updateBalancingState(eq(IDLE), any(SubchannelPicker.class));
+  }
+
+  private static Attributes eagAttrsWithToken(String token) {
+    return LB_BACKEND_ATTRS.toBuilder().set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, token).build();
+  }
+
   @Test
   public void retrieveModeFromLbConfig_pickFirst() throws Exception {
     String lbConfig = "{\"childPolicy\" : [{\"pick_first\" : {}}, {\"round_robin\" : {}}]}";

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

@@ -0,0 +1,124 @@
+/*
+ * 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.grpclb;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.AdditionalAnswers.delegatesTo;
+import static org.mockito.Matchers.same;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import io.grpc.Attributes;
+import io.grpc.CallOptions;
+import io.grpc.ClientStreamTracer;
+import io.grpc.Metadata;
+import io.grpc.internal.GrpcAttributes;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link TokenAttachingTracerFactory}. */
+@RunWith(JUnit4.class)
+public class TokenAttachingTracerFactoryTest {
+  private static final ClientStreamTracer fakeTracer = new ClientStreamTracer() {};
+
+  private final ClientStreamTracer.Factory delegate = mock(
+      ClientStreamTracer.Factory.class,
+      delegatesTo(
+          new ClientStreamTracer.Factory() {
+            @Override
+            public ClientStreamTracer newClientStreamTracer(
+                ClientStreamTracer.StreamInfo info, Metadata headers) {
+              return fakeTracer;
+            }
+          }));
+
+  @Test
+  public void hasToken() {
+    TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(delegate);
+    ClientStreamTracer.StreamInfo info = new ClientStreamTracer.StreamInfo() {
+        @Override
+        public Attributes getTransportAttrs() {
+          Attributes eagAttrs = Attributes.newBuilder()
+              .set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, "token0001").build();
+          return Attributes.newBuilder()
+              .set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, eagAttrs).build();
+        }
+
+        @Override
+        public CallOptions getCallOptions() {
+          return CallOptions.DEFAULT;
+        }
+      };
+    Metadata headers = new Metadata();
+    // Preexisting token should be replaced
+    headers.put(GrpclbConstants.TOKEN_METADATA_KEY, "preexisting-token");
+
+    ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
+    verify(delegate).newClientStreamTracer(same(info), same(headers));
+    assertThat(tracer).isSameAs(fakeTracer);
+    assertThat(headers.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("token0001");
+  }
+
+  @Test
+  public void noToken() {
+    TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(delegate);
+    ClientStreamTracer.StreamInfo info = new ClientStreamTracer.StreamInfo() {
+        @Override
+        public Attributes getTransportAttrs() {
+          return Attributes.newBuilder()
+              .set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, Attributes.EMPTY).build();
+        }
+
+        @Override
+        public CallOptions getCallOptions() {
+          return CallOptions.DEFAULT;
+        }
+      };
+
+    Metadata headers = new Metadata();
+    // Preexisting token should be removed
+    headers.put(GrpclbConstants.TOKEN_METADATA_KEY, "preexisting-token");
+
+    ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
+    verify(delegate).newClientStreamTracer(same(info), same(headers));
+    assertThat(tracer).isSameAs(fakeTracer);
+    assertThat(headers.get(GrpclbConstants.TOKEN_METADATA_KEY)).isNull();
+  }
+
+  @Test
+  public void nullDelegate() {
+    TokenAttachingTracerFactory factory = new TokenAttachingTracerFactory(null);
+    ClientStreamTracer.StreamInfo info = new ClientStreamTracer.StreamInfo() {
+        @Override
+        public Attributes getTransportAttrs() {
+          return Attributes.newBuilder()
+              .set(GrpcAttributes.ATTR_CLIENT_EAG_ATTRS, Attributes.EMPTY).build();
+        }
+
+        @Override
+        public CallOptions getCallOptions() {
+          return CallOptions.DEFAULT;
+        }
+      };
+    Metadata headers = new Metadata();
+
+    ClientStreamTracer tracer = factory.newClientStreamTracer(info, headers);
+    assertThat(tracer).isNotNull();
+    assertThat(headers.get(GrpclbConstants.TOKEN_METADATA_KEY)).isNull();
+  }
+}