grpc-js: Properly unwrap subchannel in WRR picker

Author: murgatroid99

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;
     }

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) {