Java 类io.grpc.LoadBalancer.Subchannel 实例源码

private Subchannel nextSubchannel(GrpcURL refUrl) {
  if (size == 0) {
    throw new NoSuchElementException();
  }
  synchronized (this) {
    Subchannel val = list.get(index);
    boolean discard = discard(refUrl, val);
    index++;
    if (index >= size) {
      index = 0;
    }
    if (discard) {
      return nextSubchannel(refUrl);
    } else {
      return val;
    }
  }
}

private ConnectivityState getAggregatedState() {
  Set<ConnectivityState> states = EnumSet.noneOf(ConnectivityState.class);
  for (Subchannel subchannel : getSubchannels()) {
    states.add(getSubchannelStateInfoRef(subchannel).value.getState());
  }
  if (states.contains(READY)) {
    return READY;
  }
  if (states.contains(CONNECTING)) {
    return CONNECTING;
  }
  if (states.contains(IDLE)) {
    // This subchannel IDLE is not because of channel IDLE_TIMEOUT, in which case LB is already
    // shutdown.
    // RRLB will request connection immediately on subchannel IDLE.
    return CONNECTING;
  }
  return TRANSIENT_FAILURE;
}

@Test
public void pickAfterStateChangeAfterResolution() throws Exception {
  loadBalancer.handleResolvedAddressGroups(servers, affinity);
  verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
  Subchannel subchannel = pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel();
  reset(mockHelper);

  InOrder inOrder = inOrder(mockHelper);

  Status error = Status.UNAVAILABLE.withDescription("boom!");
  loadBalancer.handleSubchannelState(subchannel,
      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));
  inOrder.verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
  assertEquals(Status.OK, pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus());

  loadBalancer.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
  inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
  assertEquals(subchannel, pickerCaptor.getValue().pickSubchannel(mockArgs).getSubchannel());

  verifyNoMoreInteractions(mockHelper);
}

@Test
public void updateSubchannelAddresses_newAddressConnects() {
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata()); // Create LB
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  Subchannel subchannel = helper.createSubchannel(servers.get(0), Attributes.EMPTY);

  subchannel.requestConnection();
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  helper.updateSubchannelAddresses(subchannel, servers.get(1));

  subchannel.requestConnection();
  MockClientTransportInfo t1 = newTransports.poll();
  t1.listener.transportReady();
}

@Test
public void updateSubchannelAddresses_existingAddressDoesNotConnect() {
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata()); // Create LB
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  Subchannel subchannel = helper.createSubchannel(servers.get(0), Attributes.EMPTY);

  subchannel.requestConnection();
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  List<SocketAddress> changedList = new ArrayList<SocketAddress>(servers.get(0).getAddresses());
  changedList.add(new FakeSocketAddress("aDifferentServer"));
  helper.updateSubchannelAddresses(subchannel, new EquivalentAddressGroup(changedList));

  subchannel.requestConnection();
  assertNull(newTransports.poll());
}

@Test
public void pickAfterResolved() throws Exception {
  final Subchannel readySubchannel = subchannels.values().iterator().next();
  loadBalancer.handleResolvedAddressGroups(servers, affinity);
  loadBalancer.handleSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));

  verify(mockHelper, times(3)).createSubchannel(eagCaptor.capture(),
      any(Attributes.class));

  assertThat(eagCaptor.getAllValues()).containsAllIn(subchannels.keySet());
  for (Subchannel subchannel : subchannels.values()) {
    verify(subchannel).requestConnection();
    verify(subchannel, never()).shutdown();
  }

  verify(mockHelper, times(2))
      .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());

  assertEquals(CONNECTING, stateCaptor.getAllValues().get(0));
  assertEquals(READY, stateCaptor.getAllValues().get(1));
  assertThat(pickerCaptor.getValue().getList()).containsExactly(readySubchannel);

  verifyNoMoreInteractions(mockHelper);
}

@Test
public void pickerRoundRobin() throws Exception {
  Subchannel subchannel = mock(Subchannel.class);
  Subchannel subchannel1 = mock(Subchannel.class);
  Subchannel subchannel2 = mock(Subchannel.class);

  Picker picker = new Picker(Collections.unmodifiableList(
      Lists.<Subchannel>newArrayList(subchannel, subchannel1, subchannel2)), null);

  assertThat(picker.getList()).containsExactly(subchannel, subchannel1, subchannel2);

  assertEquals(subchannel, picker.pickSubchannel(mockArgs).getSubchannel());
  assertEquals(subchannel1, picker.pickSubchannel(mockArgs).getSubchannel());
  assertEquals(subchannel2, picker.pickSubchannel(mockArgs).getSubchannel());
  assertEquals(subchannel, picker.pickSubchannel(mockArgs).getSubchannel());
}

@Test
public void nameResolutionErrorWithActiveChannels() throws Exception {
  final Subchannel readySubchannel = subchannels.values().iterator().next();
  loadBalancer.handleResolvedAddressGroups(servers, affinity);
  loadBalancer.handleSubchannelState(readySubchannel, ConnectivityStateInfo.forNonError(READY));
  loadBalancer.handleNameResolutionError(Status.NOT_FOUND.withDescription("nameResolutionError"));

  verify(mockHelper, times(3)).createSubchannel(any(EquivalentAddressGroup.class),
      any(Attributes.class));
  verify(mockHelper, times(3))
      .updateBalancingState(stateCaptor.capture(), pickerCaptor.capture());

  Iterator<ConnectivityState> stateIterator = stateCaptor.getAllValues().iterator();
  assertEquals(CONNECTING, stateIterator.next());
  assertEquals(READY, stateIterator.next());
  assertEquals(TRANSIENT_FAILURE, stateIterator.next());

  LoadBalancer.PickResult pickResult = pickerCaptor.getValue().pickSubchannel(mockArgs);
  assertEquals(readySubchannel, pickResult.getSubchannel());
  assertEquals(Status.OK.getCode(), pickResult.getStatus().getCode());

  LoadBalancer.PickResult pickResult2 = pickerCaptor.getValue().pickSubchannel(mockArgs);
  assertEquals(readySubchannel, pickResult2.getSubchannel());
  verifyNoMoreInteractions(mockHelper);
}

@Override
public PickResult pickSubchannel(PickSubchannelArgs args) {
  Map<String, Object> affinity =
      args.getCallOptions().getOption(GrpcCallOptions.CALLOPTIONS_CUSTOME_KEY);
  GrpcURL refUrl = (GrpcURL) affinity.get(GrpcCallOptions.GRPC_REF_URL);
  if (size > 0) {
    Subchannel subchannel = nextSubchannel(refUrl);
    affinity.put(GrpcCallOptions.GRPC_NAMERESOVER_ATTRIBUTES, nameResovleCache);
    return PickResult.withSubchannel(subchannel);
  }
  if (status != null) {
    return PickResult.withError(status);
  }
  return PickResult.withNoResult();
}

void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo newState) {
  if (newState.getState() == SHUTDOWN || !(subchannels.values().contains(subchannel))) {
    return;
  }
  if (newState.getState() == IDLE) {
    subchannel.requestConnection();
  }
  subchannel.getAttributes().get(STATE_INFO).set(newState);
  maybeStartFallbackTimer();
  maybeUpdatePicker();
}

/**
 * Start the fallback timer if it's not already started and all connections are lost.
 */
private void maybeStartFallbackTimer() {
  if (fallbackTimer != null) {
    return;
  }
  if (fallbackBackendList.isEmpty()) {
    return;
  }
  if (balancerWorking) {
    return;
  }
  if (usingFallbackBackends) {
    return;
  }
  int numReadySubchannels = 0;
  for (Subchannel subchannel : subchannels.values()) {
    if (subchannel.getAttributes().get(STATE_INFO).get().getState() == READY) {
      numReadySubchannels++;
    }
  }
  if (numReadySubchannels > 0) {
    return;
  }
  logger.log(Level.FINE, "[{0}] Starting fallback timer.", new Object[] {logId});
  fallbackTimer = new FallbackModeTask();
  fallbackTimer.schedule();
}

void shutdown() {
  shutdownLbComm();
  for (Subchannel subchannel : subchannels.values()) {
    subchannel.shutdown();
  }
  subchannels = Collections.emptyMap();
  cancelFallbackTimer();
}

/**
 * Make and use a picker out of the current lists and the states of subchannels if they have
 * changed since the last picker created.
 */
private void maybeUpdatePicker() {
  List<RoundRobinEntry> pickList = new ArrayList<RoundRobinEntry>(backendList.size());
  Status error = null;
  boolean hasIdle = false;
  for (BackendEntry entry : backendList) {
    Subchannel subchannel = entry.result.getSubchannel();
    Attributes attrs = subchannel.getAttributes();
    ConnectivityStateInfo stateInfo = attrs.get(STATE_INFO).get();
    if (stateInfo.getState() == READY) {
      pickList.add(entry);
    } else if (stateInfo.getState() == TRANSIENT_FAILURE) {
      error = stateInfo.getStatus();
    } else if (stateInfo.getState() == IDLE) {
      hasIdle = true;
    }
  }
  ConnectivityState state;
  if (pickList.isEmpty()) {
    if (error != null && !hasIdle) {
      logger.log(Level.FINE, "[{0}] No ready Subchannel. Using error: {1}",
          new Object[] {logId, error});
      pickList.add(new ErrorEntry(error));
      state = TRANSIENT_FAILURE;
    } else {
      logger.log(Level.FINE, "[{0}] No ready Subchannel and still connecting", logId);
      pickList.add(BUFFER_ENTRY);
      state = CONNECTING;
    }
  } else {
    logger.log(
        Level.FINE, "[{0}] Using drop list {1} and pick list {2}",
        new Object[] {logId, dropList, pickList});
    state = READY;
  }
  maybeUpdatePicker(state, new RoundRobinPicker(dropList, pickList));
}

@After
public void tearDown() {
  try {
    if (balancer != null) {
      channelExecutor.execute(new Runnable() {
          @Override
          public void run() {
            balancer.shutdown();
          }
        });
    }
    for (ManagedChannel channel : oobChannelTracker) {
      assertTrue(channel + " is shutdown", channel.isShutdown());
      // balancer should have closed the LB stream, terminating the OOB channel.
      assertTrue(channel + " is terminated", channel.isTerminated());
    }
    for (Subchannel subchannel: subchannelTracker) {
      verify(subchannel).shutdown();
    }
    // No timer should linger after shutdown
    assertEquals(0, fakeClock.numPendingTasks());
  } finally {
    if (fakeLbServer != null) {
      fakeLbServer.shutdownNow();
    }
  }
}

@Test
public void roundRobinPickerNoDrop() {
  GrpclbClientLoadRecorder loadRecorder = new GrpclbClientLoadRecorder(timeProvider);
  Subchannel subchannel = mock(Subchannel.class);
  BackendEntry b1 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0001");
  BackendEntry b2 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0002");

  List<BackendEntry> pickList = Arrays.asList(b1, b2);
  RoundRobinPicker picker = new RoundRobinPicker(Collections.<DropEntry>emptyList(), pickList);

  PickSubchannelArgs args1 = mock(PickSubchannelArgs.class);
  Metadata headers1 = new Metadata();
  // The existing token on the headers will be replaced
  headers1.put(GrpclbConstants.TOKEN_METADATA_KEY, "LBTOKEN__OLD");
  when(args1.getHeaders()).thenReturn(headers1);
  assertSame(b1.result, picker.pickSubchannel(args1));
  verify(args1).getHeaders();
  assertThat(headers1.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");

  PickSubchannelArgs args2 = mock(PickSubchannelArgs.class);
  Metadata headers2 = new Metadata();
  when(args2.getHeaders()).thenReturn(headers2);
  assertSame(b2.result, picker.pickSubchannel(args2));
  verify(args2).getHeaders();
  assertThat(headers2.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0002");

  PickSubchannelArgs args3 = mock(PickSubchannelArgs.class);
  Metadata headers3 = new Metadata();
  when(args3.getHeaders()).thenReturn(headers3);
  assertSame(b1.result, picker.pickSubchannel(args3));
  verify(args3).getHeaders();
  assertThat(headers3.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");

  verify(subchannel, never()).getAttributes();
}

private List<Subchannel> fallbackTestVerifyUseOfBalancerBackendLists(
    InOrder inOrder, Helper helper, List<ServerEntry> servers) {
  ArrayList<EquivalentAddressGroup> addrs = new ArrayList<EquivalentAddressGroup>();
  ArrayList<String> tokens = new ArrayList<String>();
  for (ServerEntry server : servers) {
    addrs.add(new EquivalentAddressGroup(server.addr));
    tokens.add(server.token);
  }
  return fallbackTestVerifyUseOfBackendLists(inOrder, helper, addrs, tokens);
}

private void deliverSubchannelState(
    final Subchannel subchannel, final ConnectivityStateInfo newState) {
  channelExecutor.execute(new Runnable() {
      @Override
      public void run() {
        balancer.handleSubchannelState(subchannel, newState);
      }
    });
}

@Override
public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
  if (subchannels.get(subchannel.getAddresses()) != subchannel) {
    return;
  }
  if (stateInfo.getState() == IDLE) {
    subchannel.requestConnection();
  }
  getSubchannelStateInfoRef(subchannel).value = stateInfo;
  updateBalancingState(getAggregatedState(), getAggregatedError());
}

/**
 * Filters out non-ready subchannels.
 */
private static List<Subchannel> filterNonFailingSubchannels(
    Collection<Subchannel> subchannels) {
  List<Subchannel> readySubchannels = new ArrayList<Subchannel>(subchannels.size());
  for (Subchannel subchannel : subchannels) {
    if (getSubchannelStateInfoRef(subchannel).value.getState() == READY) {
      readySubchannels.add(subchannel);
    }
  }
  return readySubchannels;
}

/**
 * If all subchannels are TRANSIENT_FAILURE, return the Status associated with an arbitrary
 * subchannel otherwise, return null.
 */
@Nullable
private Status getAggregatedError() {
  Status status = null;
  for (Subchannel subchannel : getSubchannels()) {
    ConnectivityStateInfo stateInfo = getSubchannelStateInfoRef(subchannel).value;
    if (stateInfo.getState() != TRANSIENT_FAILURE) {
      return null;
    }
    status = stateInfo.getStatus();
  }
  return status;
}

private Subchannel nextSubchannel() {
  if (list.isEmpty()) {
    throw new NoSuchElementException();
  }
  int size = list.size();

  int i = indexUpdater.incrementAndGet(this);
  if (i >= size) {
    int oldi = i;
    i %= size;
    indexUpdater.compareAndSet(this, oldi, i);
  }
  return list.get(i);
}

@Test
public void subchannelsWhenChannelShutdownNow() {
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);
  Subchannel sub1 = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  Subchannel sub2 = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  sub1.requestConnection();
  sub2.requestConnection();

  assertEquals(2, transports.size());
  MockClientTransportInfo ti1 = transports.poll();
  MockClientTransportInfo ti2 = transports.poll();

  ti1.listener.transportReady();
  ti2.listener.transportReady();

  channel.shutdownNow();
  verify(ti1.transport).shutdownNow(any(Status.class));
  verify(ti2.transport).shutdownNow(any(Status.class));

  ti1.listener.transportShutdown(Status.UNAVAILABLE.withDescription("shutdown now"));
  ti2.listener.transportShutdown(Status.UNAVAILABLE.withDescription("shutdown now"));
  ti1.listener.transportTerminated();

  assertFalse(channel.isTerminated());
  ti2.listener.transportTerminated();
  assertTrue(channel.isTerminated());
}

@Test
public void subchannelsNoConnectionShutdown() {
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);
  Subchannel sub1 = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  Subchannel sub2 = helper.createSubchannel(addressGroup, Attributes.EMPTY);

  channel.shutdown();
  verify(mockLoadBalancer).shutdown();
  sub1.shutdown();
  assertFalse(channel.isTerminated());
  sub2.shutdown();
  assertTrue(channel.isTerminated());
  verify(mockTransportFactory, never()).newClientTransport(any(SocketAddress.class), anyString(),
      anyString(), any(ProxyParameters.class));
}

private void subtestRefreshNameResolutionWhenConnectionFailed(boolean isOobChannel) {
  FakeNameResolverFactory nameResolverFactory = new FakeNameResolverFactory(true);
  createChannel(nameResolverFactory, NO_INTERCEPTOR);
  FakeNameResolverFactory.FakeNameResolver resolver = nameResolverFactory.resolvers.get(0);

  if (isOobChannel) {
    OobChannel oobChannel = (OobChannel) helper.createOobChannel(addressGroup, "oobAuthority");
    oobChannel.getSubchannel().requestConnection();
  } else {
    Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
    subchannel.requestConnection();
  }

  MockClientTransportInfo transportInfo = transports.poll();
  assertNotNull(transportInfo);

  // Transport closed when connecting
  assertEquals(0, resolver.refreshCalled);
  transportInfo.listener.transportShutdown(Status.UNAVAILABLE);
  assertEquals(1, resolver.refreshCalled);

  timer.forwardNanos(RECONNECT_BACKOFF_INTERVAL_NANOS);
  transportInfo = transports.poll();
  assertNotNull(transportInfo);

  transportInfo.listener.transportReady();

  // Transport closed when ready
  assertEquals(1, resolver.refreshCalled);
  transportInfo.listener.transportShutdown(Status.UNAVAILABLE);
  assertEquals(2, resolver.refreshCalled);
}

@Test
public void pickerReturnsStreamTracer_noDelay() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);
  Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  subchannel.requestConnection();
  MockClientTransportInfo transportInfo = transports.poll();
  transportInfo.listener.transportReady();
  ClientTransport mockTransport = transportInfo.transport;
  when(mockTransport.newStream(
          any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mockStream);

  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(
      PickResult.withSubchannel(subchannel, factory2));
  helper.updateBalancingState(READY, mockPicker);

  CallOptions callOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory1);
  ClientCall<String, Integer> call = channel.newCall(method, callOptions);
  call.start(mockCallListener, new Metadata());

  verify(mockPicker).pickSubchannel(any(PickSubchannelArgs.class));
  verify(mockTransport).newStream(same(method), any(Metadata.class), callOptionsCaptor.capture());
  assertEquals(
      Arrays.asList(factory1, factory2),
      callOptionsCaptor.getValue().getStreamTracerFactories());
  // The factories are safely not stubbed because we do not expect any usage of them.
  verifyZeroInteractions(factory1);
  verifyZeroInteractions(factory2);
}

@Test
public void pickerReturnsStreamTracer_delayed() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);

  CallOptions callOptions = CallOptions.DEFAULT.withStreamTracerFactory(factory1);
  ClientCall<String, Integer> call = channel.newCall(method, callOptions);
  call.start(mockCallListener, new Metadata());

  Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  subchannel.requestConnection();
  MockClientTransportInfo transportInfo = transports.poll();
  transportInfo.listener.transportReady();
  ClientTransport mockTransport = transportInfo.transport;
  when(mockTransport.newStream(
          any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mockStream);
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class))).thenReturn(
      PickResult.withSubchannel(subchannel, factory2));

  helper.updateBalancingState(READY, mockPicker);
  assertEquals(1, executor.runDueTasks());

  verify(mockPicker).pickSubchannel(any(PickSubchannelArgs.class));
  verify(mockTransport).newStream(same(method), any(Metadata.class), callOptionsCaptor.capture());
  assertEquals(
      Arrays.asList(factory1, factory2),
      callOptionsCaptor.getValue().getStreamTracerFactories());
  // The factories are safely not stubbed because we do not expect any usage of them.
  verifyZeroInteractions(factory1);
  verifyZeroInteractions(factory2);
}

@Test
public void updateBalancingStateDoesUpdatePicker() {
  ClientStream mockStream = mock(ClientStream.class);
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);

  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata());

  // Make the transport available with subchannel2
  Subchannel subchannel1 = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  Subchannel subchannel2 = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  subchannel2.requestConnection();

  MockClientTransportInfo transportInfo = transports.poll();
  ConnectionClientTransport mockTransport = transportInfo.transport;
  ManagedClientTransport.Listener transportListener = transportInfo.listener;
  when(mockTransport.newStream(same(method), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mockStream);
  transportListener.transportReady();

  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel1));
  helper.updateBalancingState(READY, mockPicker);

  executor.runDueTasks();
  verify(mockTransport, never())
      .newStream(any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class));
  verify(mockStream, never()).start(any(ClientStreamListener.class));


  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel2));
  helper.updateBalancingState(READY, mockPicker);

  executor.runDueTasks();
  verify(mockTransport).newStream(same(method), any(Metadata.class), any(CallOptions.class));
  verify(mockStream).start(any(ClientStreamListener.class));
}

@Test
public void channelsAndSubchannels_instrumented_name() throws Exception {
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);
  assertEquals(target, getStats(channel).target);

  Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  assertEquals(target, getStats((AbstractSubchannel) subchannel).target);
}

@Test
public void channelStat_callEndSuccess() throws Exception {
  // set up
  Metadata headers = new Metadata();
  ClientStream mockStream = mock(ClientStream.class);
  createChannel(new FakeNameResolverFactory(true), NO_INTERCEPTOR);

  // Start a call with a call executor
  CallOptions options =
      CallOptions.DEFAULT.withExecutor(executor.getScheduledExecutorService());
  ClientCall<String, Integer> call = channel.newCall(method, options);
  call.start(mockCallListener, headers);

  // Make the transport available
  Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  subchannel.requestConnection();
  MockClientTransportInfo transportInfo = transports.poll();
  ConnectionClientTransport mockTransport = transportInfo.transport;
  ManagedClientTransport.Listener transportListener = transportInfo.listener;
  when(mockTransport.newStream(same(method), same(headers), any(CallOptions.class)))
      .thenReturn(mockStream);
  transportListener.transportReady();
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel));
  helper.updateBalancingState(READY, mockPicker);

  executor.runDueTasks();
  verify(mockStream).start(streamListenerCaptor.capture());
  // end set up

  // the actual test
  ClientStreamListener streamListener = streamListenerCaptor.getValue();
  call.halfClose();
  assertEquals(0, getStats(channel).callsSucceeded);
  assertEquals(0, getStats(channel).callsFailed);
  streamListener.closed(Status.OK, new Metadata());
  executor.runDueTasks();
  assertEquals(1, getStats(channel).callsSucceeded);
  assertEquals(0, getStats(channel).callsFailed);
}

@Before
public void setUp() {
  MockitoAnnotations.initMocks(this);

  for (int i = 0; i < 3; i++) {
    SocketAddress addr = new FakeSocketAddress("server" + i);
    EquivalentAddressGroup eag = new EquivalentAddressGroup(addr);
    servers.add(eag);
    Subchannel sc = mock(Subchannel.class);
    when(sc.getAddresses()).thenReturn(eag);
    subchannels.put(eag, sc);
  }

  when(mockHelper.createSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class)))
      .then(new Answer<Subchannel>() {
        @Override
        public Subchannel answer(InvocationOnMock invocation) throws Throwable {
          Object[] args = invocation.getArguments();
          Subchannel subchannel = subchannels.get(args[0]);
          when(subchannel.getAttributes()).thenReturn((Attributes) args[1]);
          return subchannel;
        }
      });

  loadBalancer = (RoundRobinLoadBalancer) RoundRobinLoadBalancerFactory.getInstance()
      .newLoadBalancer(mockHelper);
}

@Test
public void pickAfterStateChange() throws Exception {
  InOrder inOrder = inOrder(mockHelper);
  loadBalancer.handleResolvedAddressGroups(servers, Attributes.EMPTY);
  Subchannel subchannel = loadBalancer.getSubchannels().iterator().next();
  Ref<ConnectivityStateInfo> subchannelStateInfo = subchannel.getAttributes().get(
      STATE_INFO);

  inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), isA(Picker.class));
  assertThat(subchannelStateInfo.value).isEqualTo(ConnectivityStateInfo.forNonError(IDLE));

  loadBalancer.handleSubchannelState(subchannel,
      ConnectivityStateInfo.forNonError(READY));
  inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
  assertNull(pickerCaptor.getValue().getStatus());
  assertThat(subchannelStateInfo.value).isEqualTo(
      ConnectivityStateInfo.forNonError(READY));

  Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯");
  loadBalancer.handleSubchannelState(subchannel,
      ConnectivityStateInfo.forTransientFailure(error));
  assertThat(subchannelStateInfo.value).isEqualTo(
      ConnectivityStateInfo.forTransientFailure(error));
  inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
  assertNull(pickerCaptor.getValue().getStatus());

  loadBalancer.handleSubchannelState(subchannel,
      ConnectivityStateInfo.forNonError(IDLE));
  inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
  assertNull(pickerCaptor.getValue().getStatus());
  assertThat(subchannelStateInfo.value).isEqualTo(
      ConnectivityStateInfo.forNonError(IDLE));

  verify(subchannel, times(2)).requestConnection();
  verify(mockHelper, times(3)).createSubchannel(any(EquivalentAddressGroup.class),
      any(Attributes.class));
  verifyNoMoreInteractions(mockHelper);
}

@Test
public void pickerEmptyList() throws Exception {
  Picker picker = new Picker(Lists.<Subchannel>newArrayList(), Status.UNKNOWN);

  assertEquals(null, picker.pickSubchannel(mockArgs).getSubchannel());
  assertEquals(Status.UNKNOWN,
      picker.pickSubchannel(mockArgs).getStatus());
}

GrpcRoutePicker(List<Subchannel> list, Status status, Attributes nameResovleCache) {
  this.list = list;
  this.size = list.size();
  this.status = status;
  this.nameResovleCache = nameResovleCache;
}

/**
 * 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.loadRecorder = checkNotNull(loadRecorder, "loadRecorder");
  this.token = checkNotNull(token, "token");
}

/**
 * Creates a BackendEntry whose usage will not be reported.
 */
BackendEntry(Subchannel subchannel) {
  this.result = PickResult.withSubchannel(subchannel);
  this.loadRecorder = null;
  this.token = null;
}

@Test
public void roundRobinPickerWithDrop() {
  assertTrue(DROP_PICK_RESULT.isDrop());
  GrpclbClientLoadRecorder loadRecorder = new GrpclbClientLoadRecorder(timeProvider);
  Subchannel subchannel = mock(Subchannel.class);
  // 1 out of 2 requests are to be dropped
  DropEntry d = new DropEntry(loadRecorder, "LBTOKEN0003");
  List<DropEntry> dropList = Arrays.asList(null, d);

  BackendEntry b1 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0001");
  BackendEntry b2 = new BackendEntry(subchannel, loadRecorder, "LBTOKEN0002");
  List<BackendEntry> pickList = Arrays.asList(b1, b2);
  RoundRobinPicker picker = new RoundRobinPicker(dropList, pickList);

  // dropList[0], pickList[0]
  PickSubchannelArgs args1 = mock(PickSubchannelArgs.class);
  Metadata headers1 = new Metadata();
  headers1.put(GrpclbConstants.TOKEN_METADATA_KEY, "LBTOKEN__OLD");
  when(args1.getHeaders()).thenReturn(headers1);
  assertSame(b1.result, picker.pickSubchannel(args1));
  verify(args1).getHeaders();
  assertThat(headers1.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");

  // dropList[1]: drop
  PickSubchannelArgs args2 = mock(PickSubchannelArgs.class);
  Metadata headers2 = new Metadata();
  when(args2.getHeaders()).thenReturn(headers2);
  assertSame(DROP_PICK_RESULT, picker.pickSubchannel(args2));
  verify(args2, never()).getHeaders();

  // dropList[0], pickList[1]
  PickSubchannelArgs args3 = mock(PickSubchannelArgs.class);
  Metadata headers3 = new Metadata();
  when(args3.getHeaders()).thenReturn(headers3);
  assertSame(b2.result, picker.pickSubchannel(args3));
  verify(args3).getHeaders();
  assertThat(headers3.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0002");

  // dropList[1]: drop
  PickSubchannelArgs args4 = mock(PickSubchannelArgs.class);
  Metadata headers4 = new Metadata();
  when(args4.getHeaders()).thenReturn(headers4);
  assertSame(DROP_PICK_RESULT, picker.pickSubchannel(args4));
  verify(args4, never()).getHeaders();

  // dropList[0], pickList[0]
  PickSubchannelArgs args5 = mock(PickSubchannelArgs.class);
  Metadata headers5 = new Metadata();
  when(args5.getHeaders()).thenReturn(headers5);
  assertSame(b1.result, picker.pickSubchannel(args5));
  verify(args5).getHeaders();
  assertThat(headers5.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");

  verify(subchannel, never()).getAttributes();
}

@Test
public void resetGrpclbWhenSwitchingAwayFromGrpclb() {
  InOrder inOrder = inOrder(helper);
  List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
  Attributes grpclbResolutionAttrs = Attributes.newBuilder()
      .set(GrpclbConstants.ATTR_LB_POLICY, LbPolicy.GRPCLB).build();
  deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);

  assertSame(LbPolicy.GRPCLB, balancer.getLbPolicy());
  assertNull(balancer.getDelegate());
  verify(helper).createOobChannel(addrsEq(grpclbResolutionList.get(0)), eq(lbAuthority(0)));
  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> backends = Arrays.asList(new ServerEntry("127.0.0.1", 2000, "token0001"));
  inOrder.verify(helper, never())
      .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
  lbResponseObserver.onNext(buildInitialResponse());
  lbResponseObserver.onNext(buildLbResponse(backends));

  inOrder.verify(helper).createSubchannel(
      eq(new EquivalentAddressGroup(backends.get(0).addr)), any(Attributes.class));
  assertEquals(1, mockSubchannels.size());
  Subchannel subchannel = mockSubchannels.poll();
  verify(subchannel).requestConnection();

  // Switch to round-robin. GRPCLB streams and connections should be closed. 
  List<EquivalentAddressGroup> roundRobinResolutionList =
      createResolvedServerAddresses(false, false, false);
  Attributes roundRobinResolutionAttrs = Attributes.newBuilder()
      .set(GrpclbConstants.ATTR_LB_POLICY, LbPolicy.ROUND_ROBIN).build();
  verify(lbRequestObserver, never()).onCompleted();
  verify(subchannel, never()).shutdown();
  assertFalse(oobChannel.isShutdown());
  deliverResolvedAddresses(roundRobinResolutionList, roundRobinResolutionAttrs);

  verify(lbRequestObserver).onCompleted();
  verify(subchannel).shutdown();
  assertTrue(oobChannel.isShutdown());
  assertTrue(oobChannel.isTerminated());
  assertSame(LbPolicy.ROUND_ROBIN, balancer.getLbPolicy());
  assertSame(roundRobinBalancer, balancer.getDelegate());
  assertNull(balancer.getGrpclbState());
}

private List<Subchannel> fallbackTestVerifyUseOfFallbackBackendLists(
    InOrder inOrder, Helper helper, List<EquivalentAddressGroup> addrs) {
  return fallbackTestVerifyUseOfBackendLists(inOrder, helper, addrs, null);
}

/**
 * Returns a transport out of a PickResult, or {@code null} if the result is "buffer".
 */
@Nullable
static ClientTransport getTransportFromPickResult(PickResult result, boolean isWaitForReady) {
  final ClientTransport transport;
  Subchannel subchannel = result.getSubchannel();
  if (subchannel != null) {
    transport = ((AbstractSubchannel) subchannel).obtainActiveTransport();
  } else {
    transport = null;
  }
  if (transport != null) {
    final ClientStreamTracer.Factory streamTracerFactory = result.getStreamTracerFactory();
    if (streamTracerFactory == null) {
      return transport;
    }
    return new ClientTransport() {
      @Override
      public ClientStream newStream(
          MethodDescriptor<?, ?> method, Metadata headers, CallOptions callOptions) {
        return transport.newStream(
            method, headers, callOptions.withStreamTracerFactory(streamTracerFactory));
      }

      @Override
      public void ping(PingCallback callback, Executor executor) {
        transport.ping(callback, executor);
      }

      @Override
      public LogId getLogId() {
        return transport.getLogId();
      }

      @Override
      public ListenableFuture<TransportStats> getStats() {
        return transport.getStats();
      }
    };
  }
  if (!result.getStatus().isOk() && (result.isDrop() || !isWaitForReady)) {
    return new FailingClientTransport(result.getStatus());
  }
  return null;
}

@VisibleForTesting
Subchannel getSubchannel() {
  return subchannelImpl;
}