From c67f7183b20f23f09eccb638c2d5ac488307cfbb Mon Sep 17 00:00:00 2001
From: Michael Lumish <mlumish@google.com>
Date: Fri, 22 Aug 2025 15:25:58 -0700
Subject: [PATCH] grpc-js: Properly unwrap subchannel in WRR picker

---
 .../src/load-balancer-weighted-round-robin.ts | 20 +++++++++++++------
 packages/grpc-js/src/orca.ts                  |  4 ++++
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/packages/grpc-js/src/load-balancer-weighted-round-robin.ts b/packages/grpc-js/src/load-balancer-weighted-round-robin.ts
index b043e186..cdeabc36 100644
--- a/packages/grpc-js/src/load-balancer-weighted-round-robin.ts
+++ b/packages/grpc-js/src/load-balancer-weighted-round-robin.ts
@@ -25,7 +25,7 @@ import { ChannelControlHelper, createChildChannelControlHelper, LoadBalancer, re
 import { LeafLoadBalancer } from './load-balancer-pick-first';
 import * as logging from './logging';
 import { createMetricsReader, MetricsListener, OrcaOobMetricsSubchannelWrapper } from './orca';
-import { PickArgs, Picker, PickResult, QueuePicker, UnavailablePicker } from './picker';
+import { PickArgs, Picker, PickResult, PickResultType, QueuePicker, UnavailablePicker } from './picker';
 import { PriorityQueue } from './priority-queue';
 import { Endpoint, endpointToString } from './subchannel-address';
 
@@ -209,11 +209,19 @@ class WeightedRoundRobinPicker implements Picker {
       deadline: entry.deadline + entry.period
     })
     const childPick = entry.picker.pick(pickArgs);
-    if (this.metricsHandler) {
-      return {
-        ...childPick,
-        onCallEnded: createMetricsReader(loadReport => this.metricsHandler!(loadReport, entry.endpointName), childPick.onCallEnded)
-      };
+    if (childPick.pickResultType === PickResultType.COMPLETE) {
+      if (this.metricsHandler) {
+        return {
+          ...childPick,
+          onCallEnded: createMetricsReader(loadReport => this.metricsHandler!(loadReport, entry.endpointName), childPick.onCallEnded)
+        };
+      } else {
+        const subchannelWrapper = childPick.subchannel as OrcaOobMetricsSubchannelWrapper;
+        return {
+          ...childPick,
+          subchannel: subchannelWrapper.getWrappedSubchannel()
+        }
+      }
     } else {
       return childPick;
     }
diff --git a/packages/grpc-js/src/orca.ts b/packages/grpc-js/src/orca.ts
index 16372fb0..23490733 100644
--- a/packages/grpc-js/src/orca.ts
+++ b/packages/grpc-js/src/orca.ts
@@ -338,6 +338,10 @@ export class OrcaOobMetricsSubchannelWrapper extends BaseSubchannelWrapper {
     super(child);
     this.addDataWatcher(new OobMetricsDataWatcher(metricsListener, intervalMs));
   }
+
+  getWrappedSubchannel(): SubchannelInterface {
+    return this.child;
+  }
 }
 
 function createOobMetricsDataProducer(subchannel: Subchannel) {