#20175

Committed 20 Feb 2026 07:37AM UTC coverage: 88.707% (+0.001%) from 88.706%

Build # #20175

Build Type

push

github

Committed by

web-flow

Commit Message

unwrap ForwardingSubchannel during Picks (#12658)

This PR ensures that Load Balancing (LB) policies unwrap
`ForwardingSubchannel` instances before returning them in a
`PickResult`.

**Rationale:** Currently, the identity of a subchannel is "awkward"
because decorators break object identity. This forces the core channel
to use internal workarounds like `getInternalSubchannel()` to find the
underlying implementation. Removing these wrappers during the pick
process is a critical prerequisite for deleting Subchannel Attributes.

By enforcing unwrapping, `ManagedChannelImpl` can rely on the fact that
a returned subchannel is the same instance it originally created. This
allows the channel to use strongly-typed fields for state management
(via "blind casting") rather than abusing attributes to re-discover
information that should already be known. This also paves the way for
the eventual removal of the `getInternalSubchannel()` internal API.

**New APIs:** To ensure we don't "drop data on the floor" during the
unwrapping process, this PR adds two new non-static APIs to PickResult:
- copyWithSubchannel()
- copyWithStreamTracerFactory()

Unlike static factory methods, these instance methods follow a
"copy-and-update" pattern that preserves all existing pick-level
metadata (such as authority overrides or drop status) while only
swapping the specific field required.

Run Details

35450 of 39963 relevant lines covered (88.71%)

0.89 hits per line

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

89.15

/../services/src/main/java/io/grpc/protobuf/services/HealthCheckingLoadBalancerFactory.java

/*
 * Copyright 2018 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.protobuf.services;

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.LoadBalancer.HEALTH_CONSUMER_LISTENER_ARG_KEY;

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.common.base.Supplier;
import io.grpc.CallOptions;
import io.grpc.ChannelLogger;
import io.grpc.ChannelLogger.ChannelLogLevel;
import io.grpc.ClientCall;
import io.grpc.ConnectivityStateInfo;
import io.grpc.LoadBalancer;
import io.grpc.LoadBalancer.CreateSubchannelArgs;
import io.grpc.LoadBalancer.Helper;
import io.grpc.LoadBalancer.Subchannel;
import io.grpc.LoadBalancer.SubchannelStateListener;
import io.grpc.Metadata;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.SynchronizationContext;
import io.grpc.SynchronizationContext.ScheduledHandle;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthCheckResponse.ServingStatus;
import io.grpc.health.v1.HealthGrpc;
import io.grpc.internal.BackoffPolicy;
import io.grpc.internal.ServiceConfigUtil;
import io.grpc.util.ForwardingLoadBalancer;
import io.grpc.util.ForwardingLoadBalancerHelper;
import io.grpc.util.ForwardingSubchannel;
import io.grpc.util.HealthProducerHelper;
import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;

/**
 * Wraps a {@link LoadBalancer} and implements the client-side health-checking
 * (https://github.com/grpc/proposal/blob/master/A17-client-side-health-checking.md).  The
 * Subchannel received by the states wrapped LoadBalancer will be determined by health-checking.
 *
 * <p>Note the original LoadBalancer must call {@code Helper.createSubchannel()} from the
 * SynchronizationContext, or it will throw.
 */
final class HealthCheckingLoadBalancerFactory extends LoadBalancer.Factory {
  private static final Logger logger =
      Logger.getLogger(HealthCheckingLoadBalancerFactory.class.getName());

  private final LoadBalancer.Factory delegateFactory;
  private final BackoffPolicy.Provider backoffPolicyProvider;
  private final Supplier<Stopwatch> stopwatchSupplier;

  public HealthCheckingLoadBalancerFactory(
      LoadBalancer.Factory delegateFactory, BackoffPolicy.Provider backoffPolicyProvider,
      Supplier<Stopwatch> stopwatchSupplier) {
    this.delegateFactory = checkNotNull(delegateFactory, "delegateFactory");
    this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
    this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier");
  }

  @Override
  public LoadBalancer newLoadBalancer(Helper helper) {
    HelperImpl wrappedHelper = new HelperImpl(helper);
    LoadBalancer delegateBalancer = delegateFactory.newLoadBalancer(wrappedHelper);
    return new HealthCheckingLoadBalancer(wrappedHelper, delegateBalancer);
  }

  private final class HelperImpl extends ForwardingLoadBalancerHelper {
    private final Helper delegate;
    private final SynchronizationContext syncContext;

    @Nullable String healthCheckedService;

    final HashSet<HealthCheckState> hcStates = new HashSet<>();

    HelperImpl(Helper delegate) {
      this.delegate = new HealthProducerHelper(checkNotNull(delegate, "delegate"));
      this.syncContext = checkNotNull(delegate.getSynchronizationContext(), "syncContext");
    }

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

    @Override
    public Subchannel createSubchannel(CreateSubchannelArgs args) {
      // HealthCheckState is not thread-safe, we are requiring the original LoadBalancer calls
      // createSubchannel() from the SynchronizationContext.
      syncContext.throwIfNotInThisSynchronizationContext();
      LoadBalancer.SubchannelStateListener healthConsumerListener =
          args.getOption(HEALTH_CONSUMER_LISTENER_ARG_KEY);
      HealthCheckState hcState = new HealthCheckState(
          this, syncContext, delegate.getScheduledExecutorService(), healthConsumerListener);
      if (healthConsumerListener != null) {
        args = args.toBuilder().addOption(HEALTH_CONSUMER_LISTENER_ARG_KEY, hcState).build();
      }
      Subchannel delegate = super.createSubchannel(args);
      hcState.setSubchannel(delegate);
      hcStates.add(hcState);
      Subchannel subchannel = new SubchannelImpl(delegate, hcState);
      if (healthCheckedService != null) {
        hcState.setServiceName(healthCheckedService);
      }
      return subchannel;
    }

    void setHealthCheckedService(@Nullable String service) {
      healthCheckedService = service;
      for (HealthCheckState hcState : hcStates) {
        hcState.setServiceName(service);
      }
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this).add("delegate", delegate()).toString();
    }

    @Override
    public void updateBalancingState(
        io.grpc.ConnectivityState newState, LoadBalancer.SubchannelPicker newPicker) {
      delegate().updateBalancingState(newState, new HealthCheckPicker(newPicker));
    }

    private final class HealthCheckPicker extends LoadBalancer.SubchannelPicker {
      private final LoadBalancer.SubchannelPicker delegate;

      HealthCheckPicker(LoadBalancer.SubchannelPicker delegate) {
        this.delegate = delegate;
      }

      @Override
      public LoadBalancer.PickResult pickSubchannel(LoadBalancer.PickSubchannelArgs args) {
        LoadBalancer.PickResult result = delegate.pickSubchannel(args);
        LoadBalancer.Subchannel subchannel = result.getSubchannel();
        if (subchannel instanceof SubchannelImpl) {
          return result.copyWithSubchannel(((SubchannelImpl) subchannel).delegate());
        }
        return result;
      }
    }
  }

  @VisibleForTesting
  static final class SubchannelImpl extends ForwardingSubchannel {
    final Subchannel delegate;
    final HealthCheckState hcState;

    SubchannelImpl(Subchannel delegate, HealthCheckState hcState) {
      this.delegate = checkNotNull(delegate, "delegate");
      this.hcState = checkNotNull(hcState, "hcState");
    }

    @Override
    protected Subchannel delegate() {
      return delegate;
    }

    @Override
    public void start(final SubchannelStateListener listener) {
      if (hcState.stateListener == null) {
        hcState.init(listener);
        delegate().start(hcState);
      } else {
        delegate().start(listener);
      }
    }
  }

  private static final class HealthCheckingLoadBalancer extends ForwardingLoadBalancer {
    final LoadBalancer delegate;
    final HelperImpl helper;

    HealthCheckingLoadBalancer(HelperImpl helper, LoadBalancer delegate) {
      this.helper = checkNotNull(helper, "helper");
      this.delegate = checkNotNull(delegate, "delegate");
    }

    @Override
    protected LoadBalancer delegate() {
      return delegate;
    }

    @Override
    public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
      Map<String, ?> healthCheckingConfig =
          resolvedAddresses
              .getAttributes()
              .get(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG);
      String serviceName = ServiceConfigUtil.getHealthCheckedServiceName(healthCheckingConfig);
      helper.setHealthCheckedService(serviceName);
      delegate.handleResolvedAddresses(resolvedAddresses);
    }

    @Override
    public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
      Map<String, ?> healthCheckingConfig =
          resolvedAddresses
              .getAttributes()
              .get(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG);
      String serviceName = ServiceConfigUtil.getHealthCheckedServiceName(healthCheckingConfig);
      helper.setHealthCheckedService(serviceName);
      return delegate.acceptResolvedAddresses(resolvedAddresses);
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this).add("delegate", delegate()).toString();
    }
  }

  
  // All methods are run from syncContext
  private final class HealthCheckState implements SubchannelStateListener {
    private final Runnable retryTask = new Runnable() {
        @Override
        public void run() {
          startRpc();
        }
      };

    private final SynchronizationContext syncContext;
    private final ScheduledExecutorService timerService;
    private final HelperImpl helperImpl;
    private Subchannel subchannel;
    private ChannelLogger subchannelLogger;

    // In dual stack new pick first, this becomes health listener from create subchannel args.
    private SubchannelStateListener stateListener;

    // Set when RPC started. Cleared when the RPC has closed or abandoned.
    @Nullable
    private HcStream activeRpc;

    // The service name that should be used for health checking
    private String serviceName;
    private BackoffPolicy backoffPolicy;
    // The state from the underlying Subchannel
    private ConnectivityStateInfo rawState = ConnectivityStateInfo.forNonError(IDLE);
    // The state concluded from health checking
    private ConnectivityStateInfo concludedState = ConnectivityStateInfo.forNonError(IDLE);
    // true if a health check stream should be kept.  When true, either there is an active RPC, or a
    // retry is pending.
    private boolean running;
    // true if server returned UNIMPLEMENTED
    private boolean disabled;
    private ScheduledHandle retryTimer;

    HealthCheckState(
        HelperImpl helperImpl, SynchronizationContext syncContext,
        ScheduledExecutorService timerService,
        @Nullable SubchannelStateListener healthListener) {
      this.helperImpl = checkNotNull(helperImpl, "helperImpl");
      this.syncContext = checkNotNull(syncContext, "syncContext");
      this.timerService = checkNotNull(timerService, "timerService");
      this.stateListener = healthListener;
    }

    void setSubchannel(Subchannel subchannel) {
      this.subchannel = checkNotNull(subchannel, "subchannel");
      this.subchannelLogger = checkNotNull(subchannel.getChannelLogger(), "subchannelLogger");
    }

    // Only called in old pick first. Can be removed after migration.
    void init(SubchannelStateListener listener) {
      checkState(this.stateListener == null, "init() already called");
      this.stateListener = checkNotNull(listener, "listener");
    }

    void setServiceName(@Nullable String newServiceName) {
      if (Objects.equal(newServiceName, serviceName)) {
        return;
      }
      serviceName = newServiceName;
      // If service name has changed while there is active RPC, cancel it so that
      // a new call will be made with the new name.
      String cancelMsg =
          serviceName == null ? "Health check disabled by service config"
          : "Switching to new service name: " + newServiceName;
      stopRpc(cancelMsg);
      adjustHealthCheck();
    }

    @Override
    public void onSubchannelState(ConnectivityStateInfo rawState) {
      if (Objects.equal(this.rawState.getState(), READY)
          && !Objects.equal(rawState.getState(), READY)) {
        // A connection was lost.  We will reset disabled flag because health check
        // may be available on the new connection.
        disabled = false;
      }
      if (Objects.equal(rawState.getState(), SHUTDOWN)) {
        helperImpl.hcStates.remove(this);
      }
      this.rawState = rawState;
      adjustHealthCheck();
    }

    private boolean isRetryTimerPending() {
      return retryTimer != null && retryTimer.isPending();
    }

    // Start or stop health check according to the current states.
    private void adjustHealthCheck() {
      if (!disabled && serviceName != null && Objects.equal(rawState.getState(), READY)) {
        running = true;
        if (activeRpc == null && !isRetryTimerPending()) {
          startRpc();
        }
      } else {
        running = false;
        // Prerequisites for health checking not met.
        // Make sure it's stopped.
        stopRpc("Client stops health check");
        backoffPolicy = null;
        gotoState(rawState);
      }
    }

    private void startRpc() {
      checkState(serviceName != null, "serviceName is null");
      checkState(activeRpc == null, "previous health-checking RPC has not been cleaned up");
      checkState(subchannel != null, "init() not called");
      // Optimization suggested by @markroth: if we are already READY and starting the health
      // checking RPC, either because health check is just enabled or has switched to a new service
      // name, we don't go to CONNECTING, otherwise there will be artificial delays on RPCs
      // waiting for the health check to respond.
      if (!Objects.equal(concludedState.getState(), READY)) {
        subchannelLogger.log(
            ChannelLogLevel.INFO, "CONNECTING: Starting health-check for \"{0}\"", serviceName);
        gotoState(ConnectivityStateInfo.forNonError(CONNECTING));
      }
      activeRpc = new HcStream();
      activeRpc.start();
    }

    private void stopRpc(String msg) {
      if (activeRpc != null) {
        activeRpc.cancel(msg);
        // Abandon this RPC.  We are not interested in anything from this RPC any more.
        activeRpc = null;
      }
      if (retryTimer != null) {
        retryTimer.cancel();
        retryTimer = null;
      }
    }

    private void gotoState(ConnectivityStateInfo newState) {
      checkState(subchannel != null, "init() not called");
      if (!Objects.equal(concludedState, newState)) {
        concludedState = newState;
        stateListener.onSubchannelState(concludedState);
      }
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(this)
          .add("running", running)
          .add("disabled", disabled)
          .add("activeRpc", activeRpc)
          .add("serviceName", serviceName)
          .add("rawState", rawState)
          .add("concludedState", concludedState)
          .toString();
    }

    private class HcStream extends ClientCall.Listener<HealthCheckResponse> {
      private final ClientCall<HealthCheckRequest, HealthCheckResponse> call;
      private final String callServiceName;
      private final Stopwatch stopwatch;
      private boolean callHasResponded;

      HcStream() {
        stopwatch = stopwatchSupplier.get().start();
        callServiceName = serviceName;
        call = subchannel.asChannel().newCall(HealthGrpc.getWatchMethod(), CallOptions.DEFAULT);
      }

      void start() {
        call.start(this, new Metadata());
        call.sendMessage(HealthCheckRequest.newBuilder().setService(serviceName).build());
        call.halfClose();
        call.request(1);
      }

      void cancel(String msg) {
        call.cancel(msg, null);
      }

      @Override
      public void onMessage(final HealthCheckResponse response) {
        syncContext.execute(new Runnable() {
            @Override
            public void run() {
              if (activeRpc == HcStream.this) {
                handleResponse(response);
              }
            }
          });
      }

      @Override
      public void onClose(final Status status, Metadata trailers) {
        syncContext.execute(new Runnable() {
            @Override
            public void run() {
              if (activeRpc == HcStream.this) {
                activeRpc = null;
                handleStreamClosed(status);
              }
            }
          });
      }

      void handleResponse(HealthCheckResponse response) {
        callHasResponded = true;
        backoffPolicy = null;
        ServingStatus status = response.getStatus();
        // running == true means the Subchannel's state (rawState) is READY
        if (Objects.equal(status, ServingStatus.SERVING)) {
          subchannelLogger.log(ChannelLogLevel.INFO, "READY: health-check responded SERVING");
          gotoState(ConnectivityStateInfo.forNonError(READY));
        } else {
          subchannelLogger.log(
              ChannelLogLevel.INFO, "TRANSIENT_FAILURE: health-check responded {0}", status);
          String errorDescription =  "Health-check service responded "
              + status + " for '" + callServiceName + "'";
          gotoState(ConnectivityStateInfo.forTransientFailure(
              Status.UNAVAILABLE.withDescription(errorDescription)));
        }
        call.request(1);
      }

      void handleStreamClosed(Status status) {
        if (Objects.equal(status.getCode(), Code.UNIMPLEMENTED)) {
          disabled = true;
          logger.log(
              Level.SEVERE, "Health-check with {0} is disabled. Server returned: {1}",
              new Object[] {subchannel.getAllAddresses(), status});
          subchannelLogger.log(ChannelLogLevel.ERROR, "Health-check disabled: {0}", status);
          subchannelLogger.log(ChannelLogLevel.INFO, "{0} (no health-check)", rawState);
          gotoState(rawState);
          return;
        }
        long delayNanos = 0;
        subchannelLogger.log(
            ChannelLogLevel.INFO, "TRANSIENT_FAILURE: health-check stream closed with {0}", status);
        String errorDescription = "Health-check stream unexpectedly closed with "
            + status + " for '" + callServiceName + "'";
        gotoState(ConnectivityStateInfo.forTransientFailure(
            Status.UNAVAILABLE.withDescription(errorDescription)));
        // Use backoff only when server has not responded for the previous call
        if (!callHasResponded) {
          if (backoffPolicy == null) {
            backoffPolicy = backoffPolicyProvider.get();
          }
          delayNanos = backoffPolicy.nextBackoffNanos() - stopwatch.elapsed(TimeUnit.NANOSECONDS);
        }
        if (delayNanos <= 0) {
          startRpc();
        } else {
          checkState(!isRetryTimerPending(), "Retry double scheduled");
          subchannelLogger.log(
              ChannelLogLevel.DEBUG, "Will retry health-check after {0} ns", delayNanos);
          retryTimer = syncContext.schedule(
              retryTask, delayNanos, TimeUnit.NANOSECONDS, timerService);
        }
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper(this)
            .add("callStarted", call != null)
            .add("serviceName", callServiceName)
            .add("hasResponded", callHasResponded)
            .toString();
      }
    }
  }
}

1	/*
2	* Copyright 2018 The gRPC Authors
3	*
4	* Licensed under the Apache License, Version 2.0 (the "License");
5	* you may not use this file except in compliance with the License.
6	* You may obtain a copy of the License at
7	*
8	* http://www.apache.org/licenses/LICENSE-2.0
9	*
10	* Unless required by applicable law or agreed to in writing, software
11	* distributed under the License is distributed on an "AS IS" BASIS,
12	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	* See the License for the specific language governing permissions and
14	* limitations under the License.
15	*/
16
17	package io.grpc.protobuf.services;
18
19	import static com.google.common.base.Preconditions.checkNotNull;
20	import static com.google.common.base.Preconditions.checkState;
21	import static io.grpc.ConnectivityState.CONNECTING;
22	import static io.grpc.ConnectivityState.IDLE;
23	import static io.grpc.ConnectivityState.READY;
24	import static io.grpc.ConnectivityState.SHUTDOWN;
25	import static io.grpc.LoadBalancer.HEALTH_CONSUMER_LISTENER_ARG_KEY;
26
27	import com.google.common.annotations.VisibleForTesting;
28	import com.google.common.base.MoreObjects;
29	import com.google.common.base.Objects;
30	import com.google.common.base.Stopwatch;
31	import com.google.common.base.Supplier;
32	import io.grpc.CallOptions;
33	import io.grpc.ChannelLogger;
34	import io.grpc.ChannelLogger.ChannelLogLevel;
35	import io.grpc.ClientCall;
36	import io.grpc.ConnectivityStateInfo;
37	import io.grpc.LoadBalancer;
38	import io.grpc.LoadBalancer.CreateSubchannelArgs;
39	import io.grpc.LoadBalancer.Helper;
40	import io.grpc.LoadBalancer.Subchannel;
41	import io.grpc.LoadBalancer.SubchannelStateListener;
42	import io.grpc.Metadata;
43	import io.grpc.Status;
44	import io.grpc.Status.Code;
45	import io.grpc.SynchronizationContext;
46	import io.grpc.SynchronizationContext.ScheduledHandle;
47	import io.grpc.health.v1.HealthCheckRequest;
48	import io.grpc.health.v1.HealthCheckResponse;
49	import io.grpc.health.v1.HealthCheckResponse.ServingStatus;
50	import io.grpc.health.v1.HealthGrpc;
51	import io.grpc.internal.BackoffPolicy;
52	import io.grpc.internal.ServiceConfigUtil;
53	import io.grpc.util.ForwardingLoadBalancer;
54	import io.grpc.util.ForwardingLoadBalancerHelper;
55	import io.grpc.util.ForwardingSubchannel;
56	import io.grpc.util.HealthProducerHelper;
57	import java.util.HashSet;
58	import java.util.Map;
59	import java.util.concurrent.ScheduledExecutorService;
60	import java.util.concurrent.TimeUnit;
61	import java.util.logging.Level;
62	import java.util.logging.Logger;
63	import javax.annotation.Nullable;
64
65	/**
66	* Wraps a {@link LoadBalancer} and implements the client-side health-checking
67	* (https://github.com/grpc/proposal/blob/master/A17-client-side-health-checking.md). The
68	* Subchannel received by the states wrapped LoadBalancer will be determined by health-checking.
69	*
70	* <p>Note the original LoadBalancer must call {@code Helper.createSubchannel()} from the
71	* SynchronizationContext, or it will throw.
72	*/
73	final class HealthCheckingLoadBalancerFactory extends LoadBalancer.Factory {
74	private static final Logger logger =	1✔
75	Logger.getLogger(HealthCheckingLoadBalancerFactory.class.getName());	1✔
76
77	private final LoadBalancer.Factory delegateFactory;
78	private final BackoffPolicy.Provider backoffPolicyProvider;
79	private final Supplier<Stopwatch> stopwatchSupplier;
80
81	public HealthCheckingLoadBalancerFactory(
82	LoadBalancer.Factory delegateFactory, BackoffPolicy.Provider backoffPolicyProvider,
83	Supplier<Stopwatch> stopwatchSupplier) {	1✔
84	this.delegateFactory = checkNotNull(delegateFactory, "delegateFactory");	1✔
85	this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");	1✔
86	this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier");	1✔
87	}	1✔
88
89	@Override
90	public LoadBalancer newLoadBalancer(Helper helper) {
91	HelperImpl wrappedHelper = new HelperImpl(helper);	1✔
92	LoadBalancer delegateBalancer = delegateFactory.newLoadBalancer(wrappedHelper);	1✔
93	return new HealthCheckingLoadBalancer(wrappedHelper, delegateBalancer);	1✔
94	}
95
96	private final class HelperImpl extends ForwardingLoadBalancerHelper {
97	private final Helper delegate;
98	private final SynchronizationContext syncContext;
99
100	@Nullable String healthCheckedService;
101
102	final HashSet<HealthCheckState> hcStates = new HashSet<>();	1✔
103
104	HelperImpl(Helper delegate) {	1✔
105	this.delegate = new HealthProducerHelper(checkNotNull(delegate, "delegate"));	1✔
106	this.syncContext = checkNotNull(delegate.getSynchronizationContext(), "syncContext");	1✔
107	}	1✔
108
109	@Override
110	protected Helper delegate() {
111	return delegate;	1✔
112	}
113
114	@Override
115	public Subchannel createSubchannel(CreateSubchannelArgs args) {
116	// HealthCheckState is not thread-safe, we are requiring the original LoadBalancer calls
117	// createSubchannel() from the SynchronizationContext.
118	syncContext.throwIfNotInThisSynchronizationContext();	1✔
119	LoadBalancer.SubchannelStateListener healthConsumerListener =	1✔
120	args.getOption(HEALTH_CONSUMER_LISTENER_ARG_KEY);	1✔
121	HealthCheckState hcState = new HealthCheckState(	1✔
122	this, syncContext, delegate.getScheduledExecutorService(), healthConsumerListener);	1✔
123	if (healthConsumerListener != null) {	1✔
124	args = args.toBuilder().addOption(HEALTH_CONSUMER_LISTENER_ARG_KEY, hcState).build();	1✔
125	}
126	Subchannel delegate = super.createSubchannel(args);	1✔
127	hcState.setSubchannel(delegate);	1✔
128	hcStates.add(hcState);	1✔
129	Subchannel subchannel = new SubchannelImpl(delegate, hcState);	1✔
130	if (healthCheckedService != null) {	1✔
131	hcState.setServiceName(healthCheckedService);	1✔
132	}
133	return subchannel;	1✔
134	}
135
136	void setHealthCheckedService(@Nullable String service) {
137	healthCheckedService = service;	1✔
138	for (HealthCheckState hcState : hcStates) {	1✔
139	hcState.setServiceName(service);	1✔
140	}	1✔
141	}	1✔
142
143	@Override
144	public String toString() {
145	return MoreObjects.toStringHelper(this).add("delegate", delegate()).toString();	×
146	}
147
148	@Override
149	public void updateBalancingState(
150	io.grpc.ConnectivityState newState, LoadBalancer.SubchannelPicker newPicker) {
151	delegate().updateBalancingState(newState, new HealthCheckPicker(newPicker));	1✔
152	}	1✔
153
154	private final class HealthCheckPicker extends LoadBalancer.SubchannelPicker {
155	private final LoadBalancer.SubchannelPicker delegate;
156
157	HealthCheckPicker(LoadBalancer.SubchannelPicker delegate) {	1✔
158	this.delegate = delegate;	1✔
159	}	1✔
160
161	@Override
162	public LoadBalancer.PickResult pickSubchannel(LoadBalancer.PickSubchannelArgs args) {
163	LoadBalancer.PickResult result = delegate.pickSubchannel(args);	1✔
164	LoadBalancer.Subchannel subchannel = result.getSubchannel();	1✔
165	if (subchannel instanceof SubchannelImpl) {	1✔
166	return result.copyWithSubchannel(((SubchannelImpl) subchannel).delegate());	1✔
167	}
168	return result;	×
169	}
170	}
171	}
172
173	@VisibleForTesting
174	static final class SubchannelImpl extends ForwardingSubchannel {
175	final Subchannel delegate;
176	final HealthCheckState hcState;
177
178	SubchannelImpl(Subchannel delegate, HealthCheckState hcState) {	1✔
179	this.delegate = checkNotNull(delegate, "delegate");	1✔
180	this.hcState = checkNotNull(hcState, "hcState");	1✔
181	}	1✔
182
183	@Override
184	protected Subchannel delegate() {
185	return delegate;	1✔
186	}
187
188	@Override
189	public void start(final SubchannelStateListener listener) {
190	if (hcState.stateListener == null) {	1✔
191	hcState.init(listener);	1✔
192	delegate().start(hcState);	1✔
193	} else {
194	delegate().start(listener);	1✔
195	}
196	}	1✔
197	}
198
199	private static final class HealthCheckingLoadBalancer extends ForwardingLoadBalancer {
200	final LoadBalancer delegate;
201	final HelperImpl helper;
202
203	HealthCheckingLoadBalancer(HelperImpl helper, LoadBalancer delegate) {	1✔
204	this.helper = checkNotNull(helper, "helper");	1✔
205	this.delegate = checkNotNull(delegate, "delegate");	1✔
206	}	1✔
207
208	@Override
209	protected LoadBalancer delegate() {
210	return delegate;	1✔
211	}
212
213	@Override
214	public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
215	Map<String, ?> healthCheckingConfig =	×
216	resolvedAddresses
217	.getAttributes()	×
218	.get(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG);	×
219	String serviceName = ServiceConfigUtil.getHealthCheckedServiceName(healthCheckingConfig);	×
220	helper.setHealthCheckedService(serviceName);	×
221	delegate.handleResolvedAddresses(resolvedAddresses);	×
222	}	×
223
224	@Override
225	public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
226	Map<String, ?> healthCheckingConfig =	1✔
227	resolvedAddresses
228	.getAttributes()	1✔
229	.get(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG);	1✔
230	String serviceName = ServiceConfigUtil.getHealthCheckedServiceName(healthCheckingConfig);	1✔
231	helper.setHealthCheckedService(serviceName);	1✔
232	return delegate.acceptResolvedAddresses(resolvedAddresses);	1✔
233	}
234
235	@Override
236	public String toString() {
237	return MoreObjects.toStringHelper(this).add("delegate", delegate()).toString();	×
238	}
239	}
240
241
242	// All methods are run from syncContext
243	private final class HealthCheckState implements SubchannelStateListener {
244	private final Runnable retryTask = new Runnable() {	1✔
245	@Override
246	public void run() {
247	startRpc();	1✔
248	}	1✔
249	};
250
251	private final SynchronizationContext syncContext;
252	private final ScheduledExecutorService timerService;
253	private final HelperImpl helperImpl;
254	private Subchannel subchannel;
255	private ChannelLogger subchannelLogger;
256
257	// In dual stack new pick first, this becomes health listener from create subchannel args.
258	private SubchannelStateListener stateListener;
259
260	// Set when RPC started. Cleared when the RPC has closed or abandoned.
261	@Nullable
262	private HcStream activeRpc;
263
264	// The service name that should be used for health checking
265	private String serviceName;
266	private BackoffPolicy backoffPolicy;
267	// The state from the underlying Subchannel
268	private ConnectivityStateInfo rawState = ConnectivityStateInfo.forNonError(IDLE);	1✔
269	// The state concluded from health checking
270	private ConnectivityStateInfo concludedState = ConnectivityStateInfo.forNonError(IDLE);	1✔
271	// true if a health check stream should be kept. When true, either there is an active RPC, or a
272	// retry is pending.
273	private boolean running;
274	// true if server returned UNIMPLEMENTED
275	private boolean disabled;
276	private ScheduledHandle retryTimer;
277
278	HealthCheckState(
279	HelperImpl helperImpl, SynchronizationContext syncContext,
280	ScheduledExecutorService timerService,
281	@Nullable SubchannelStateListener healthListener) {	1✔
282	this.helperImpl = checkNotNull(helperImpl, "helperImpl");	1✔
283	this.syncContext = checkNotNull(syncContext, "syncContext");	1✔
284	this.timerService = checkNotNull(timerService, "timerService");	1✔
285	this.stateListener = healthListener;	1✔
286	}	1✔
287
288	void setSubchannel(Subchannel subchannel) {
289	this.subchannel = checkNotNull(subchannel, "subchannel");	1✔
290	this.subchannelLogger = checkNotNull(subchannel.getChannelLogger(), "subchannelLogger");	1✔
291	}	1✔
292
293	// Only called in old pick first. Can be removed after migration.
294	void init(SubchannelStateListener listener) {
295	checkState(this.stateListener == null, "init() already called");	1✔
296	this.stateListener = checkNotNull(listener, "listener");	1✔
297	}	1✔
298
299	void setServiceName(@Nullable String newServiceName) {
300	if (Objects.equal(newServiceName, serviceName)) {	1✔
301	return;	1✔
302	}
303	serviceName = newServiceName;	1✔
304	// If service name has changed while there is active RPC, cancel it so that
305	// a new call will be made with the new name.
306	String cancelMsg =
307	serviceName == null ? "Health check disabled by service config"	1✔
308	: "Switching to new service name: " + newServiceName;	1✔
309	stopRpc(cancelMsg);	1✔
310	adjustHealthCheck();	1✔
311	}	1✔
312
313	@Override
314	public void onSubchannelState(ConnectivityStateInfo rawState) {
315	if (Objects.equal(this.rawState.getState(), READY)	1✔
316	&& !Objects.equal(rawState.getState(), READY)) {	1✔
317	// A connection was lost. We will reset disabled flag because health check
318	// may be available on the new connection.
319	disabled = false;	1✔
320	}
321	if (Objects.equal(rawState.getState(), SHUTDOWN)) {	1✔
322	helperImpl.hcStates.remove(this);	1✔
323	}
324	this.rawState = rawState;	1✔
325	adjustHealthCheck();	1✔
326	}	1✔
327
328	private boolean isRetryTimerPending() {
329	return retryTimer != null && retryTimer.isPending();	1✔
330	}
331
332	// Start or stop health check according to the current states.
333	private void adjustHealthCheck() {
334	if (!disabled && serviceName != null && Objects.equal(rawState.getState(), READY)) {	1✔
335	running = true;	1✔
336	if (activeRpc == null && !isRetryTimerPending()) {	1✔
337	startRpc();	1✔
338	}
339	} else {
340	running = false;	1✔
341	// Prerequisites for health checking not met.
342	// Make sure it's stopped.
343	stopRpc("Client stops health check");	1✔
344	backoffPolicy = null;	1✔
345	gotoState(rawState);	1✔
346	}
347	}	1✔
348
349	private void startRpc() {
350	checkState(serviceName != null, "serviceName is null");	1✔
351	checkState(activeRpc == null, "previous health-checking RPC has not been cleaned up");	1✔
352	checkState(subchannel != null, "init() not called");	1✔
353	// Optimization suggested by @markroth: if we are already READY and starting the health
354	// checking RPC, either because health check is just enabled or has switched to a new service
355	// name, we don't go to CONNECTING, otherwise there will be artificial delays on RPCs
356	// waiting for the health check to respond.
357	if (!Objects.equal(concludedState.getState(), READY)) {	1✔
358	subchannelLogger.log(	1✔
359	ChannelLogLevel.INFO, "CONNECTING: Starting health-check for \"{0}\"", serviceName);
360	gotoState(ConnectivityStateInfo.forNonError(CONNECTING));	1✔
361	}
362	activeRpc = new HcStream();	1✔
363	activeRpc.start();	1✔
364	}	1✔
365
366	private void stopRpc(String msg) {
367	if (activeRpc != null) {	1✔
368	activeRpc.cancel(msg);	1✔
369	// Abandon this RPC. We are not interested in anything from this RPC any more.
370	activeRpc = null;	1✔
371	}
372	if (retryTimer != null) {	1✔
373	retryTimer.cancel();	1✔
374	retryTimer = null;	1✔
375	}
376	}	1✔
377
378	private void gotoState(ConnectivityStateInfo newState) {
379	checkState(subchannel != null, "init() not called");	1✔
380	if (!Objects.equal(concludedState, newState)) {	1✔
381	concludedState = newState;	1✔
382	stateListener.onSubchannelState(concludedState);	1✔
383	}
384	}	1✔
385
386	@Override
387	public String toString() {
388	return MoreObjects.toStringHelper(this)	×
389	.add("running", running)	×
390	.add("disabled", disabled)	×
391	.add("activeRpc", activeRpc)	×
392	.add("serviceName", serviceName)	×
393	.add("rawState", rawState)	×
394	.add("concludedState", concludedState)	×
395	.toString();	×
396	}
397
398	private class HcStream extends ClientCall.Listener<HealthCheckResponse> {
399	private final ClientCall<HealthCheckRequest, HealthCheckResponse> call;
400	private final String callServiceName;
401	private final Stopwatch stopwatch;
402	private boolean callHasResponded;
403
404	HcStream() {	1✔
405	stopwatch = stopwatchSupplier.get().start();	1✔
406	callServiceName = serviceName;	1✔
407	call = subchannel.asChannel().newCall(HealthGrpc.getWatchMethod(), CallOptions.DEFAULT);	1✔
408	}	1✔
409
410	void start() {
411	call.start(this, new Metadata());	1✔
412	call.sendMessage(HealthCheckRequest.newBuilder().setService(serviceName).build());	1✔
413	call.halfClose();	1✔
414	call.request(1);	1✔
415	}	1✔
416
417	void cancel(String msg) {
418	call.cancel(msg, null);	1✔
419	}	1✔
420
421	@Override
422	public void onMessage(final HealthCheckResponse response) {
423	syncContext.execute(new Runnable() {	1✔
424	@Override
425	public void run() {
426	if (activeRpc == HcStream.this) {	1✔
427	handleResponse(response);	1✔
428	}
429	}	1✔
430	});
431	}	1✔
432
433	@Override
434	public void onClose(final Status status, Metadata trailers) {
435	syncContext.execute(new Runnable() {	1✔
436	@Override
437	public void run() {
438	if (activeRpc == HcStream.this) {	1✔
439	activeRpc = null;	1✔
440	handleStreamClosed(status);	1✔
441	}
442	}	1✔
443	});
444	}	1✔
445
446	void handleResponse(HealthCheckResponse response) {
447	callHasResponded = true;	1✔
448	backoffPolicy = null;	1✔
449	ServingStatus status = response.getStatus();	1✔
450	// running == true means the Subchannel's state (rawState) is READY
451	if (Objects.equal(status, ServingStatus.SERVING)) {	1✔
452	subchannelLogger.log(ChannelLogLevel.INFO, "READY: health-check responded SERVING");	1✔
453	gotoState(ConnectivityStateInfo.forNonError(READY));	1✔
454	} else {
455	subchannelLogger.log(	1✔
456	ChannelLogLevel.INFO, "TRANSIENT_FAILURE: health-check responded {0}", status);
457	String errorDescription = "Health-check service responded "	1✔
458	+ status + " for '" + callServiceName + "'";
459	gotoState(ConnectivityStateInfo.forTransientFailure(	1✔
460	Status.UNAVAILABLE.withDescription(errorDescription)));	1✔
461	}
462	call.request(1);	1✔
463	}	1✔
464
465	void handleStreamClosed(Status status) {
466	if (Objects.equal(status.getCode(), Code.UNIMPLEMENTED)) {	1✔
467	disabled = true;	1✔
468	logger.log(	1✔
469	Level.SEVERE, "Health-check with {0} is disabled. Server returned: {1}",
470	new Object[] {subchannel.getAllAddresses(), status});	1✔
471	subchannelLogger.log(ChannelLogLevel.ERROR, "Health-check disabled: {0}", status);	1✔
472	subchannelLogger.log(ChannelLogLevel.INFO, "{0} (no health-check)", rawState);	1✔
473	gotoState(rawState);	1✔
474	return;	1✔
475	}
476	long delayNanos = 0;	1✔
477	subchannelLogger.log(	1✔
478	ChannelLogLevel.INFO, "TRANSIENT_FAILURE: health-check stream closed with {0}", status);
479	String errorDescription = "Health-check stream unexpectedly closed with "	1✔
480	+ status + " for '" + callServiceName + "'";
481	gotoState(ConnectivityStateInfo.forTransientFailure(	1✔
482	Status.UNAVAILABLE.withDescription(errorDescription)));	1✔
483	// Use backoff only when server has not responded for the previous call
484	if (!callHasResponded) {	1✔
485	if (backoffPolicy == null) {	1✔
486	backoffPolicy = backoffPolicyProvider.get();	1✔
487	}
488	delayNanos = backoffPolicy.nextBackoffNanos() - stopwatch.elapsed(TimeUnit.NANOSECONDS);	1✔
489	}
490	if (delayNanos <= 0) {	1✔
491	startRpc();	1✔
492	} else {
493	checkState(!isRetryTimerPending(), "Retry double scheduled");	1✔
494	subchannelLogger.log(	1✔
495	ChannelLogLevel.DEBUG, "Will retry health-check after {0} ns", delayNanos);	1✔
496	retryTimer = syncContext.schedule(	1✔
497	retryTask, delayNanos, TimeUnit.NANOSECONDS, timerService);	1✔
498	}
499	}	1✔
500
501	@Override
502	public String toString() {
503	return MoreObjects.toStringHelper(this)	×
504	.add("callStarted", call != null)	×
505	.add("serviceName", callServiceName)	×
506	.add("hasResponded", callHasResponded)	×
507	.toString();	×
508	}
509	}
510	}
511	}

grpc / grpc-java / #20175

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