io.grpc.LoadBalancer.Subchannel Java Examples

@Deprecated
@Test
public void helper_createSubchannel_old_delegates() {
  class OverrideCreateSubchannel extends NoopHelper {
    boolean ran;

    @Override
    public Subchannel createSubchannel(List<EquivalentAddressGroup> addrsIn, Attributes attrsIn) {
      assertThat(addrsIn).hasSize(1);
      assertThat(addrsIn.get(0)).isSameInstanceAs(eag);
      assertThat(attrsIn).isSameInstanceAs(attrs);
      ran = true;
      return subchannel;
    }
  }

  OverrideCreateSubchannel helper = new OverrideCreateSubchannel();
  assertThat(helper.createSubchannel(eag, attrs)).isSameInstanceAs(subchannel);
  assertThat(helper.ran).isTrue();
}
 

@Override
public void returnSubchannel(Subchannel subchannel) {
  CacheEntry prev = cache.get(subchannel.getAddresses());
  if (prev != null) {
    // Returning the same Subchannel twice has no effect.
    // Returning a different Subchannel for an already cached EAG will cause the
    // latter Subchannel to be shutdown immediately.
    if (prev.subchannel != subchannel) {
      subchannel.shutdown();
    }
    return;
  }
  final ShutdownSubchannelTask shutdownTask = new ShutdownSubchannelTask(subchannel);
  ScheduledHandle shutdownTimer =
      helper.getSynchronizationContext().schedule(
          shutdownTask, SHUTDOWN_TIMEOUT_MS, TimeUnit.MILLISECONDS,
          helper.getScheduledExecutorService());
  CacheEntry entry = new CacheEntry(subchannel, shutdownTimer);
  cache.put(subchannel.getAddresses(), entry);
}
 

@Test
public void handleResolvedAddressGroups_keepOldBalancer() {
  final List<EquivalentAddressGroup> servers =
      Collections.singletonList(new EquivalentAddressGroup(new SocketAddress(){}));
  Helper helper = new TestHelper() {
    @Override
    public Subchannel createSubchannel(CreateSubchannelArgs args) {
      assertThat(args.getAddresses()).isEqualTo(servers);
      return new TestSubchannel(args);
    }
  };
  AutoConfiguredLoadBalancer lb = lbf.newLoadBalancer(helper);
  LoadBalancer oldDelegate = lb.getDelegate();

  Status handleResult = lb.tryHandleResolvedAddresses(
      ResolvedAddresses.newBuilder()
          .setAddresses(servers)
          .setAttributes(Attributes.EMPTY)
          .setLoadBalancingPolicyConfig(null)
          .build());

  assertThat(handleResult.getCode()).isEqualTo(Status.Code.OK);
  assertThat(lb.getDelegate()).isSameInstanceAs(oldDelegate);
}
 

@Test
public void clear() {
  Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
  Subchannel subchannel2 = pool.takeOrCreateSubchannel(EAG2, ATTRS2);
  Subchannel subchannel3 = pool.takeOrCreateSubchannel(EAG2, ATTRS2);

  pool.returnSubchannel(subchannel1);
  pool.returnSubchannel(subchannel2);

  verify(subchannel1, never()).shutdown();
  verify(subchannel2, never()).shutdown();
  pool.clear();
  verify(subchannel1).shutdown();
  verify(subchannel2).shutdown();

  verify(subchannel3, never()).shutdown();
  assertThat(clock.numPendingTasks()).isEqualTo(0);
}
 

@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 roundRobinPickerWithIdleEntry_noDrop() {
  Subchannel subchannel = mock(Subchannel.class);
  IdleSubchannelEntry entry = new IdleSubchannelEntry(subchannel, syncContext);

  RoundRobinPicker picker =
      new RoundRobinPicker(Collections.<DropEntry>emptyList(), Collections.singletonList(entry));
  PickSubchannelArgs args = mock(PickSubchannelArgs.class);

  verify(subchannel, never()).requestConnection();
  assertThat(picker.pickSubchannel(args)).isSameInstanceAs(PickResult.withNoResult());
  verify(subchannel).requestConnection();
  assertThat(picker.pickSubchannel(args)).isSameInstanceAs(PickResult.withNoResult());
  // Only the first pick triggers requestConnection()
  verify(subchannel).requestConnection();
}
 

@Override
public Subchannel createSubchannel(List<EquivalentAddressGroup> addrs, Attributes attrs) {
  // HealthCheckState is not thread-safe, we are requiring the original LoadBalancer calls
  // createSubchannel() from the SynchronizationContext.
  syncContext.throwIfNotInThisSynchronizationContext();
  HealthCheckState hcState = new HealthCheckState(
      this, delegateBalancer, syncContext, delegate.getScheduledExecutorService());
  hcStates.add(hcState);
  Subchannel subchannel = super.createSubchannel(
      addrs, attrs.toBuilder().set(KEY_HEALTH_CHECK_STATE, hcState).build());
  hcState.init(subchannel);
  if (healthCheckedService != null) {
    hcState.setServiceName(healthCheckedService);
  }
  return subchannel;
}
 

@Before
public void setUp() {
  doAnswer(new Answer<Subchannel>() {
      @Override
      public Subchannel answer(InvocationOnMock invocation) throws Throwable {
        Subchannel subchannel = mock(Subchannel.class);
        CreateSubchannelArgs args = (CreateSubchannelArgs) invocation.getArguments()[0];
        when(subchannel.getAllAddresses()).thenReturn(args.getAddresses());
        when(subchannel.getAttributes()).thenReturn(args.getAttributes());
        mockSubchannels.add(subchannel);
        return subchannel;
      }
    }).when(helper).createSubchannel(any(CreateSubchannelArgs.class));
  when(helper.getSynchronizationContext()).thenReturn(syncContext);
  when(helper.getScheduledExecutorService()).thenReturn(clock.getScheduledExecutorService());
  pool.registerListener(listener);
}
 

@Test
public void subchannelChannel_failWhenNotReady() {
  createChannel();
  Subchannel subchannel = createSubchannelSafely(helper, addressGroup, Attributes.EMPTY);
  Channel sChannel = subchannel.asChannel();
  Metadata headers = new Metadata();

  subchannel.requestConnection();
  verify(mockTransportFactory)
      .newClientTransport(any(SocketAddress.class), any(ClientTransportOptions.class));
  MockClientTransportInfo transportInfo = transports.poll();
  ConnectionClientTransport mockTransport = transportInfo.transport;

  assertEquals(0, balancerRpcExecutor.numPendingTasks());

  // Subchannel is still CONNECTING, but not READY yet
  ClientCall<String, Integer> call = sChannel.newCall(method, CallOptions.DEFAULT);
  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.NOT_READY_ERROR), any(Metadata.class));
}
 

@Test
public void helper_updateSubchannelAddresses_delegates() {
  class OverrideUpdateSubchannel extends NoopHelper {
    boolean ran;

    @Override
    public void updateSubchannelAddresses(
        Subchannel subchannelIn, List<EquivalentAddressGroup> addrsIn) {
      assertThat(subchannelIn).isSameAs(emptySubchannel);
      assertThat(addrsIn).hasSize(1);
      assertThat(addrsIn.get(0)).isSameAs(eag);
      ran = true;
    }
  }

  OverrideUpdateSubchannel helper = new OverrideUpdateSubchannel();
  helper.updateSubchannelAddresses(emptySubchannel, eag);
  assertThat(helper.ran).isTrue();
}
 

private static Subchannel createUnstartedSubchannel(
    final Helper helper, final EquivalentAddressGroup addressGroup, final Attributes attrs) {
  final AtomicReference<Subchannel> resultCapture = new AtomicReference<>();
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          Subchannel s = helper.createSubchannel(CreateSubchannelArgs.newBuilder()
              .setAddresses(addressGroup)
              .setAttributes(attrs)
              .build());
          resultCapture.set(s);
        }
      });
  return resultCapture.get();
}
 

@Test
public void subchannelsNoConnectionShutdown() {
  createChannel();
  Subchannel sub1 =
      createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
  Subchannel sub2 =
      createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);

  channel.shutdown();
  verify(mockLoadBalancer).shutdown();
  shutdownSafely(helper, sub1);
  assertFalse(channel.isTerminated());
  shutdownSafely(helper, sub2);
  assertTrue(channel.isTerminated());
  verify(mockTransportFactory, never())
      .newClientTransport(
          any(SocketAddress.class), any(ClientTransportOptions.class), any(ChannelLogger.class));
}
 

@Test
public void subchannelStateChange_updateChildPolicyWrapper() {
  ChildPolicyWrapper childPolicyWrapper = factory.createOrGet("foo.google.com");
  ChildPolicyReportingHelper childPolicyReportingHelper = childPolicyWrapper.getHelper();
  FakeSubchannel fakeSubchannel = new FakeSubchannel();
  when(helper.createSubchannel(any(CreateSubchannelArgs.class))).thenReturn(fakeSubchannel);
  Subchannel subchannel =
      childPolicyReportingHelper
          .createSubchannel(
              CreateSubchannelArgs.newBuilder()
                  .setAddresses(new EquivalentAddressGroup(mock(SocketAddress.class)))
                  .build());
  subchannel.start(new SubchannelStateListener() {
    @Override
    public void onSubchannelState(ConnectivityStateInfo newState) {
      // no-op
    }
  });

  fakeSubchannel.updateState(ConnectivityStateInfo.forNonError(ConnectivityState.CONNECTING));

  assertThat(childPolicyWrapper.getConnectivityStateInfo())
      .isEqualTo(ConnectivityStateInfo.forNonError(ConnectivityState.CONNECTING));
}
 

@Override
public Subchannel createSubchannel(CreateSubchannelArgs args) {
  syncContext.throwIfNotInThisSynchronizationContext();
  OrcaReportingState orcaState = args.getOption(ORCA_REPORTING_STATE_KEY);
  boolean augmented = false;
  if (orcaState == null) {
    // Only the first load balancing policy requesting ORCA reports instantiates an
    // OrcaReportingState.
    orcaState = new OrcaReportingState(this, syncContext,
        delegate().getScheduledExecutorService());
    args = args.toBuilder().addOption(ORCA_REPORTING_STATE_KEY, orcaState).build();
    augmented = true;
  }
  orcaStates.add(orcaState);
  orcaState.listeners.add(this);
  Subchannel subchannel = super.createSubchannel(args);
  if (augmented) {
    subchannel = new SubchannelImpl(subchannel, orcaState);
  }
  if (orcaConfig != null) {
    orcaState.setReportingConfig(this, orcaConfig);
  }
  return subchannel;
}
 

/**
 * Returns the subchannel associated to the stickiness value if available in both the
 * registry and the round robin list, otherwise associates the given subchannel with the
 * stickiness key in the registry and returns the given subchannel.
 */
@Nonnull
Subchannel maybeRegister(
    String stickinessValue, @Nonnull Subchannel subchannel) {
  final Ref<Subchannel> newSubchannelRef = subchannel.getAttributes().get(STICKY_REF);
  while (true) {
    Ref<Subchannel> existingSubchannelRef =
        stickinessMap.putIfAbsent(stickinessValue, newSubchannelRef);
    if (existingSubchannelRef == null) {
      // new entry
      addToEvictionQueue(stickinessValue);
      return subchannel;
    } else {
      // existing entry
      Subchannel existingSubchannel = existingSubchannelRef.value;
      if (existingSubchannel != null && isReady(existingSubchannel)) {
        return existingSubchannel;
      }
    }
    // existingSubchannelRef is not null but no longer valid, replace it
    if (stickinessMap.replace(stickinessValue, existingSubchannelRef, newSubchannelRef)) {
      return subchannel;
    }
    // another thread concurrently removed or updated the entry, try again
  }
}
 

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

  ReadyPicker picker = new ReadyPicker(Collections.unmodifiableList(
      Lists.<Subchannel>newArrayList(subchannel, subchannel1, subchannel2)),
      0 /* startIndex */, null /* stickinessState */);

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

@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);
}
 

/**
 * Updates picker with the list of active subchannels (state == READY).
 */
@SuppressWarnings("ReferenceEquality")
private void updateBalancingState() {
  List<Subchannel> activeList = filterNonFailingSubchannels(getSubchannels());
  if (activeList.isEmpty()) {
    // No READY subchannels, determine aggregate state and error status
    boolean isConnecting = false;
    Status aggStatus = EMPTY_OK;
    for (Subchannel subchannel : getSubchannels()) {
      ConnectivityStateInfo stateInfo = getSubchannelStateInfoRef(subchannel).value;
      // 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.
      if (stateInfo.getState() == CONNECTING || stateInfo.getState() == IDLE) {
        isConnecting = true;
      }
      if (aggStatus == EMPTY_OK || !aggStatus.isOk()) {
        aggStatus = stateInfo.getStatus();
      }
    }
    updateBalancingState(isConnecting ? CONNECTING : TRANSIENT_FAILURE,
        // If all subchannels are TRANSIENT_FAILURE, return the Status associated with
        // an arbitrary subchannel, otherwise return OK.
        new EmptyPicker(aggStatus));
  } else {
    // initialize the Picker to a random start index to ensure that a high frequency of Picker
    // churn does not skew subchannel selection.
    int startIndex = random.nextInt(activeList.size());
    updateBalancingState(READY, new ReadyPicker(activeList, startIndex, stickinessState));
  }
}
 

@Before
@SuppressWarnings("unchecked")
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.getAllAddresses()).thenReturn(Arrays.asList(eag));
    subchannels.put(Arrays.asList(eag), sc);
  }

  when(mockHelper.createSubchannel(any(List.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 = new RoundRobinLoadBalancer(mockHelper);
}
 

ReadyPicker(List<Subchannel> list, int startIndex,
    @Nullable RoundRobinLoadBalancer.StickinessState stickinessState) {
  Preconditions.checkArgument(!list.isEmpty(), "empty list");
  this.list = list;
  this.stickinessState = stickinessState;
  this.index = startIndex - 1;
}
 

@Test
public void returnDuplicateAddressSubchannel() {
  Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
  Subchannel subchannel2 = pool.takeOrCreateSubchannel(EAG1, ATTRS2);
  Subchannel subchannel3 = pool.takeOrCreateSubchannel(EAG2, ATTRS1);
  assertThat(subchannel1).isNotSameInstanceAs(subchannel2);

  assertThat(clock.getPendingTasks(SHUTDOWN_TASK_FILTER)).isEmpty();
  pool.returnSubchannel(subchannel2, READY_STATE);
  assertThat(clock.getPendingTasks(SHUTDOWN_TASK_FILTER)).hasSize(1);

  // If the subchannel being returned has an address that is the same as a subchannel in the pool,
  // the returned subchannel will be shut down.
  verify(subchannel1, never()).shutdown();
  pool.returnSubchannel(subchannel1, READY_STATE);
  assertThat(clock.getPendingTasks(SHUTDOWN_TASK_FILTER)).hasSize(1);
  verify(subchannel1).shutdown();

  pool.returnSubchannel(subchannel3, READY_STATE);
  assertThat(clock.getPendingTasks(SHUTDOWN_TASK_FILTER)).hasSize(2);
  // Returning the same subchannel twice has no effect.
  pool.returnSubchannel(subchannel3, READY_STATE);
  assertThat(clock.getPendingTasks(SHUTDOWN_TASK_FILTER)).hasSize(2);

  verify(subchannel2, never()).shutdown();
  verify(subchannel3, never()).shutdown();
}
 

@After
public void wrapUp() {
  if (mockSubchannels.isEmpty()) {
    return;
  }
  // Sanity checks
  for (Subchannel subchannel : mockSubchannels) {
    verify(subchannel, atMost(1)).shutdown();
  }
  verify(listener, atLeast(0))
      .onSubchannelState(any(Subchannel.class), any(ConnectivityStateInfo.class));
  verifyNoMoreInteractions(listener);
}
 

void deliverSubchannelState(final Subchannel subchannel, ConnectivityState state) {
  SubchannelPicker picker = new SubchannelPicker() {
    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      return PickResult.withSubchannel(subchannel);
    }
  };
  helper.updateBalancingState(state, picker);
}
 

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

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

  final Subchannel subchannel1 =
      createSubchannelSafely(helper, addressGroup, Attributes.EMPTY, subchannelStateListener);
  requestConnectionSafely(helper, subchannel1);

  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();

  Subchannel wrappedSubchannel1 = new ForwardingSubchannel() {
      @Override
      protected Subchannel delegate() {
        return subchannel1;
      }
    };
  when(mockPicker.pickSubchannel(any(PickSubchannelArgs.class)))
      .thenReturn(PickResult.withSubchannel(wrappedSubchannel1));
  updateBalancingStateSafely(helper, READY, mockPicker);

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

/**
 * 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<>(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) {
      pickList.add(new ErrorEntry(error));
      state = TRANSIENT_FAILURE;
    } else {
      pickList.add(BUFFER_ENTRY);
      state = CONNECTING;
    }
  } else {
    state = READY;
  }
  maybeUpdatePicker(state, new RoundRobinPicker(dropList, pickList));
}
 

/**
 * For ROUND_ROBIN: creates a BackendEntry whose usage will be reported to load recorder.
 */
BackendEntry(Subchannel subchannel, GrpclbClientLoadRecorder loadRecorder, String token) {
  this.subchannel = checkNotNull(subchannel, "subchannel");
  this.result =
      PickResult.withSubchannel(subchannel, checkNotNull(loadRecorder, "loadRecorder"));
  this.token = checkNotNull(token, "token");
}
 

@Test
public void subchannelExpireAfterReturned() {
  Subchannel subchannel1 = pool.takeOrCreateSubchannel(EAG1, ATTRS1);
  assertThat(subchannel1).isNotNull();
  verify(helper).createSubchannel(eq(Arrays.asList(EAG1)), same(ATTRS1));

  Subchannel subchannel2 = pool.takeOrCreateSubchannel(EAG2, ATTRS2);
  assertThat(subchannel2).isNotNull();
  assertThat(subchannel2).isNotSameAs(subchannel1);
  verify(helper).createSubchannel(eq(Arrays.asList(EAG2)), same(ATTRS2));

  pool.returnSubchannel(subchannel1);

  // subchannel1 is 1ms away from expiration.
  clock.forwardTime(SHUTDOWN_TIMEOUT_MS - 1, MILLISECONDS);
  verify(subchannel1, never()).shutdown();

  pool.returnSubchannel(subchannel2);

  // subchannel1 expires. subchannel2 is (SHUTDOWN_TIMEOUT_MS - 1) away from expiration.
  clock.forwardTime(1, MILLISECONDS);
  verify(subchannel1).shutdown();

  // subchanne2 expires.
  clock.forwardTime(SHUTDOWN_TIMEOUT_MS - 1, MILLISECONDS);
  verify(subchannel2).shutdown();

  assertThat(clock.numPendingTasks()).isEqualTo(0);
}
 

@Test
public void roundRobinPickerNoDrop() {
  GrpclbClientLoadRecorder loadRecorder =
      new GrpclbClientLoadRecorder(fakeClock.getTimeProvider());
  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();
}
 

@Test
public void balancerShutdown() {
  Attributes resolutionAttrs = attrsWithHealthCheckService("TeeService");
  hcLbEventDelivery.handleResolvedAddressGroups(resolvedAddressList, resolutionAttrs);

  verify(origLb).handleResolvedAddressGroups(same(resolvedAddressList), same(resolutionAttrs));
  verifyNoMoreInteractions(origLb);

  Subchannel subchannel = createSubchannel(0, Attributes.EMPTY);
  assertThat(subchannel).isSameAs(subchannels[0]);

  // Trigger the health check
  hcLbEventDelivery.handleSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));

  HealthImpl healthImpl = healthImpls[0];
  assertThat(healthImpl.calls).hasSize(1);
  ServerSideCall serverCall = healthImpl.calls.poll();
  assertThat(serverCall.cancelled).isFalse();

  verify(origLb).handleSubchannelState(
      same(subchannel), eq(ConnectivityStateInfo.forNonError(CONNECTING)));

  // Shut down the balancer
  hcLbEventDelivery.shutdown();
  verify(origLb).shutdown();

  // Health check stream should be cancelled
  assertThat(serverCall.cancelled).isTrue();

  // LoadBalancer API requires no more callbacks on LoadBalancer after shutdown() is called.
  verifyNoMoreInteractions(origLb);

  // No more health check call is made or scheduled
  assertThat(healthImpl.calls).isEmpty();
  assertThat(clock.getPendingTasks()).isEmpty();
}
 

@Override
public Subchannel takeOrCreateSubchannel(
    EquivalentAddressGroup eag, Attributes defaultAttributes) {
  CacheEntry entry = cache.remove(eag);
  Subchannel subchannel;
  if (entry == null) {
    subchannel = helper.createSubchannel(eag, defaultAttributes);
  } else {
    subchannel = entry.subchannel;
    entry.shutdownTimer.cancel();
  }
  return subchannel;
}