From 2378ddcd51fb7a9529fe4bfce6f8802c6f1878f4 Mon Sep 17 00:00:00 2001
From: Yash Tibrewal <yashkt@google.com>
Date: Tue, 9 Apr 2019 22:23:10 -0700
Subject: [PATCH] Add service config to call context

---
 .../ext/filters/client_channel/client_channel.cc | 16 +++++++++++++---
 .../client_channel/resolver_result_parsing.h     |  1 +
 src/core/lib/channel/context.h                   |  2 ++
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/core/ext/filters/client_channel/client_channel.cc b/src/core/ext/filters/client_channel/client_channel.cc
index cd552739732..049c8d0a13f 100644
--- a/src/core/ext/filters/client_channel/client_channel.cc
+++ b/src/core/ext/filters/client_channel/client_channel.cc
@@ -170,6 +170,7 @@ struct client_channel_channel_data {
   bool received_service_config_data;
   grpc_core::RefCountedPtr<ServerRetryThrottleData> retry_throttle_data;
   grpc_core::RefCountedPtr<ClientChannelMethodParamsTable> method_params_table;
+  grpc_core::RefCountedPtr<grpc_core::ServiceConfig> service_config;
 
   //
   // Fields used in the control plane.  Protected by combiner.
@@ -276,10 +277,12 @@ class ServiceConfigSetter {
   ServiceConfigSetter(
       channel_data* chand,
       RefCountedPtr<ServerRetryThrottleData> retry_throttle_data,
-      RefCountedPtr<ClientChannelMethodParamsTable> method_params_table)
+      RefCountedPtr<ClientChannelMethodParamsTable> method_params_table,
+      RefCountedPtr<ServiceConfig> service_config)
       : chand_(chand),
         retry_throttle_data_(std::move(retry_throttle_data)),
-        method_params_table_(std::move(method_params_table)) {
+        method_params_table_(std::move(method_params_table)),
+        service_config_(std::move(service_config)) {
     GRPC_CHANNEL_STACK_REF(chand->owning_stack, "ServiceConfigSetter");
     GRPC_CLOSURE_INIT(&closure_, SetServiceConfigData, this,
                       grpc_combiner_scheduler(chand->data_plane_combiner));
@@ -294,6 +297,7 @@ class ServiceConfigSetter {
     chand->received_service_config_data = true;
     chand->retry_throttle_data = std::move(self->retry_throttle_data_);
     chand->method_params_table = std::move(self->method_params_table_);
+    chand->service_config = std::move(self->service_config_);
     // Apply service config to queued picks.
     for (QueuedPick* pick = chand->queued_picks; pick != nullptr;
          pick = pick->next) {
@@ -307,6 +311,7 @@ class ServiceConfigSetter {
   channel_data* chand_;
   RefCountedPtr<ServerRetryThrottleData> retry_throttle_data_;
   RefCountedPtr<ClientChannelMethodParamsTable> method_params_table_;
+  RefCountedPtr<ServiceConfig> service_config_;
   grpc_closure closure_;
 };
 
@@ -505,7 +510,7 @@ static bool process_resolver_result_locked(
   // plane combiner.  Destroys itself when done.
   grpc_core::New<grpc_core::ServiceConfigSetter>(
       chand, resolver_result.retry_throttle_data(),
-      resolver_result.method_params_table());
+      resolver_result.method_params_table(), resolver_result.service_config());
   // Swap out the data used by cc_get_channel_info().
   gpr_mu_lock(&chand->info_mu);
   chand->info_lb_policy_name = resolver_result.lb_policy_name();
@@ -970,6 +975,7 @@ struct call_data {
 
   grpc_core::RefCountedPtr<ServerRetryThrottleData> retry_throttle_data;
   grpc_core::RefCountedPtr<ClientChannelMethodParams> method_params;
+  grpc_core::RefCountedPtr<grpc_core::ServiceConfig> service_config;
 
   grpc_core::RefCountedPtr<grpc_core::SubchannelCall> subchannel_call;
 
@@ -2813,6 +2819,10 @@ static void apply_service_config_to_call_locked(grpc_call_element* elem) {
     gpr_log(GPR_INFO, "chand=%p calld=%p: applying service config to call",
             chand, calld);
   }
+  if(chand->service_config != nullptr) {
+    calld->service_config = chand->service_config;
+    calld->call_context[GRPC_SERVICE_CONFIG].value = &calld->service_config;
+  }
   if (chand->retry_throttle_data != nullptr) {
     calld->retry_throttle_data = chand->retry_throttle_data->Ref();
   }
diff --git a/src/core/ext/filters/client_channel/resolver_result_parsing.h b/src/core/ext/filters/client_channel/resolver_result_parsing.h
index 1a46278f38b..ad40c4991fe 100644
--- a/src/core/ext/filters/client_channel/resolver_result_parsing.h
+++ b/src/core/ext/filters/client_channel/resolver_result_parsing.h
@@ -64,6 +64,7 @@ class ProcessedResolverResult {
   RefCountedPtr<LoadBalancingPolicy::Config> lb_policy_config() {
     return std::move(lb_policy_config_);
   }
+  RefCountedPtr<ServiceConfig> service_config() { return service_config_; }
 
  private:
   // Finds the service config; extracts LB config and (maybe) retry throttle
diff --git a/src/core/lib/channel/context.h b/src/core/lib/channel/context.h
index 81b84f1ca05..722c22c4623 100644
--- a/src/core/lib/channel/context.h
+++ b/src/core/lib/channel/context.h
@@ -35,6 +35,8 @@ typedef enum {
   /// Reserved for traffic_class_context.
   GRPC_CONTEXT_TRAFFIC,
 
+  GRPC_SERVICE_CONFIG,
+
   GRPC_CONTEXT_COUNT
 } grpc_context_index;