grpclb: use a standalone Context for gRPCLB control plane RPCs (grpc#8154)

Inject a standalone Context that is independent of application RPCs to GrpclbLoadBalancer for control plane RPCs. The control plane RPC should be independent and not impacted by the lifetime of Context used for application RPCs.

Author: voidzcy

grpclb/BUILD.bazel

@@ -9,6 +9,7 @@ java_library(
     deps = [
         ":load_balancer_java_grpc",
         "//api",
+        "//context",
         "//core:internal",
         "//core:util",
         "//stub",

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

@@ -23,6 +23,7 @@
 import com.google.common.base.Stopwatch;
 import io.grpc.Attributes;
 import io.grpc.ChannelLogger.ChannelLogLevel;
+import io.grpc.Context;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
 import io.grpc.Status;
@@ -45,6 +46,7 @@ class GrpclbLoadBalancer extends LoadBalancer {
   private static final GrpclbConfig DEFAULT_CONFIG = GrpclbConfig.create(Mode.ROUND_ROBIN);
 
   private final Helper helper;
+  private final Context context;
   private final TimeProvider time;
   private final Stopwatch stopwatch;
   private final SubchannelPool subchannelPool;
@@ -58,11 +60,13 @@ class GrpclbLoadBalancer extends LoadBalancer {
 
   GrpclbLoadBalancer(
       Helper helper,
+      Context context,
       SubchannelPool subchannelPool,
       TimeProvider time,
       Stopwatch stopwatch,
       BackoffPolicy.Provider backoffPolicyProvider) {
     this.helper = checkNotNull(helper, "helper");
+    this.context = checkNotNull(context, "context");
     this.time = checkNotNull(time, "time provider");
     this.stopwatch = checkNotNull(stopwatch, "stopwatch");
     this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider");
@@ -128,7 +132,7 @@ private void recreateStates() {
     checkState(grpclbState == null, "Should've been cleared");
     grpclbState =
         new GrpclbState(
-            config, helper, subchannelPool, time, stopwatch, backoffPolicyProvider);
+            config, helper, context, subchannelPool, time, stopwatch, backoffPolicyProvider);
   }
 
   @Override

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

@@ -17,6 +17,7 @@
 package io.grpc.grpclb;
 
 import com.google.common.base.Stopwatch;
+import io.grpc.Context;
 import io.grpc.Internal;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancerProvider;
@@ -62,6 +63,7 @@ public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
     return
         new GrpclbLoadBalancer(
             helper,
+            Context.ROOT,
             new CachedSubchannelPool(helper),
             TimeProvider.SYSTEM_TIME_PROVIDER,
             Stopwatch.createUnstarted(),

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

@@ -35,6 +35,7 @@
 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.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
@@ -116,6 +117,7 @@ enum Mode {
 
   private final String serviceName;
   private final Helper helper;
+  private final Context context;
   private final SynchronizationContext syncContext;
   @Nullable
   private final SubchannelPool subchannelPool;
@@ -163,12 +165,14 @@ enum Mode {
   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");
@@ -342,7 +346,12 @@ private void startLbRpc() {
     checkState(lbStream == null, "previous lbStream has not been cleared yet");
     LoadBalancerGrpc.LoadBalancerStub stub = LoadBalancerGrpc.newStub(lbCommChannel);
     lbStream = new LbStream(stub);
-    lbStream.start();
+    Context prevContext = context.attach();
+    try {
+      lbStream.start();
+    } finally {
+      context.detach(prevContext);
+    }
     stopwatch.reset().start();
 
     LoadBalanceRequest initRequest = LoadBalanceRequest.newBuilder()

grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java

@@ -54,6 +54,8 @@
 import io.grpc.ClientStreamTracer;
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
+import io.grpc.Context;
+import io.grpc.Context.CancellableContext;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
@@ -227,6 +229,7 @@ public Void answer(InvocationOnMock invocation) {
     when(backoffPolicyProvider.get()).thenReturn(backoffPolicy1, backoffPolicy2);
     balancer = new GrpclbLoadBalancer(
         helper,
+        Context.ROOT,
         subchannelPool,
         fakeClock.getTimeProvider(),
         fakeClock.getStopwatchSupplier().get(),
@@ -2474,6 +2477,43 @@ public void grpclbWorking_lbSendsFallbackMessage() {
         .inOrder();
   }
 
+  @Test
+  public void useIndependentRpcContext() {
+    // Simulates making RPCs within the context of an inbound RPC.
+    CancellableContext cancellableContext = Context.current().withCancellation();
+    Context prevContext = cancellableContext.attach();
+    try {
+      List<EquivalentAddressGroup> backendList = createResolvedBackendAddresses(2);
+      List<EquivalentAddressGroup> grpclbBalancerList = createResolvedBalancerAddresses(2);
+      deliverResolvedAddresses(backendList, grpclbBalancerList);
+
+      List<SocketAddress> addrs = new ArrayList<>();
+      addrs.addAll(grpclbBalancerList.get(0).getAddresses());
+      addrs.addAll(grpclbBalancerList.get(1).getAddresses());
+      Attributes attr = grpclbBalancerList.get(0).getAttributes();
+      EquivalentAddressGroup oobChannelEag = new EquivalentAddressGroup(addrs, attr);
+      verify(helper).createOobChannel(eq(oobChannelEag), eq(lbAuthority(0)));
+      verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
+      StreamObserver<LoadBalanceResponse> lbResponseObserver = lbResponseObserverCaptor.getValue();
+      assertEquals(1, lbRequestObservers.size());
+      StreamObserver<LoadBalanceRequest> lbRequestObserver = lbRequestObservers.poll();
+      verify(lbRequestObserver).onNext(
+          eq(LoadBalanceRequest.newBuilder()
+              .setInitialRequest(
+                  InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build())
+              .build()));
+      lbResponseObserver.onNext(buildInitialResponse());
+
+      // The inbound RPC finishes and closes its context. The outbound RPC's control plane RPC
+      // should not be impacted (no retry).
+      cancellableContext.close();
+      assertEquals(0, fakeClock.numPendingTasks(LB_RPC_RETRY_TASK_FILTER));
+      verifyNoMoreInteractions(mockLbService);
+    } finally {
+      cancellableContext.detach(prevContext);
+    }
+  }
+
   private void deliverSubchannelState(
       final Subchannel subchannel, final ConnectivityStateInfo newState) {
     ((FakeSubchannel) subchannel).updateState(newState);