#20125

Committed 23 Dec 2025 04:32AM UTC coverage: 88.706% (-0.01%) from 88.72%

Build # #20125

Build Type

push

github

Committed by

web-flow

Commit Message

grpclb: pick_first delegation (#12568)

**Summary of Changes**
This pull request refactors the grpclb load balancer's PICK_FIRST mode
to delegate its logic to a standard pick_first load balancing policy.

The key changes are as follows:
1. **`grpclb/build.gradle`**

Added dependency on `grpc-util` module to access
`ForwardingLoadBalancerHelper`

2. **`grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java`**
- New imports:
LoadBalancer, LoadBalancerProvider, LoadBalancerRegistry,
ResolvedAddresses, FixedResultPicker, ForwardingLoadBalancerHelper
- New fields for PICK_FIRST delegation:
    - pickFirstLbProvider - Provider for creating child pick_first LB
    - pickFirstLb - The child LoadBalancer instance
- pickFirstLbState / pickFirstLbPicker - Track child LB's state and
picker
    - currentPickFirstLoadRecorder - Load recorder for token attachment
- Key behavioral changes:
- updateServerList() PICK_FIRST case: Instead of creating a single
subchannel, it now:
- Creates the child pick_first LB once and then updates it with new
addresses on subsequent updates.
        - Passes addresses to child LB via acceptResolvedAddresses()
- maybeUpdatePicker() PICK_FIRST case: Uses child LB's state and picker
wrapped with ChildLbPickerEntry
- RoundRobinEntry.picked() signature change: Changed from
picked(Metadata) to picked(PickSubchannelArgs) to allow child picker
delegation
- New ChildLbPickerEntry class: Wraps child LB's picker and attaches
TokenAttachingTracerFactory for token propagation
- New PickFirstLbHelper class: Forwarding helper that intercepts
updateBalancingState() to store child state and trigger grpclb picker
updates
- Updated shutdown(), requestConnection(), maybeUseFallbackBackends():
Handle the new child LB delegation model

3. **`grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java`**

- Updated tests to reflect the new delegation behavior:
- Initial state is now CONNECTING (not IDLE) since standard pick_first
eagerly connects
- Tests ... (continued)

Run Details

35463 of 39978 relevant lines covered (88.71%)

0.89 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.79

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

/*
 * Copyright 2017 The gRPC Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.grpc.grpclb;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static io.grpc.ConnectivityState.CONNECTING;
import static io.grpc.ConnectivityState.IDLE;
import static io.grpc.ConnectivityState.READY;
import static io.grpc.ConnectivityState.SHUTDOWN;
import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import com.google.common.base.Objects;
import com.google.common.base.Stopwatch;
import com.google.protobuf.util.Durations;
import io.grpc.Attributes;
import io.grpc.ChannelLogger;
import io.grpc.ChannelLogger.ChannelLogLevel;
import io.grpc.ConnectivityState;
import io.grpc.ConnectivityStateInfo;
import io.grpc.Context;
import io.grpc.EquivalentAddressGroup;
import io.grpc.LoadBalancer;
import io.grpc.LoadBalancer.FixedResultPicker;
import io.grpc.LoadBalancer.Helper;
import io.grpc.LoadBalancer.PickResult;
import io.grpc.LoadBalancer.PickSubchannelArgs;
import io.grpc.LoadBalancer.ResolvedAddresses;
import io.grpc.LoadBalancer.Subchannel;
import io.grpc.LoadBalancer.SubchannelPicker;
import io.grpc.LoadBalancerProvider;
import io.grpc.LoadBalancerRegistry;
import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.Status;
import io.grpc.SynchronizationContext;
import io.grpc.SynchronizationContext.ScheduledHandle;
import io.grpc.grpclb.SubchannelPool.PooledSubchannelStateListener;
import io.grpc.internal.BackoffPolicy;
import io.grpc.internal.TimeProvider;
import io.grpc.lb.v1.ClientStats;
import io.grpc.lb.v1.InitialLoadBalanceRequest;
import io.grpc.lb.v1.InitialLoadBalanceResponse;
import io.grpc.lb.v1.LoadBalanceRequest;
import io.grpc.lb.v1.LoadBalanceResponse;
import io.grpc.lb.v1.LoadBalanceResponse.LoadBalanceResponseTypeCase;
import io.grpc.lb.v1.LoadBalancerGrpc;
import io.grpc.lb.v1.Server;
import io.grpc.lb.v1.ServerList;
import io.grpc.stub.StreamObserver;
import io.grpc.util.ForwardingLoadBalancerHelper;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * The states of a GRPCLB working session of {@link GrpclbLoadBalancer}.  Created when
 * GrpclbLoadBalancer switches to GRPCLB mode.  Closed and discarded when GrpclbLoadBalancer
 * switches away from GRPCLB mode.
 */
@NotThreadSafe
final class GrpclbState {
  static final long FALLBACK_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(10);
  private static final Attributes LB_PROVIDED_BACKEND_ATTRS =
      Attributes.newBuilder().set(GrpclbConstants.ATTR_LB_PROVIDED_BACKEND, true).build();

  // Temporary workaround to reduce log spam for a grpclb server that incessantly sends updates
  // Tracked by b/198440401
  static final boolean SHOULD_LOG_SERVER_LISTS =
      Boolean.parseBoolean(System.getProperty("io.grpc.grpclb.LogServerLists", "true"));

  @VisibleForTesting
  static final PickResult DROP_PICK_RESULT =
      PickResult.withDrop(Status.UNAVAILABLE.withDescription("Dropped as requested by balancer"));
  @VisibleForTesting
  static final Status NO_AVAILABLE_BACKENDS_STATUS =
      Status.UNAVAILABLE.withDescription("LoadBalancer responded without any backends");
  @VisibleForTesting
  static final Status BALANCER_TIMEOUT_STATUS =
      Status.UNAVAILABLE.withDescription("Timeout waiting for remote balancer");
  @VisibleForTesting
  static final Status BALANCER_REQUESTED_FALLBACK_STATUS =
      Status.UNAVAILABLE.withDescription("Fallback requested by balancer");
  @VisibleForTesting
  static final Status NO_FALLBACK_BACKENDS_STATUS =
      Status.UNAVAILABLE.withDescription("Unable to fallback, no fallback addresses found");
  // This error status should never be propagated to RPC failures, as "no backend or balancer
  // addresses found" should be directly handled as a name resolution error. So in cases of no
  // balancer address, fallback should never fail.
  private static final Status NO_LB_ADDRESS_PROVIDED_STATUS =
      Status.UNAVAILABLE.withDescription("No balancer address found");


  @VisibleForTesting
  static final RoundRobinEntry BUFFER_ENTRY = new RoundRobinEntry() {
      @Override
      public PickResult picked(PickSubchannelArgs args) {
        return PickResult.withNoResult();
      }

      @Override
      public String toString() {
        return "BUFFER_ENTRY";
      }
    };
  @VisibleForTesting
  static final String NO_USE_AUTHORITY_SUFFIX = "-notIntendedToBeUsed";

  enum Mode {
    ROUND_ROBIN,
    PICK_FIRST,
  }

  private final String serviceName;
  private final long fallbackTimeoutMs;
  private final Helper helper;
  private final Context context;
  private final SynchronizationContext syncContext;
  @Nullable
  private final SubchannelPool subchannelPool;
  private final TimeProvider time;
  private final Stopwatch stopwatch;
  private final ScheduledExecutorService timerService;

  private static final Attributes.Key<AtomicReference<ConnectivityStateInfo>> STATE_INFO =
      Attributes.Key.create("io.grpc.grpclb.GrpclbLoadBalancer.stateInfo");
  private final BackoffPolicy.Provider backoffPolicyProvider;
  private final ChannelLogger logger;

  // Scheduled only once.  Never reset.
  @Nullable
  private ScheduledHandle fallbackTimer;
  private List<EquivalentAddressGroup> fallbackBackendList = Collections.emptyList();
  private boolean usingFallbackBackends;
  // Reason to fallback, will be used as RPC's error message if fail to fallback (e.g., no
  // fallback addresses found).
  @Nullable
  private Status fallbackReason;
  // True if the current balancer has returned a serverlist.  Will be reset to false when lost
  // connection to a balancer.
  private boolean balancerWorking;
  @Nullable
  private BackoffPolicy lbRpcRetryPolicy;
  @Nullable
  private ScheduledHandle lbRpcRetryTimer;

  @Nullable
  private ManagedChannel lbCommChannel;

  @Nullable
  private LbStream lbStream;
  private Map<List<EquivalentAddressGroup>, Subchannel> subchannels = Collections.emptyMap();
  private final GrpclbConfig config;

  // Has the same size as the round-robin list from the balancer.
  // A drop entry from the round-robin list becomes a DropEntry here.
  // A backend entry from the robin-robin list becomes a null here.
  private List<DropEntry> dropList = Collections.emptyList();
  // Contains only non-drop, i.e., backends from the round-robin list from the balancer.
  private List<BackendEntry> backendList = Collections.emptyList();
  private RoundRobinPicker currentPicker =
      new RoundRobinPicker(Collections.<DropEntry>emptyList(), Arrays.asList(BUFFER_ENTRY));
  private boolean requestConnectionPending;

  // Child LoadBalancer and state for PICK_FIRST mode delegation.
  private final LoadBalancerProvider pickFirstLbProvider;
  @Nullable
  private LoadBalancer pickFirstLb;
  private ConnectivityState pickFirstLbState = CONNECTING;
  private SubchannelPicker pickFirstLbPicker = new FixedResultPicker(PickResult.withNoResult());
  @Nullable
  private GrpclbClientLoadRecorder currentPickFirstLoadRecorder;

  GrpclbState(
      GrpclbConfig config,
      Helper helper,
      Context context,
      SubchannelPool subchannelPool,
      TimeProvider time,
      Stopwatch stopwatch,
      BackoffPolicy.Provider backoffPolicyProvider) {
    this.config = checkNotNull(config, "config");
    this.helper = checkNotNull(helper, "helper");
    this.context = checkNotNull(context, "context");
    this.syncContext = checkNotNull(helper.getSynchronizationContext(), "syncContext");
    if (config.getMode() == Mode.ROUND_ROBIN) {
      this.subchannelPool = checkNotNull(subchannelPool, "subchannelPool");
      subchannelPool.registerListener(
          new PooledSubchannelStateListener() {
            @Override
            public void onSubchannelState(
                Subchannel subchannel, ConnectivityStateInfo newState) {
              handleSubchannelState(subchannel, newState);
            }
          });
    } else {
      this.subchannelPool = null;
    }
    this.pickFirstLbProvider = checkNotNull(
        LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first"),
        "pick_first balancer not available");
    this.time = checkNotNull(time, "time provider");
    this.stopwatch = checkNotNull(stopwatch, "stopwatch");
    this.timerService = checkNotNull(helper.getScheduledExecutorService(), "timerService");
    this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
    if (config.getServiceName() != null) {
      this.serviceName = config.getServiceName();
    } else {
      this.serviceName = checkNotNull(helper.getAuthority(), "helper returns null authority");
    }
    this.fallbackTimeoutMs = config.getFallbackTimeoutMs();
    this.logger = checkNotNull(helper.getChannelLogger(), "logger");
    logger.log(ChannelLogLevel.INFO, "[grpclb-<{0}>] Created", serviceName);
  }

  void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo newState) {
    if (newState.getState() == SHUTDOWN || !subchannels.containsValue(subchannel)) {
      return;
    }
    if (config.getMode() == Mode.ROUND_ROBIN && newState.getState() == IDLE) {
      subchannel.requestConnection();
    }
    if (newState.getState() == TRANSIENT_FAILURE || newState.getState() == IDLE) {
      helper.refreshNameResolution();
    }

    AtomicReference<ConnectivityStateInfo> stateInfoRef =
        subchannel.getAttributes().get(STATE_INFO);
    // If all RR servers are unhealthy, it's possible that at least one connection is CONNECTING at
    // every moment which causes RR to stay in CONNECTING. It's better to keep the TRANSIENT_FAILURE
    // state in that case so that fail-fast RPCs can fail fast.
    boolean keepState =
        config.getMode() == Mode.ROUND_ROBIN
        && stateInfoRef.get().getState() == TRANSIENT_FAILURE
        && (newState.getState() == CONNECTING || newState.getState() == IDLE);
    if (!keepState) {
      stateInfoRef.set(newState);
      maybeUseFallbackBackends();
      maybeUpdatePicker();
    }
  }

  /**
   * Handle new addresses of the balancer and backends from the resolver, and create connection if
   * not yet connected.
   */
  void handleAddresses(
      List<EquivalentAddressGroup> newLbAddressGroups,
      List<EquivalentAddressGroup> newBackendServers) {
    logger.log(
        ChannelLogLevel.DEBUG,
        "[grpclb-<{0}>] Resolved addresses: lb addresses {1}, backends: {2}",
        serviceName,
        newLbAddressGroups,
        newBackendServers);
    fallbackBackendList = newBackendServers;
    if (newLbAddressGroups.isEmpty()) {
      // No balancer address: close existing balancer connection and prepare to enter fallback
      // mode. If there is no successful backend connection, it enters fallback mode immediately.
      // Otherwise, fallback does not happen until backend connections are lost. This behavior
      // might be different from other languages (e.g., existing balancer connection is not
      // closed in C-core), but we aren't changing it at this time.
      shutdownLbComm();
      if (!usingFallbackBackends) {
        fallbackReason = NO_LB_ADDRESS_PROVIDED_STATUS;
        cancelFallbackTimer();
        maybeUseFallbackBackends();
      }
    } else {
      startLbComm(newLbAddressGroups);
      // Avoid creating a new RPC just because the addresses were updated, as it can cause a
      // stampeding herd. The current RPC may be on a connection to an address not present in
      // newLbAddressGroups, but we're considering that "okay". If we detected the RPC is to an
      // outdated backend, we could choose to re-create the RPC.
      if (lbStream == null) {
        cancelLbRpcRetryTimer();
        startLbRpc();
      }
      // Start the fallback timer if it's never started and we are not already using fallback
      // backends.
      if (fallbackTimer == null && !usingFallbackBackends) {
        fallbackTimer =
            syncContext.schedule(
                new FallbackModeTask(BALANCER_TIMEOUT_STATUS),
                fallbackTimeoutMs,
                TimeUnit.MILLISECONDS,
                timerService);
      }
    }
    if (usingFallbackBackends) {
      // Populate the new fallback backends to round-robin list.
      useFallbackBackends();
    }
    maybeUpdatePicker();
  }

  void requestConnection() {
    requestConnectionPending = true;
    // For PICK_FIRST mode with delegation, forward to the child LB.
    if (config.getMode() == Mode.PICK_FIRST && pickFirstLb != null) {
      pickFirstLb.requestConnection();
      requestConnectionPending = false;
      return;
    }
    for (RoundRobinEntry entry : currentPicker.pickList) {
      if (entry instanceof IdleSubchannelEntry) {
        ((IdleSubchannelEntry) entry).subchannel.requestConnection();
        requestConnectionPending = false;
      }
    }
  }

  private void maybeUseFallbackBackends() {
    if (balancerWorking || usingFallbackBackends) {
      return;
    }
    // Balancer RPC should have either been broken or timed out.
    checkState(fallbackReason != null, "no reason to fallback");
    // For PICK_FIRST mode with delegation, check the child LB's state.
    if (config.getMode() == Mode.PICK_FIRST) {
      if (pickFirstLb != null && pickFirstLbState == READY) {
        return;
      }
      // For PICK_FIRST, we don't have individual subchannel states to use as fallback reason.
    } else {
      for (Subchannel subchannel : subchannels.values()) {
        ConnectivityStateInfo stateInfo = subchannel.getAttributes().get(STATE_INFO).get();
        if (stateInfo.getState() == READY) {
          return;
        }
        // If we do have balancer-provided backends, use one of its error in the error message if
        // fail to fallback.
        if (stateInfo.getState() == TRANSIENT_FAILURE) {
          fallbackReason = stateInfo.getStatus();
        }
      }
    }
    // Fallback conditions met
    useFallbackBackends();
  }

  /**
   * Populate backend servers to be used from the fallback backends.
   */
  private void useFallbackBackends() {
    usingFallbackBackends = true;
    logger.log(ChannelLogLevel.INFO, "[grpclb-<{0}>] Using fallback backends", serviceName);

    List<DropEntry> newDropList = new ArrayList<>();
    List<BackendAddressGroup> newBackendAddrList = new ArrayList<>();
    for (EquivalentAddressGroup eag : fallbackBackendList) {
      newDropList.add(null);
      newBackendAddrList.add(new BackendAddressGroup(eag, null));
    }
    updateServerList(newDropList, newBackendAddrList, null);
  }

  private void shutdownLbComm() {
    if (lbCommChannel != null) {
      lbCommChannel.shutdown();
      lbCommChannel = null;
    }
    shutdownLbRpc();
  }

  private void shutdownLbRpc() {
    if (lbStream != null) {
      lbStream.close(Status.CANCELLED.withDescription("balancer shutdown").asException());
      // lbStream will be set to null in LbStream.cleanup()
    }
  }

  private void startLbComm(List<EquivalentAddressGroup> overrideAuthorityEags) {
    checkNotNull(overrideAuthorityEags, "overrideAuthorityEags");
    assert !overrideAuthorityEags.isEmpty();
    String doNotUseAuthority = overrideAuthorityEags.get(0).getAttributes()
        .get(EquivalentAddressGroup.ATTR_AUTHORITY_OVERRIDE) + NO_USE_AUTHORITY_SUFFIX;
    if (lbCommChannel == null) {
      lbCommChannel = helper.createOobChannel(overrideAuthorityEags, doNotUseAuthority);
      logger.log(
          ChannelLogLevel.DEBUG,
          "[grpclb-<{0}>] Created grpclb channel: EAG={1}",
          serviceName,
          overrideAuthorityEags);
    } else {
      helper.updateOobChannelAddresses(lbCommChannel, overrideAuthorityEags);
    }
  }

  private void startLbRpc() {
    checkState(lbStream == null, "previous lbStream has not been cleared yet");
    LoadBalancerGrpc.LoadBalancerStub stub = LoadBalancerGrpc.newStub(lbCommChannel);
    lbStream = new LbStream(stub);
    Context prevContext = context.attach();
    try {
      lbStream.start();
    } finally {
      context.detach(prevContext);
    }
    stopwatch.reset().start();

    LoadBalanceRequest initRequest = LoadBalanceRequest.newBuilder()
        .setInitialRequest(InitialLoadBalanceRequest.newBuilder()
            .setName(serviceName).build())
        .build();
    logger.log(
        ChannelLogLevel.DEBUG,
        "[grpclb-<{0}>] Sent initial grpclb request {1}", serviceName, initRequest);
    try {
      lbStream.lbRequestWriter.onNext(initRequest);
    } catch (Exception e) {
      lbStream.close(e);
    }
  }

  private void cancelFallbackTimer() {
    if (fallbackTimer != null) {
      fallbackTimer.cancel();
    }
  }

  private void cancelLbRpcRetryTimer() {
    if (lbRpcRetryTimer != null) {
      lbRpcRetryTimer.cancel();
      lbRpcRetryTimer = null;
    }
  }

  void shutdown() {
    logger.log(ChannelLogLevel.INFO, "[grpclb-<{0}>] Shutdown", serviceName);
    shutdownLbComm();
    switch (config.getMode()) {
      case ROUND_ROBIN:
        // We close the subchannels through subchannelPool instead of helper just for convenience of
        // testing.
        for (Subchannel subchannel : subchannels.values()) {
          returnSubchannelToPool(subchannel);
        }
        subchannelPool.clear();
        break;
      case PICK_FIRST:
        // Shutdown the child pick_first LB which manages its own subchannels.
        if (pickFirstLb != null) {
          pickFirstLb.shutdown();
          pickFirstLb = null;
        }
        break;
      default:
        throw new AssertionError("Missing case for " + config.getMode());
    }
    subchannels = Collections.emptyMap();
    cancelFallbackTimer();
    cancelLbRpcRetryTimer();
  }

  void propagateError(Status status) {
    logger.log(ChannelLogLevel.DEBUG, "[grpclb-<{0}>] Error: {1}", serviceName, status);
    if (backendList.isEmpty()) {
      Status error =
          Status.UNAVAILABLE.withCause(status.getCause()).withDescription(status.getDescription());
      maybeUpdatePicker(
          TRANSIENT_FAILURE, new RoundRobinPicker(dropList, Arrays.asList(new ErrorEntry(error))));
    }
  }

  private void returnSubchannelToPool(Subchannel subchannel) {
    subchannelPool.returnSubchannel(subchannel, subchannel.getAttributes().get(STATE_INFO).get());
  }

  @VisibleForTesting
  @Nullable
  GrpclbClientLoadRecorder getLoadRecorder() {
    if (lbStream == null) {
      return null;
    }
    return lbStream.loadRecorder;
  }

  /**
   * Populate backend servers to be used based on the given list of addresses.
   */
  private void updateServerList(
      List<DropEntry> newDropList, List<BackendAddressGroup> newBackendAddrList,
      @Nullable GrpclbClientLoadRecorder loadRecorder) {
    HashMap<List<EquivalentAddressGroup>, Subchannel> newSubchannelMap =
        new HashMap<>();
    List<BackendEntry> newBackendList = new ArrayList<>();

    switch (config.getMode()) {
      case ROUND_ROBIN:
        for (BackendAddressGroup backendAddr : newBackendAddrList) {
          EquivalentAddressGroup eag = backendAddr.getAddresses();
          List<EquivalentAddressGroup> eagAsList = Collections.singletonList(eag);
          Subchannel subchannel = newSubchannelMap.get(eagAsList);
          if (subchannel == null) {
            subchannel = subchannels.get(eagAsList);
            if (subchannel == null) {
              subchannel = subchannelPool.takeOrCreateSubchannel(eag, createSubchannelAttrs());
              subchannel.requestConnection();
            }
            newSubchannelMap.put(eagAsList, subchannel);
          }
          BackendEntry entry;
          // Only picks with tokens are reported to LoadRecorder
          if (backendAddr.getToken() == null) {
            entry = new BackendEntry(subchannel);
          } else {
            entry = new BackendEntry(subchannel, loadRecorder, backendAddr.getToken());
          }
          newBackendList.add(entry);
        }
        // Close Subchannels whose addresses have been delisted
        for (Map.Entry<List<EquivalentAddressGroup>, Subchannel> entry : subchannels.entrySet()) {
          List<EquivalentAddressGroup> eagList = entry.getKey();
          if (!newSubchannelMap.containsKey(eagList)) {
            returnSubchannelToPool(entry.getValue());
          }
        }
        subchannels = Collections.unmodifiableMap(newSubchannelMap);
        break;
      case PICK_FIRST:
        // Delegate to child pick_first LB for address management.
        // Shutdown existing child LB if addresses become empty.
        if (newBackendAddrList.isEmpty()) {
          if (pickFirstLb != null) {
            pickFirstLb.shutdown();
            pickFirstLb = null;
          }
          break;
        }
        List<EquivalentAddressGroup> eagList = new ArrayList<>();
        // Attach tokens to EAG attributes for TokenAttachingTracerFactory to retrieve.
        for (BackendAddressGroup bag : newBackendAddrList) {
          EquivalentAddressGroup origEag = bag.getAddresses();
          Attributes eagAttrs = origEag.getAttributes();
          if (bag.getToken() != null) {
            eagAttrs = eagAttrs.toBuilder()
                .set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, bag.getToken()).build();
          }
          eagList.add(new EquivalentAddressGroup(origEag.getAddresses(), eagAttrs));
        }

        if (pickFirstLb == null) {
          pickFirstLb = pickFirstLbProvider.newLoadBalancer(new PickFirstLbHelper());
        }

        // Pass addresses to child LB.
        pickFirstLb.acceptResolvedAddresses(
            ResolvedAddresses.newBuilder()
                .setAddresses(eagList)
                .build());
        if (requestConnectionPending) {
          pickFirstLb.requestConnection();
          requestConnectionPending = false;
        }
        // Store the load recorder for token attachment.
        currentPickFirstLoadRecorder = loadRecorder;
        break;
      default:
        throw new AssertionError("Missing case for " + config.getMode());
    }

    dropList = Collections.unmodifiableList(newDropList);
    backendList = Collections.unmodifiableList(newBackendList);
  }

  @VisibleForTesting
  class FallbackModeTask implements Runnable {
    private final Status reason;

    private FallbackModeTask(Status reason) {
      this.reason = reason;
    }

    @Override
    public void run() {
      // Timer should have been cancelled if entered fallback early.
      checkState(!usingFallbackBackends, "already in fallback");
      fallbackReason = reason;
      maybeUseFallbackBackends();
      maybeUpdatePicker();
    }
  }

  @VisibleForTesting
  class LbRpcRetryTask implements Runnable {
    @Override
    public void run() {
      startLbRpc();
    }
  }

  @VisibleForTesting
  static class LoadReportingTask implements Runnable {
    private final LbStream stream;

    LoadReportingTask(LbStream stream) {
      this.stream = stream;
    }

    @Override
    public void run() {
      stream.loadReportTimer = null;
      stream.sendLoadReport();
    }
  }

  private class LbStream implements StreamObserver<LoadBalanceResponse> {
    final GrpclbClientLoadRecorder loadRecorder;
    final LoadBalancerGrpc.LoadBalancerStub stub;
    StreamObserver<LoadBalanceRequest> lbRequestWriter;

    // These fields are only accessed from helper.runSerialized()
    boolean initialResponseReceived;
    boolean closed;
    long loadReportIntervalMillis = -1;
    ScheduledHandle loadReportTimer;

    LbStream(LoadBalancerGrpc.LoadBalancerStub stub) {
      this.stub = checkNotNull(stub, "stub");
      // Stats data only valid for current LbStream.  We do not carry over data from previous
      // stream.
      loadRecorder = new GrpclbClientLoadRecorder(time);
    }

    void start() {
      lbRequestWriter = stub.withWaitForReady().balanceLoad(this);
    }

    @Override public void onNext(final LoadBalanceResponse response) {
      syncContext.execute(new Runnable() {
          @Override
          public void run() {
            handleResponse(response);
          }
        });
    }

    @Override public void onError(final Throwable error) {
      syncContext.execute(new Runnable() {
          @Override
          public void run() {
            handleStreamClosed(Status.fromThrowable(error)
                .augmentDescription("Stream to GRPCLB LoadBalancer had an error"));
          }
        });
    }

    @Override public void onCompleted() {
      syncContext.execute(new Runnable() {
          @Override
          public void run() {
            handleStreamClosed(
                Status.UNAVAILABLE.withDescription("Stream to GRPCLB LoadBalancer was closed"));
          }
        });
    }

    // Following methods must be run in helper.runSerialized()

    private void sendLoadReport() {
      if (closed) {
        return;
      }
      ClientStats stats = loadRecorder.generateLoadReport();
      // TODO(zhangkun83): flow control?
      try {
        lbRequestWriter.onNext(LoadBalanceRequest.newBuilder().setClientStats(stats).build());
        scheduleNextLoadReport();
      } catch (Exception e) {
        close(e);
      }
    }

    private void scheduleNextLoadReport() {
      if (loadReportIntervalMillis > 0) {
        loadReportTimer = syncContext.schedule(
            new LoadReportingTask(this), loadReportIntervalMillis, TimeUnit.MILLISECONDS,
            timerService);
      }
    }

    private void handleResponse(LoadBalanceResponse response) {
      if (closed) {
        return;
      }

      LoadBalanceResponseTypeCase typeCase = response.getLoadBalanceResponseTypeCase();
      if (!initialResponseReceived) {
        logger.log(
            ChannelLogLevel.INFO,
            "[grpclb-<{0}>] Got an LB initial response: {1}", serviceName, response);
        if (typeCase != LoadBalanceResponseTypeCase.INITIAL_RESPONSE) {
          logger.log(
              ChannelLogLevel.WARNING,
              "[grpclb-<{0}>] Received a response without initial response",
              serviceName);
          return;
        }
        initialResponseReceived = true;
        InitialLoadBalanceResponse initialResponse = response.getInitialResponse();
        loadReportIntervalMillis =
            Durations.toMillis(initialResponse.getClientStatsReportInterval());
        scheduleNextLoadReport();
        return;
      }
      if (SHOULD_LOG_SERVER_LISTS) {
        logger.log(
            ChannelLogLevel.DEBUG, "[grpclb-<{0}>] Got an LB response: {1}", serviceName, response);
      } else {
        logger.log(ChannelLogLevel.DEBUG, "[grpclb-<{0}>] Got an LB response", serviceName);
      }

      if (typeCase == LoadBalanceResponseTypeCase.FALLBACK_RESPONSE) {
        // Force entering fallback requested by balancer.
        cancelFallbackTimer();
        fallbackReason = BALANCER_REQUESTED_FALLBACK_STATUS;
        useFallbackBackends();
        maybeUpdatePicker();
        return;
      } else if (typeCase != LoadBalanceResponseTypeCase.SERVER_LIST) {
        logger.log(
            ChannelLogLevel.WARNING,
            "[grpclb-<{0}>] Ignoring unexpected response type: {1}",
            serviceName,
            typeCase);
        return;
      }

      balancerWorking = true;
      // TODO(zhangkun83): handle delegate from initialResponse
      ServerList serverList = response.getServerList();
      List<DropEntry> newDropList = new ArrayList<>();
      List<BackendAddressGroup> newBackendAddrList = new ArrayList<>();
      // Construct the new collections. Create new Subchannels when necessary.
      for (Server server : serverList.getServersList()) {
        String token = server.getLoadBalanceToken();
        if (server.getDrop()) {
          newDropList.add(new DropEntry(loadRecorder, token));
        } else {
          newDropList.add(null);
          InetSocketAddress address;
          try {
            address = new InetSocketAddress(
                InetAddress.getByAddress(server.getIpAddress().toByteArray()), server.getPort());
          } catch (UnknownHostException e) {
            propagateError(
                Status.UNAVAILABLE
                    .withDescription("Invalid backend address: " + server)
                    .withCause(e));
            continue;
          }
          // ALTS code can use the presence of ATTR_LB_PROVIDED_BACKEND to select ALTS instead of
          // TLS, with Netty.
          EquivalentAddressGroup eag =
              new EquivalentAddressGroup(address, LB_PROVIDED_BACKEND_ATTRS);
          newBackendAddrList.add(new BackendAddressGroup(eag, token));
        }
      }
      // Exit fallback as soon as a new server list is received from the balancer.
      usingFallbackBackends = false;
      fallbackReason = null;
      cancelFallbackTimer();
      updateServerList(newDropList, newBackendAddrList, loadRecorder);
      maybeUpdatePicker();
    }

    private void handleStreamClosed(Status error) {
      checkArgument(!error.isOk(), "unexpected OK status");
      if (closed) {
        return;
      }
      closed = true;
      cleanUp();
      propagateError(error);
      balancerWorking = false;
      fallbackReason = error;
      cancelFallbackTimer();
      maybeUseFallbackBackends();
      maybeUpdatePicker();

      long delayNanos = 0;
      if (initialResponseReceived || lbRpcRetryPolicy == null) {
        // Reset the backoff sequence if balancer has sent the initial response, or backoff sequence
        // has never been initialized.
        lbRpcRetryPolicy = backoffPolicyProvider.get();
      }
      // Backoff only when balancer wasn't working previously.
      if (!initialResponseReceived) {
        // The back-off policy determines the interval between consecutive RPC upstarts, thus the
        // actual delay may be smaller than the value from the back-off policy, or even negative,
        // depending how much time was spent in the previous RPC.
        delayNanos =
            lbRpcRetryPolicy.nextBackoffNanos() - stopwatch.elapsed(TimeUnit.NANOSECONDS);
      }
      if (delayNanos <= 0) {
        startLbRpc();
      } else {
        lbRpcRetryTimer =
            syncContext.schedule(new LbRpcRetryTask(), delayNanos, TimeUnit.NANOSECONDS,
                timerService);
      }

      helper.refreshNameResolution();
    }

    void close(Exception error) {
      if (closed) {
        return;
      }
      closed = true;
      cleanUp();
      lbRequestWriter.onError(error);
    }

    private void cleanUp() {
      if (loadReportTimer != null) {
        loadReportTimer.cancel();
        loadReportTimer = null;
      }
      if (lbStream == this) {
        lbStream = null;
      }
    }
  }

  /**
   * 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;
    ConnectivityState state;
    // For PICK_FIRST mode with delegation, check if child LB exists instead of backendList.
    boolean hasBackends = config.getMode() == Mode.PICK_FIRST
        ? pickFirstLb != null
        : !backendList.isEmpty();
    if (!hasBackends) {
      // Note balancer (is working) may enforce using fallback backends, and that fallback may
      // fail. So we should check if currently in fallback first.
      if (usingFallbackBackends) {
        Status error =
            NO_FALLBACK_BACKENDS_STATUS
                .withCause(fallbackReason.getCause())
                .augmentDescription(fallbackReason.getDescription());
        pickList = Collections.<RoundRobinEntry>singletonList(new ErrorEntry(error));
        state = TRANSIENT_FAILURE;
      } else if (balancerWorking)  {
        pickList =
            Collections.<RoundRobinEntry>singletonList(
                new ErrorEntry(NO_AVAILABLE_BACKENDS_STATUS));
        state = TRANSIENT_FAILURE;
      } else {  // still waiting for balancer
        pickList = Collections.singletonList(BUFFER_ENTRY);
        state = CONNECTING;
      }
      maybeUpdatePicker(state, new RoundRobinPicker(dropList, pickList));
      return;
    }
    switch (config.getMode()) {
      case ROUND_ROBIN:
        pickList = new ArrayList<>(backendList.size());
        Status error = null;
        boolean hasPending = false;
        for (BackendEntry entry : backendList) {
          Subchannel subchannel = entry.subchannel;
          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 {
            hasPending = true;
          }
        }
        if (pickList.isEmpty()) {
          if (hasPending) {
            pickList.add(BUFFER_ENTRY);
            state = CONNECTING;
          } else {
            pickList.add(new ErrorEntry(error));
            state = TRANSIENT_FAILURE;
          }
        } else {
          state = READY;
        }
        break;
      case PICK_FIRST: {
        // Use child LB's state and picker. Wrap the picker for token attachment.
        state = pickFirstLbState;
        TokenAttachingTracerFactory tracerFactory =
            new TokenAttachingTracerFactory(currentPickFirstLoadRecorder);
        pickList = Collections.<RoundRobinEntry>singletonList(
            new ChildLbPickerEntry(pickFirstLbPicker, tracerFactory));
        break;
      }
      default:
        throw new AssertionError("Missing case for " + config.getMode());
    }
    maybeUpdatePicker(state, new RoundRobinPicker(dropList, pickList));
  }

  /**
   * Update the given picker to the helper if it's different from the current one.
   */
  private void maybeUpdatePicker(ConnectivityState state, RoundRobinPicker picker) {
    // Discard the new picker if we are sure it won't make any difference, in order to save
    // re-processing pending streams, and avoid unnecessary resetting of the pointer in
    // RoundRobinPicker.
    if (picker.dropList.equals(currentPicker.dropList)
        && picker.pickList.equals(currentPicker.pickList)) {
      return;
    }
    currentPicker = picker;
    helper.updateBalancingState(state, picker);
  }

  private static Attributes createSubchannelAttrs() {
    return Attributes.newBuilder()
        .set(STATE_INFO,
            new AtomicReference<>(
                ConnectivityStateInfo.forNonError(IDLE)))
        .build();
  }

  @VisibleForTesting
  static final class DropEntry {
    private final GrpclbClientLoadRecorder loadRecorder;
    private final String token;

    DropEntry(GrpclbClientLoadRecorder loadRecorder, String token) {
      this.loadRecorder = checkNotNull(loadRecorder, "loadRecorder");
      this.token = checkNotNull(token, "token");
    }

    PickResult picked() {
      loadRecorder.recordDroppedRequest(token);
      return DROP_PICK_RESULT;
    }

    @Override
    public String toString() {
      // This is printed in logs.  Only include useful information.
      return "drop(" + token + ")";
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(loadRecorder, token);
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof DropEntry)) {
        return false;
      }
      DropEntry that = (DropEntry) other;
      return Objects.equal(loadRecorder, that.loadRecorder) && Objects.equal(token, that.token);
    }
  }

  @VisibleForTesting
  interface RoundRobinEntry {
    PickResult picked(PickSubchannelArgs args);
  }

  @VisibleForTesting
  static final class BackendEntry implements RoundRobinEntry {
    final Subchannel subchannel;
    @VisibleForTesting
    final PickResult result;
    @Nullable
    private final String token;

    /**
     * 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");
    }

    /**
     * For ROUND_ROBIN/PICK_FIRST: creates a BackendEntry whose usage will not be reported.
     */
    BackendEntry(Subchannel subchannel) {
      this.subchannel = checkNotNull(subchannel, "subchannel");
      this.result = PickResult.withSubchannel(subchannel);
      this.token = null;
    }

    /**
     * For PICK_FIRST: creates a BackendEntry that includes all addresses.
     */
    BackendEntry(Subchannel subchannel, TokenAttachingTracerFactory tracerFactory) {
      this.subchannel = checkNotNull(subchannel, "subchannel");
      this.result =
          PickResult.withSubchannel(subchannel, checkNotNull(tracerFactory, "tracerFactory"));
      this.token = null;
    }

    @Override
    public PickResult picked(PickSubchannelArgs args) {
      Metadata headers = args.getHeaders();
      headers.discardAll(GrpclbConstants.TOKEN_METADATA_KEY);
      if (token != null) {
        headers.put(GrpclbConstants.TOKEN_METADATA_KEY, token);
      }
      return result;
    }

    @Override
    public String toString() {
      // This is printed in logs.  Only give out useful information.
      return "[" + subchannel.getAllAddresses().toString() + "(" + token + ")]";
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(result, token);
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof BackendEntry)) {
        return false;
      }
      BackendEntry that = (BackendEntry) other;
      return Objects.equal(result, that.result) && Objects.equal(token, that.token);
    }
  }

  @VisibleForTesting
  static final class IdleSubchannelEntry implements RoundRobinEntry {
    private final SynchronizationContext syncContext;
    private final Subchannel subchannel;
    private final AtomicBoolean connectionRequested = new AtomicBoolean(false);

    IdleSubchannelEntry(Subchannel subchannel, SynchronizationContext syncContext) {
      this.subchannel = checkNotNull(subchannel, "subchannel");
      this.syncContext = checkNotNull(syncContext, "syncContext");
    }

    @Override
    public PickResult picked(PickSubchannelArgs args) {
      if (connectionRequested.compareAndSet(false, true)) {
        syncContext.execute(new Runnable() {
            @Override
            public void run() {
              subchannel.requestConnection();
            }
          });
      }
      return PickResult.withNoResult();
    }

    @Override
    public String toString() {
      // This is printed in logs.  Only give out useful information.
      return "(idle)[" + subchannel.getAllAddresses().toString() + "]";
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(subchannel, syncContext);
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof IdleSubchannelEntry)) {
        return false;
      }
      IdleSubchannelEntry that = (IdleSubchannelEntry) other;
      return Objects.equal(subchannel, that.subchannel)
          && Objects.equal(syncContext, that.syncContext);
    }
  }

  @VisibleForTesting
  static final class ErrorEntry implements RoundRobinEntry {
    final PickResult result;

    ErrorEntry(Status status) {
      result = PickResult.withError(status);
    }

    @Override
    public PickResult picked(PickSubchannelArgs args) {
      return result;
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(result);
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof ErrorEntry)) {
        return false;
      }
      return Objects.equal(result, ((ErrorEntry) other).result);
    }

    @Override
    public String toString() {
      // This is printed in logs.  Only include useful information.
      return result.getStatus().toString();
    }
  }

  /**
   * Entry that wraps a child LB's picker for PICK_FIRST mode delegation.
   * Attaches TokenAttachingTracerFactory to the pick result for token propagation.
   */
  @VisibleForTesting
  static final class ChildLbPickerEntry implements RoundRobinEntry {
    private final SubchannelPicker childPicker;
    private final TokenAttachingTracerFactory tracerFactory;

    ChildLbPickerEntry(SubchannelPicker childPicker, TokenAttachingTracerFactory tracerFactory) {
      this.childPicker = checkNotNull(childPicker, "childPicker");
      this.tracerFactory = checkNotNull(tracerFactory, "tracerFactory");
    }

    @Override
    public PickResult picked(PickSubchannelArgs args) {
      PickResult childResult = childPicker.pickSubchannel(args);
      if (childResult.getSubchannel() == null) {
        // No subchannel (e.g., buffer, error), return as-is.
        return childResult;
      }
      // Wrap the pick result to attach tokens via the tracer factory.
      return PickResult.withSubchannel(
          childResult.getSubchannel(), tracerFactory, childResult.getAuthorityOverride());
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(childPicker, tracerFactory);
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof ChildLbPickerEntry)) {
        return false;
      }
      ChildLbPickerEntry that = (ChildLbPickerEntry) other;
      return Objects.equal(childPicker, that.childPicker)
          && Objects.equal(tracerFactory, that.tracerFactory);
    }

    @Override
    public String toString() {
      return "ChildLbPickerEntry(" + childPicker + ")";
    }

    @VisibleForTesting
    SubchannelPicker getChildPicker() {
      return childPicker;
    }
  }

  @VisibleForTesting
  static final class RoundRobinPicker extends SubchannelPicker {
    @VisibleForTesting
    final List<DropEntry> dropList;
    private int dropIndex;

    @VisibleForTesting
    final List<? extends RoundRobinEntry> pickList;
    private int pickIndex;

    // dropList can be empty, which means no drop.
    // pickList must not be empty.
    RoundRobinPicker(List<DropEntry> dropList, List<? extends RoundRobinEntry> pickList) {
      this.dropList = checkNotNull(dropList, "dropList");
      this.pickList = checkNotNull(pickList, "pickList");
      checkArgument(!pickList.isEmpty(), "pickList is empty");
    }

    @Override
    public PickResult pickSubchannel(PickSubchannelArgs args) {
      synchronized (pickList) {
        // Two-level round-robin.
        // First round-robin on dropList. If a drop entry is selected, request will be dropped.  If
        // a non-drop entry is selected, then round-robin on pickList.  This makes sure requests are
        // dropped at the same proportion as the drop entries appear on the round-robin list from
        // the balancer, while only backends from pickList are selected for the non-drop cases.
        if (!dropList.isEmpty()) {
          DropEntry drop = dropList.get(dropIndex);
          dropIndex++;
          if (dropIndex == dropList.size()) {
            dropIndex = 0;
          }
          if (drop != null) {
            return drop.picked();
          }
        }

        RoundRobinEntry pick = pickList.get(pickIndex);
        pickIndex++;
        if (pickIndex == pickList.size()) {
          pickIndex = 0;
        }
        return pick.picked(args);
      }
    }

    @Override
    public String toString() {
      if (SHOULD_LOG_SERVER_LISTS) {
        return MoreObjects.toStringHelper(RoundRobinPicker.class)
            .add("dropList", dropList)
            .add("pickList", pickList)
            .toString();
      }
      return MoreObjects.toStringHelper(RoundRobinPicker.class).toString();
    }
  }

  /**
   * Helper for the child pick_first LB in PICK_FIRST mode. Intercepts updateBalancingState()
   * to store state and trigger the grpclb picker update with drops and token attachment.
   */
  private final class PickFirstLbHelper extends ForwardingLoadBalancerHelper {

    @Override
    protected Helper delegate() {
      return helper;
    }

    @Override
    public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
      pickFirstLbState = newState;
      pickFirstLbPicker = newPicker;
      // Trigger name resolution refresh on TRANSIENT_FAILURE or IDLE, similar to ROUND_ROBIN.
      if (newState == TRANSIENT_FAILURE || newState == IDLE) {
        helper.refreshNameResolution();
      }
      maybeUseFallbackBackends();
      maybeUpdatePicker();
    }
  }
}

grpc / grpc-java / #20125

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous