unwrap ForwardingSubchannel during Picks

api/src/main/java/io/grpc/LoadBalancer.java

@@ -632,6 +632,8 @@ private PickResult(
      *                            stream is created at all in some cases.
      * @since 1.3.0
      */
+    // TODO(shivaspeaks): Need to deprecate old APIs and create new ones, 
+    // per https://github.com/grpc/grpc-java/issues/12662.
     public static PickResult withSubchannel(
         Subchannel subchannel, @Nullable ClientStreamTracer.Factory streamTracerFactory) {
       return new PickResult(
@@ -661,6 +663,28 @@ public static PickResult withSubchannel(Subchannel subchannel) {
       return withSubchannel(subchannel, null);
     }
 
+    /**
+     * Creates a new {@code PickResult} with the given {@code subchannel},
+     * but retains all other properties from this {@code PickResult}.
+     *
+     * @since 1.80.0
+     */
+    public PickResult copyWithSubchannel(Subchannel subchannel) {
+      return new PickResult(checkNotNull(subchannel, "subchannel"), streamTracerFactory,
+          status, drop, authorityOverride);
+    }
+
+    /**
+     * Creates a new {@code PickResult} with the given {@code streamTracerFactory},
+     * but retains all other properties from this {@code PickResult}.
+     *
+     * @since 1.80.0
+     */
+    public PickResult copyWithStreamTracerFactory(
+        @Nullable ClientStreamTracer.Factory streamTracerFactory) {
+      return new PickResult(subchannel, streamTracerFactory, status, drop, authorityOverride);
+    }
+
     /**
      * A decision to report a connectivity error to the RPC.  If the RPC is {@link
      * CallOptions#withWaitForReady wait-for-ready}, it will stay buffered.  Otherwise, it will fail

api/src/test/java/io/grpc/LoadBalancerTest.java

@@ -64,6 +64,26 @@ public void pickResult_withSubchannelAndTracer() {
     assertThat(result.isDrop()).isFalse();
   }
 
+  @Test
+  public void pickResult_withSubchannelReplacement() {
+    PickResult result = PickResult.withSubchannel(subchannel, tracerFactory)
+        .copyWithSubchannel(subchannel2);
+    assertThat(result.getSubchannel()).isSameInstanceAs(subchannel2);
+    assertThat(result.getStatus()).isSameInstanceAs(Status.OK);
+    assertThat(result.getStreamTracerFactory()).isSameInstanceAs(tracerFactory);
+    assertThat(result.isDrop()).isFalse();
+  }
+
+  @Test
+  public void pickResult_withStreamTracerFactory() {
+    PickResult result = PickResult.withSubchannel(subchannel)
+        .copyWithStreamTracerFactory(tracerFactory);
+    assertThat(result.getSubchannel()).isSameInstanceAs(subchannel);
+    assertThat(result.getStatus()).isSameInstanceAs(Status.OK);
+    assertThat(result.getStreamTracerFactory()).isSameInstanceAs(tracerFactory);
+    assertThat(result.isDrop()).isFalse();
+  }
+
   @Test
   public void pickResult_withNoResult() {
     PickResult result = PickResult.withNoResult();

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

@@ -144,6 +144,30 @@ void setHealthCheckedService(@Nullable String service) {
     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

util/src/main/java/io/grpc/util/HealthProducerHelper.java

@@ -22,6 +22,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import io.grpc.Attributes;
+import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.Internal;
 import io.grpc.LoadBalancer;
@@ -84,6 +85,31 @@ protected LoadBalancer.Helper delegate() {
     return delegate;
   }
 
+  @Override
+  public void updateBalancingState(
+      ConnectivityState newState, LoadBalancer.SubchannelPicker newPicker) {
+    delegate.updateBalancingState(newState, new HealthProducerPicker(newPicker));
+  }
+
+  private static final class HealthProducerPicker extends LoadBalancer.SubchannelPicker {
+    private final LoadBalancer.SubchannelPicker delegate;
+
+    HealthProducerPicker(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 HealthProducerSubchannel) {
+        return result.copyWithSubchannel(
+            ((HealthProducerSubchannel) subchannel).delegate());
+      }
+      return result;
+    }
+  }
+
   // The parent subchannel in the health check producer LB chain. It duplicates subchannel state to
   // both the state listener and health listener.
   @VisibleForTesting

util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java

@@ -442,9 +442,14 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
 
       Subchannel subchannel = pickResult.getSubchannel();
       if (subchannel != null) {
-        return PickResult.withSubchannel(subchannel, new ResultCountingClientStreamTracerFactory(
-            subchannel.getAttributes().get(ENDPOINT_TRACKER_KEY),
-            pickResult.getStreamTracerFactory()));
+        EndpointTracker tracker = subchannel.getAttributes().get(ENDPOINT_TRACKER_KEY);
+        if (subchannel instanceof OutlierDetectionSubchannel) {
+          subchannel = ((OutlierDetectionSubchannel) subchannel).delegate();
+        }
+        return pickResult.copyWithSubchannel(subchannel)
+            .copyWithStreamTracerFactory(new ResultCountingClientStreamTracerFactory(
+                tracker,
+                pickResult.getStreamTracerFactory()));
       }
 
       return pickResult;

util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java

@@ -408,7 +408,7 @@ public void delegatePick() throws Exception {
     // Make sure that we can pick the single READY subchannel.
     SubchannelPicker picker = pickerCaptor.getAllValues().get(2);
     PickResult pickResult = picker.pickSubchannel(mock(PickSubchannelArgs.class));
-    Subchannel s = ((OutlierDetectionSubchannel) pickResult.getSubchannel()).delegate();
+    Subchannel s = pickResult.getSubchannel();
     if (s instanceof HealthProducerHelper.HealthProducerSubchannel) {
       s = ((HealthProducerHelper.HealthProducerSubchannel) s).delegate();
     }

xds/src/main/java/io/grpc/xds/ClusterImplLoadBalancer.java

@@ -252,42 +252,55 @@ public Subchannel createSubchannel(CreateSubchannelArgs args) {
       args = args.toBuilder().setAddresses(addresses).setAttributes(attrsBuilder.build()).build();
       final Subchannel subchannel = delegate().createSubchannel(args);
 
-      return new ForwardingSubchannel() {
-        @Override
-        public void start(SubchannelStateListener listener) {
-          delegate().start(new SubchannelStateListener() {
-            @Override
-            public void onSubchannelState(ConnectivityStateInfo newState) {
-              // Do nothing if LB has been shutdown
-              if (xdsClient != null && newState.getState().equals(ConnectivityState.READY)) {
-                // Get locality based on the connected address attributes
-                ClusterLocality updatedClusterLocality = createClusterLocalityFromAttributes(
-                    subchannel.getConnectedAddressAttributes());
-                ClusterLocality oldClusterLocality = localityAtomicReference
-                    .getAndSet(updatedClusterLocality);
-                oldClusterLocality.release();
+      return new ClusterImplSubchannel(subchannel, localityAtomicReference);
+    }
+
+    private final class ClusterImplSubchannel extends ForwardingSubchannel {
+      private final Subchannel delegate;
+      private final AtomicReference<ClusterLocality> localityAtomicReference;
+
+      private ClusterImplSubchannel(
+          Subchannel delegate, AtomicReference<ClusterLocality> localityAtomicReference) {
+        this.delegate = delegate;
+        this.localityAtomicReference = localityAtomicReference;
+      }
+
+      @Override
+      public void start(SubchannelStateListener listener) {
+        delegate().start(
+            new SubchannelStateListener() {
+              @Override
+              public void onSubchannelState(ConnectivityStateInfo newState) {
+                // Do nothing if LB has been shutdown
+                if (xdsClient != null && newState.getState().equals(ConnectivityState.READY)) {
+                  // Get locality based on the connected address attributes
+                  ClusterLocality updatedClusterLocality =
+                      createClusterLocalityFromAttributes(
+                          delegate.getConnectedAddressAttributes());
+                  ClusterLocality oldClusterLocality =
+                      localityAtomicReference.getAndSet(updatedClusterLocality);
+                  oldClusterLocality.release();
+                }
+                listener.onSubchannelState(newState);
               }
-              listener.onSubchannelState(newState);
-            }
-          });
-        }
+            });
+      }
 
-        @Override
-        public void shutdown() {
-          localityAtomicReference.get().release();
-          delegate().shutdown();
-        }
+      @Override
+      public void shutdown() {
+        localityAtomicReference.get().release();
+        delegate().shutdown();
+      }
 
-        @Override
-        public void updateAddresses(List<EquivalentAddressGroup> addresses) {
-          delegate().updateAddresses(withAdditionalAttributes(addresses));
-        }
+      @Override
+      public void updateAddresses(List<EquivalentAddressGroup> addresses) {
+        delegate().updateAddresses(withAdditionalAttributes(addresses));
+      }
 
-        @Override
-        protected Subchannel delegate() {
-          return subchannel;
-        }
-      };
+      @Override
+      protected Subchannel delegate() {
+        return delegate;
+      }
     }
 
     private List<EquivalentAddressGroup> withAdditionalAttributes(
@@ -412,6 +425,11 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
         }
         PickResult result = delegate.pickSubchannel(args);
         if (result.getStatus().isOk() && result.getSubchannel() != null) {
+          Subchannel subchannel = result.getSubchannel();
+          if (subchannel instanceof ClusterImplLbHelper.ClusterImplSubchannel) {
+            subchannel = ((ClusterImplLbHelper.ClusterImplSubchannel) subchannel).delegate();
+            result = result.copyWithSubchannel(subchannel);
+          }
           if (enableCircuitBreaking) {
             if (inFlights.get() >= maxConcurrentRequests) {
               if (dropStats != null) {
@@ -437,8 +455,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
                   stats, inFlights, result.getStreamTracerFactory());
               ClientStreamTracer.Factory orcaTracerFactory = OrcaPerRequestUtil.getInstance()
                   .newOrcaClientStreamTracerFactory(tracerFactory, new OrcaPerRpcListener(stats));
-              result = PickResult.withSubchannel(result.getSubchannel(),
-                  orcaTracerFactory);
+              result = result.copyWithStreamTracerFactory(orcaTracerFactory);
             }
           }
           if (args.getCallOptions().getOption(XdsNameResolver.AUTO_HOST_REWRITE_KEY) != null

xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java

@@ -508,12 +508,15 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
       if (subchannel == null) {
         return pickResult;
       }
+
+      subchannel = ((WrrSubchannel) subchannel).delegate();
       if (!enableOobLoadReport) {
-        return PickResult.withSubchannel(subchannel,
-            OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(
-                reportListeners.get(pick)));
+        return pickResult.copyWithSubchannel(subchannel)
+            .copyWithStreamTracerFactory(
+                OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(
+                    reportListeners.get(pick)));
       } else {
-        return PickResult.withSubchannel(subchannel);
+        return pickResult.copyWithSubchannel(subchannel);
       }
     }
 

xds/src/main/java/io/grpc/xds/orca/OrcaOobUtil.java

@@ -36,12 +36,16 @@
 import io.grpc.ChannelLogger;
 import io.grpc.ChannelLogger.ChannelLogLevel;
 import io.grpc.ClientCall;
+import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.ExperimentalApi;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancer.PickResult;
+import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.Subchannel;
+import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.LoadBalancer.SubchannelStateListener;
 import io.grpc.Metadata;
 import io.grpc.Status;
@@ -236,6 +240,30 @@ protected Helper delegate() {
       return delegate;
     }
 
+    @Override
+    public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+      delegate.updateBalancingState(newState, new OrcaOobPicker(newPicker));
+    }
+
+    @VisibleForTesting
+    static final class OrcaOobPicker extends SubchannelPicker {
+      final SubchannelPicker delegate;
+
+      OrcaOobPicker(SubchannelPicker delegate) {
+        this.delegate = delegate;
+      }
+
+      @Override
+      public PickResult pickSubchannel(PickSubchannelArgs args) {
+        PickResult result = delegate.pickSubchannel(args);
+        Subchannel subchannel = result.getSubchannel();
+        if (subchannel instanceof SubchannelImpl) {
+          return result.copyWithSubchannel(((SubchannelImpl) subchannel).delegate());
+        }
+        return result;
+      }
+    }
+
     @Override
     public Subchannel createSubchannel(CreateSubchannelArgs args) {
       syncContext.throwIfNotInThisSynchronizationContext();