From 424c9bfe70601cccb0698341b95fa6b6a9d29283 Mon Sep 17 00:00:00 2001
From: Michael Lumish <mlumish@google.com>
Date: Mon, 27 Apr 2020 10:24:15 -0700
Subject: [PATCH] Add onCallStarted field to Pick object

---
 packages/grpc-js/src/channel.ts                   |  3 +++
 packages/grpc-js/src/load-balancer-pick-first.ts  |  2 ++
 packages/grpc-js/src/load-balancer-round-robin.ts |  2 ++
 packages/grpc-js/src/picker.ts                    | 10 ++++++++--
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/packages/grpc-js/src/channel.ts b/packages/grpc-js/src/channel.ts
index f4df2fce..b240cc43 100644
--- a/packages/grpc-js/src/channel.ts
+++ b/packages/grpc-js/src/channel.ts
@@ -303,6 +303,9 @@ export class ChannelImplementation implements Channel {
                       callStream,
                       pickResult.extraFilterFactory ?? undefined
                     );
+                    /* If we reach this point, the call stream has started
+                     * successfully */
+                    pickResult.onCallStarted?.();
                   } catch (error) {
                     if (
                       (error as NodeJS.ErrnoException).code ===
diff --git a/packages/grpc-js/src/load-balancer-pick-first.ts b/packages/grpc-js/src/load-balancer-pick-first.ts
index 6b9756ff..6289123d 100644
--- a/packages/grpc-js/src/load-balancer-pick-first.ts
+++ b/packages/grpc-js/src/load-balancer-pick-first.ts
@@ -65,6 +65,8 @@ class PickFirstPicker implements Picker {
       pickResultType: PickResultType.COMPLETE,
       subchannel: this.subchannel,
       status: null,
+      extraFilterFactory: null,
+      onCallStarted: null
     };
   }
 }
diff --git a/packages/grpc-js/src/load-balancer-round-robin.ts b/packages/grpc-js/src/load-balancer-round-robin.ts
index 93c64610..51244341 100644
--- a/packages/grpc-js/src/load-balancer-round-robin.ts
+++ b/packages/grpc-js/src/load-balancer-round-robin.ts
@@ -60,6 +60,8 @@ class RoundRobinPicker implements Picker {
       pickResultType: PickResultType.COMPLETE,
       subchannel: pickedSubchannel,
       status: null,
+      extraFilterFactory: null,
+      onCallStarted: null,
     };
   }
 
diff --git a/packages/grpc-js/src/picker.ts b/packages/grpc-js/src/picker.ts
index 9470f122..306ac3ed 100644
--- a/packages/grpc-js/src/picker.ts
+++ b/packages/grpc-js/src/picker.ts
@@ -47,6 +47,7 @@ export interface PickResult {
    * reasons filters from this factory will not see sendMetadata events.
    */
   extraFilterFactory: FilterFactory<Filter> | null;
+  onCallStarted: (() => void) | null;
 }
 
 export interface CompletePickResult extends PickResult {
@@ -54,6 +55,7 @@ export interface CompletePickResult extends PickResult {
   subchannel: Subchannel | null;
   status: null;
   extraFilterFactory: FilterFactory<Filter> | null;
+  onCallStarted: (() => void) | null;
 }
 
 export interface QueuePickResult extends PickResult {
@@ -61,6 +63,7 @@ export interface QueuePickResult extends PickResult {
   subchannel: null;
   status: null;
   extraFilterFactory: null;
+  onCallStarted: null;
 }
 
 export interface TransientFailurePickResult extends PickResult {
@@ -68,6 +71,7 @@ export interface TransientFailurePickResult extends PickResult {
   subchannel: null;
   status: StatusObject;
   extraFilterFactory: null;
+  onCallStarted: null;
 }
 
 export interface PickArgs {
@@ -105,7 +109,8 @@ export class UnavailablePicker implements Picker {
       pickResultType: PickResultType.TRANSIENT_FAILURE,
       subchannel: null,
       status: this.status,
-      extraFilterFactory: null
+      extraFilterFactory: null,
+      onCallStarted: null
     };
   }
 }
@@ -133,7 +138,8 @@ export class QueuePicker {
       pickResultType: PickResultType.QUEUE,
       subchannel: null,
       status: null,
-      extraFilterFactory: null
+      extraFilterFactory: null,
+      onCallStarted: null
     };
   }
 }