Java Code Examples for io.grpc.LoadBalancer.Subchannel#requestConnection()

@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 = createSubchannelSafely(helper, servers.get(0), Attributes.EMPTY);

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

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

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

@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 = createSubchannelSafely(helper, 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 subchannelChannel_failWaitForReady() {
  createChannel();
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  Channel sChannel = subchannel.asChannel();
  Metadata headers = new Metadata();

  // Subchannel must be READY when creating the RPC.
  subchannel.requestConnection();
  verify(mockTransportFactory)
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo = transports.poll();
  ConnectionClientTransport mockTransport = transportInfo.transport;
  ManagedClientTransport.Listener transportListener = transportInfo.listener;
  transportListener.transportReady();
  assertEquals(0, balancerRpcExecutor.numPendingTasks());

  // Wait-for-ready RPC is not allowed
  ClientCall<String, Integer> call =
      sChannel.newCall(method, CallOptions.DEFAULT.withWaitForReady());
  call.start(mockCallListener, headers);
  verify(mockTransport, never()).newStream(
      any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class));

  verifyZeroInteractions(mockCallListener);
  assertEquals(1, balancerRpcExecutor.runDueTasks());
  verify(mockCallListener).onClose(
      same(SubchannelChannel.WAIT_FOR_READY_ERROR), any(Metadata.class));
}
 

@Test
public void updateBalancingStateDoesUpdatePicker() {
  ClientStream mockStream = mock(ClientStream.class);
  createChannel();

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

  // Make the transport available with subchannel2
  Subchannel subchannel1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  Subchannel subchannel2 = createSubchannelSafely(helper, 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 pickerReturnsStreamTracer_delayed() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel();

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

  Subchannel subchannel = createSubchannelSafely(helper, 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 pickerReturnsStreamTracer_noDelay() {
  ClientStream mockStream = mock(ClientStream.class);
  ClientStreamTracer.Factory factory1 = mock(ClientStreamTracer.Factory.class);
  ClientStreamTracer.Factory factory2 = mock(ClientStreamTracer.Factory.class);
  createChannel();
  Subchannel subchannel = createSubchannelSafely(helper, 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);
}
 

private void subtestRefreshNameResolutionWhenConnectionFailed(boolean isOobChannel) {
  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(socketAddress)))
          .build();
  channelBuilder.nameResolverFactory(nameResolverFactory);
  createChannel();
  FakeNameResolverFactory.FakeNameResolver resolver = nameResolverFactory.resolvers.get(0);

  if (isOobChannel) {
    OobChannel oobChannel = (OobChannel) helper.createOobChannel(addressGroup, "oobAuthority");
    oobChannel.getSubchannel().requestConnection();
  } else {
    Subchannel subchannel = createSubchannelSafely(helper, 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);
}
 

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);
  maybeUseFallbackBackends();
  maybeUpdatePicker();
}
 

@Deprecated
@Test
@SuppressWarnings("deprecation")
public void createSubchannel_old_propagateSubchannelStatesToOldApi() {
  createChannel();

  Subchannel subchannel = helper.createSubchannel(addressGroup, Attributes.EMPTY);
  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)));

  channel.shutdown();
  verify(mockLoadBalancer).shutdown();
  subchannel.shutdown();

  verify(mockLoadBalancer, atLeast(0)).canHandleEmptyAddressListFromNameResolution();
  verify(mockLoadBalancer, atLeast(0)).handleNameResolutionError(any(Status.class));
  // handleSubchannelState() should not be called after shutdown()
  verifyNoMoreInteractions(mockLoadBalancer);
}
 

@Test
public void subchannelChannel_normalUsage() {
  createChannel();
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  verify(balancerRpcExecutorPool, never()).getObject();

  Channel sChannel = subchannel.asChannel();
  verify(balancerRpcExecutorPool).getObject();

  Metadata headers = new Metadata();
  CallOptions callOptions = CallOptions.DEFAULT.withDeadlineAfter(5, TimeUnit.SECONDS);

  // Subchannel must be READY when creating the RPC.
  subchannel.requestConnection();
  verify(mockTransportFactory)
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo = transports.poll();
  ConnectionClientTransport mockTransport = transportInfo.transport;
  ManagedClientTransport.Listener transportListener = transportInfo.listener;
  transportListener.transportReady();

  ClientCall<String, Integer> call = sChannel.newCall(method, callOptions);
  call.start(mockCallListener, headers);
  verify(mockTransport).newStream(same(method), same(headers), callOptionsCaptor.capture());

  CallOptions capturedCallOption = callOptionsCaptor.getValue();
  assertThat(capturedCallOption.getDeadline()).isSameAs(callOptions.getDeadline());
  assertThat(capturedCallOption.getOption(GrpcUtil.CALL_OPTIONS_RPC_OWNED_BY_BALANCER)).isTrue();
}
 

@Test
public void subchannelsWhenChannelShutdownNow() {
  createChannel();
  Subchannel sub1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  Subchannel sub2 = createSubchannelSafely(helper, 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());
}
 

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

@Test
public void subchannels() {
  createChannel();

  // 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();
  Subchannel sub1 = createSubchannelSafely(helper, addressGroup, attrs1);
  Subchannel sub2 = createSubchannelSafely(helper, addressGroup, attrs2);
  assertNotSame(sub1, sub2);
  assertNotSame(attrs1, attrs2);
  assertSame(attrs1, sub1.getAttributes());
  assertSame(attrs2, sub2.getAttributes());
  assertSame(addressGroup, sub1.getAddresses());
  assertSame(addressGroup, sub2.getAddresses());

  // requestConnection()
  verify(mockTransportFactory, never())
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  sub1.requestConnection();
  verify(mockTransportFactory).newClientTransport(socketAddress, clientTransportOptions);
  MockClientTransportInfo transportInfo1 = transports.poll();
  assertNotNull(transportInfo1);

  sub2.requestConnection();
  verify(mockTransportFactory, times(2))
      .newClientTransport(socketAddress, clientTransportOptions);
  MockClientTransportInfo transportInfo2 = transports.poll();
  assertNotNull(transportInfo2);

  sub1.requestConnection();
  sub2.requestConnection();
  verify(mockTransportFactory, times(2))
      .newClientTransport(socketAddress, clientTransportOptions);

  // shutdown() has a delay
  sub1.shutdown();
  timer.forwardTime(ManagedChannelImpl.SUBCHANNEL_SHUTDOWN_DELAY_SECONDS - 1, TimeUnit.SECONDS);
  sub1.shutdown();
  verify(transportInfo1.transport, never()).shutdown(any(Status.class));
  timer.forwardTime(1, TimeUnit.SECONDS);
  verify(transportInfo1.transport).shutdown(same(ManagedChannelImpl.SUBCHANNEL_SHUTDOWN_STATUS));

  // ... but not after Channel is terminating
  verify(mockLoadBalancer, never()).shutdown();
  channel.shutdown();
  verify(mockLoadBalancer).shutdown();
  verify(transportInfo2.transport, never()).shutdown(any(Status.class));

  sub2.shutdown();
  verify(transportInfo2.transport).shutdown(same(ManagedChannelImpl.SHUTDOWN_STATUS));

  // Cleanup
  transportInfo1.listener.transportShutdown(Status.UNAVAILABLE);
  transportInfo1.listener.transportTerminated();
  transportInfo2.listener.transportShutdown(Status.UNAVAILABLE);
  transportInfo2.listener.transportTerminated();
  timer.forwardTime(ManagedChannelImpl.SUBCHANNEL_SHUTDOWN_DELAY_SECONDS, TimeUnit.SECONDS);
}
 

/**
 * Verify that if all resolved addresses failed to connect, a fail-fast call will fail, while a
 * wait-for-ready call will still be buffered.
 */
@Test
public void allServersFailedToConnect() throws Exception {
  final SocketAddress addr1 = new SocketAddress() {
      @Override public String toString() {
        return "addr1";
      }
    };
  final SocketAddress addr2 = new SocketAddress() {
      @Override public String toString() {
        return "addr2";
      }
    };
  InOrder inOrder = inOrder(mockLoadBalancer);

  List<SocketAddress> resolvedAddrs = Arrays.asList(addr1, addr2);

  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(resolvedAddrs)))
          .build();
  channelBuilder.nameResolverFactory(nameResolverFactory);
  createChannel();

  // Start a wait-for-ready call
  ClientCall<String, Integer> call =
      channel.newCall(method, CallOptions.DEFAULT.withWaitForReady());
  Metadata headers = new Metadata();
  call.start(mockCallListener, headers);
  // ... and a fail-fast call
  ClientCall<String, Integer> call2 =
      channel.newCall(method, CallOptions.DEFAULT.withoutWaitForReady());
  call2.start(mockCallListener2, headers);
  executor.runDueTasks();

  // Simulate name resolution results
  EquivalentAddressGroup addressGroup = new EquivalentAddressGroup(resolvedAddrs);
  inOrder.verify(mockLoadBalancer).handleResolvedAddressGroups(
      eq(Arrays.asList(addressGroup)), eq(Attributes.EMPTY));
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel));
  subchannel.requestConnection();

  inOrder.verify(mockLoadBalancer).handleSubchannelState(
      same(subchannel), stateInfoCaptor.capture());
  assertEquals(CONNECTING, stateInfoCaptor.getValue().getState());

  // Connecting to server1, which will fail
  verify(mockTransportFactory)
      .newClientTransport(same(addr1), any(ClientTransportOptions.class));
  verify(mockTransportFactory, times(0))
      .newClientTransport(same(addr2), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo1 = transports.poll();
  transportInfo1.listener.transportShutdown(Status.UNAVAILABLE);

  // Connecting to server2, which will fail too
  verify(mockTransportFactory)
      .newClientTransport(same(addr2), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo2 = transports.poll();
  Status server2Error = Status.UNAVAILABLE.withDescription("Server2 failed to connect");
  transportInfo2.listener.transportShutdown(server2Error);

  // ... which makes the subchannel enter TRANSIENT_FAILURE. The last error Status is propagated
  // to LoadBalancer.
  inOrder.verify(mockLoadBalancer).handleSubchannelState(
      same(subchannel), stateInfoCaptor.capture());
  assertEquals(TRANSIENT_FAILURE, stateInfoCaptor.getValue().getState());
  assertSame(server2Error, stateInfoCaptor.getValue().getStatus());

  // A typical LoadBalancer would create a picker with error
  SubchannelPicker picker2 = mock(SubchannelPicker.class);
  when(picker2.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withError(server2Error));
  helper.updateBalancingState(TRANSIENT_FAILURE, picker2);
  executor.runDueTasks();

  // ... which fails the fail-fast call
  verify(mockCallListener2).onClose(same(server2Error), any(Metadata.class));
  // ... while the wait-for-ready call stays
  verifyNoMoreInteractions(mockCallListener);
  // No real stream was ever created
  verify(transportInfo1.transport, times(0))
      .newStream(any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class));
  verify(transportInfo2.transport, times(0))
      .newStream(any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class));
}
 

/**
 * Verify that if the first resolved address points to a server that cannot be connected, the call
 * will end up with the second address which works.
 */
@Test
public void firstResolvedServerFailedToConnect() throws Exception {
  final SocketAddress goodAddress = new SocketAddress() {
      @Override public String toString() {
        return "goodAddress";
      }
    };
  final SocketAddress badAddress = new SocketAddress() {
      @Override public String toString() {
        return "badAddress";
      }
    };
  InOrder inOrder = inOrder(mockLoadBalancer);

  List<SocketAddress> resolvedAddrs = Arrays.asList(badAddress, goodAddress);
  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(resolvedAddrs)))
          .build();
  channelBuilder.nameResolverFactory(nameResolverFactory);
  createChannel();

  // Start the call
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  Metadata headers = new Metadata();
  call.start(mockCallListener, headers);
  executor.runDueTasks();

  // Simulate name resolution results
  EquivalentAddressGroup addressGroup = new EquivalentAddressGroup(resolvedAddrs);
  inOrder.verify(mockLoadBalancer).handleResolvedAddressGroups(
      eq(Arrays.asList(addressGroup)), eq(Attributes.EMPTY));
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel));
  subchannel.requestConnection();
  inOrder.verify(mockLoadBalancer).handleSubchannelState(
      same(subchannel), stateInfoCaptor.capture());
  assertEquals(CONNECTING, stateInfoCaptor.getValue().getState());

  // The channel will starts with the first address (badAddress)
  verify(mockTransportFactory)
      .newClientTransport(same(badAddress), any(ClientTransportOptions.class));
  verify(mockTransportFactory, times(0))
      .newClientTransport(same(goodAddress), any(ClientTransportOptions.class));

  MockClientTransportInfo badTransportInfo = transports.poll();
  // Which failed to connect
  badTransportInfo.listener.transportShutdown(Status.UNAVAILABLE);
  inOrder.verifyNoMoreInteractions();

  // The channel then try the second address (goodAddress)
  verify(mockTransportFactory)
      .newClientTransport(same(goodAddress), any(ClientTransportOptions.class));
  MockClientTransportInfo goodTransportInfo = transports.poll();
  when(goodTransportInfo.transport.newStream(
          any(MethodDescriptor.class), any(Metadata.class), any(CallOptions.class)))
      .thenReturn(mock(ClientStream.class));

  goodTransportInfo.listener.transportReady();
  inOrder.verify(mockLoadBalancer).handleSubchannelState(
      same(subchannel), stateInfoCaptor.capture());
  assertEquals(READY, stateInfoCaptor.getValue().getState());

  // A typical LoadBalancer will call this once the subchannel becomes READY
  helper.updateBalancingState(READY, mockPicker);
  // Delayed transport uses the app executor to create real streams.
  executor.runDueTasks();

  verify(goodTransportInfo.transport).newStream(same(method), same(headers),
      same(CallOptions.DEFAULT));
  // The bad transport was never used.
  verify(badTransportInfo.transport, times(0)).newStream(any(MethodDescriptor.class),
      any(Metadata.class), any(CallOptions.class));
}
 

@Test
public void idleMode_resetsDelayedTransportPicker() {
  ClientStream mockStream = mock(ClientStream.class);
  Status pickError = Status.UNAVAILABLE.withDescription("pick result error");
  long idleTimeoutMillis = 1000L;
  channelBuilder.idleTimeout(idleTimeoutMillis, TimeUnit.MILLISECONDS);
  channelBuilder.nameResolverFactory(
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(socketAddress)))
          .build());
  createChannel();
  assertEquals(IDLE, channel.getState(false));

  // This call will be buffered in delayedTransport
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata());

  // Move channel into TRANSIENT_FAILURE, which will fail the pending call
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withError(pickError));
  helper.updateBalancingState(TRANSIENT_FAILURE, mockPicker);
  assertEquals(TRANSIENT_FAILURE, channel.getState(false));
  executor.runDueTasks();
  verify(mockCallListener).onClose(same(pickError), any(Metadata.class));

  // Move channel to idle
  timer.forwardNanos(TimeUnit.MILLISECONDS.toNanos(idleTimeoutMillis));
  assertEquals(IDLE, channel.getState(false));

  // This call should be buffered, but will move the channel out of idle
  ClientCall<String, Integer> call2 = channel.newCall(method, CallOptions.DEFAULT);
  call2.start(mockCallListener2, new Metadata());
  executor.runDueTasks();
  verifyNoMoreInteractions(mockCallListener2);

  // Get the helper created on exiting idle
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(Helper.class);
  verify(mockLoadBalancerFactory, times(2)).newLoadBalancer(helperCaptor.capture());
  Helper helper2 = helperCaptor.getValue();

  // Establish a connection
  Subchannel subchannel = createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY);
  subchannel.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(subchannel));
  helper2.updateBalancingState(READY, mockPicker);
  assertEquals(READY, channel.getState(false));
  executor.runDueTasks();

  // Verify the buffered call was drained
  verify(mockTransport).newStream(same(method), any(Metadata.class), any(CallOptions.class));
  verify(mockStream).start(any(ClientStreamListener.class));
}
 

@Test
public void enterIdle_exitsIdleIfDelayedStreamPending() {
  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(socketAddress)))
          .build();
  channelBuilder.nameResolverFactory(nameResolverFactory);
  createChannel();

  // Start a call that will be buffered in delayedTransport
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata());

  // enterIdle() will shut down the name resolver and lb policy used to get a pick for the delayed
  // call
  channel.enterIdle();
  assertEquals(IDLE, channel.getState(false));

  // enterIdle() will restart the delayed call by exiting idle. This creates a new helper.
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(Helper.class);
  verify(mockLoadBalancerFactory, times(2)).newLoadBalancer(helperCaptor.capture());
  Helper helper2 = helperCaptor.getValue();

  // Establish a connection
  Subchannel subchannel = createSubchannelSafely(helper2, addressGroup, Attributes.EMPTY);
  subchannel.requestConnection();
  ClientStream mockStream = mock(ClientStream.class);
  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(subchannel));
  helper2.updateBalancingState(READY, mockPicker);
  assertEquals(READY, channel.getState(false));

  // Verify the original call was drained
  executor.runDueTasks();
  verify(mockTransport).newStream(same(method), any(Metadata.class), any(CallOptions.class));
  verify(mockStream).start(any(ClientStreamListener.class));
}
 

@Test
public void callOptionsExecutor() {
  Metadata headers = new Metadata();
  ClientStream mockStream = mock(ClientStream.class);
  FakeClock callExecutor = new FakeClock();
  createChannel();

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

  // Make the transport available
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  verify(mockTransportFactory, never())
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  subchannel.requestConnection();
  verify(mockTransportFactory)
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  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));
  assertEquals(0, callExecutor.numPendingTasks());
  helper.updateBalancingState(READY, mockPicker);

  // Real streams are started in the call executor if they were previously buffered.
  assertEquals(1, callExecutor.runDueTasks());
  verify(mockTransport).newStream(same(method), same(headers), same(options));
  verify(mockStream).start(streamListenerCaptor.capture());

  // Call listener callbacks are also run in the call executor
  ClientStreamListener streamListener = streamListenerCaptor.getValue();
  Metadata trailers = new Metadata();
  assertEquals(0, callExecutor.numPendingTasks());
  streamListener.closed(Status.CANCELLED, trailers);
  verify(mockCallListener, never()).onClose(same(Status.CANCELLED), same(trailers));
  assertEquals(1, callExecutor.runDueTasks());
  verify(mockCallListener).onClose(same(Status.CANCELLED), same(trailers));


  transportListener.transportShutdown(Status.UNAVAILABLE);
  transportListener.transportTerminated();

  // Clean up as much as possible to allow the channel to terminate.
  subchannel.shutdown();
  timer.forwardNanos(
      TimeUnit.SECONDS.toNanos(ManagedChannelImpl.SUBCHANNEL_SHUTDOWN_DELAY_SECONDS));
}
 

@Test
public void realTransportsHoldsOffIdleness() throws Exception {
  final EquivalentAddressGroup addressGroup = servers.get(1);

  // Start a call, which goes to delayed transport
  ClientCall<String, Integer> call = channel.newCall(method, CallOptions.DEFAULT);
  call.start(mockCallListener, new Metadata());

  // Verify that we have exited the idle mode
  ArgumentCaptor<Helper> helperCaptor = ArgumentCaptor.forClass(null);
  verify(mockLoadBalancerFactory).newLoadBalancer(helperCaptor.capture());
  Helper helper = helperCaptor.getValue();
  assertTrue(channel.inUseStateAggregator.isInUse());

  // Assume LoadBalancer has received an address, then create a subchannel.
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  subchannel.requestConnection();
  MockClientTransportInfo t0 = newTransports.poll();
  t0.listener.transportReady();

  SubchannelPicker mockPicker = mock(SubchannelPicker.class);
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(subchannel));
  helper.updateBalancingState(READY, mockPicker);
  // Delayed transport creates real streams in the app executor
  executor.runDueTasks();

  // Delayed transport exits in-use, while real transport has not entered in-use yet.
  assertFalse(channel.inUseStateAggregator.isInUse());

  // Now it's in-use
  t0.listener.transportInUse(true);
  assertTrue(channel.inUseStateAggregator.isInUse());

  // As long as the transport is in-use, the channel won't go idle.
  timer.forwardTime(IDLE_TIMEOUT_SECONDS * 2, TimeUnit.SECONDS);
  assertTrue(channel.inUseStateAggregator.isInUse());

  t0.listener.transportInUse(false);
  assertFalse(channel.inUseStateAggregator.isInUse());
  // And allow the channel to go idle.
  timer.forwardTime(IDLE_TIMEOUT_SECONDS - 1, TimeUnit.SECONDS);
  verify(mockLoadBalancer, never()).shutdown();
  timer.forwardTime(1, TimeUnit.SECONDS);
  verify(mockLoadBalancer).shutdown();
}
 

@Test
public void noMoreCallbackAfterLoadBalancerShutdown() {
  FakeNameResolverFactory nameResolverFactory =
      new FakeNameResolverFactory.Builder(expectedUri)
          .setServers(Collections.singletonList(new EquivalentAddressGroup(socketAddress)))
          .build();
  channelBuilder.nameResolverFactory(nameResolverFactory);
  Status resolutionError = Status.UNAVAILABLE.withDescription("Resolution failed");
  createChannel();

  FakeNameResolverFactory.FakeNameResolver resolver = nameResolverFactory.resolvers.get(0);
  verify(mockLoadBalancerFactory).newLoadBalancer(any(Helper.class));
  verify(mockLoadBalancer).handleResolvedAddressGroups(
      eq(Arrays.asList(addressGroup)), eq(Attributes.EMPTY));

  Subchannel subchannel1 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  Subchannel subchannel2 = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  subchannel1.requestConnection();
  subchannel2.requestConnection();
  verify(mockTransportFactory, times(2))
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo1 = transports.poll();
  MockClientTransportInfo transportInfo2 = transports.poll();

  // LoadBalancer receives all sorts of callbacks
  transportInfo1.listener.transportReady();
  verify(mockLoadBalancer, times(2))
      .handleSubchannelState(same(subchannel1), stateInfoCaptor.capture());
  assertSame(CONNECTING, stateInfoCaptor.getAllValues().get(0).getState());
  assertSame(READY, stateInfoCaptor.getAllValues().get(1).getState());

  verify(mockLoadBalancer)
      .handleSubchannelState(same(subchannel2), stateInfoCaptor.capture());
  assertSame(CONNECTING, stateInfoCaptor.getValue().getState());

  resolver.listener.onError(resolutionError);
  verify(mockLoadBalancer).handleNameResolutionError(resolutionError);

  verifyNoMoreInteractions(mockLoadBalancer);

  channel.shutdown();
  verify(mockLoadBalancer).shutdown();

  // No more callback should be delivered to LoadBalancer after it's shut down
  transportInfo2.listener.transportReady();
  resolver.listener.onError(resolutionError);
  resolver.resolved();
  verifyNoMoreInteractions(mockLoadBalancer);
}