Use call context to propagare LR costs

2 files changed, 20 insertions, 3 deletions

diff --git a/include/grpc++/impl/codegen/server_context.h b/include/grpc++/impl/codegen/server_context.h
index 8c7fe0809e..43bbff6ce9 100644
--- a/include/grpc++/impl/codegen/server_context.h
+++ b/include/grpc++/impl/codegen/server_context.h
@@ -36,6 +36,10 @@
 
 #include <map>
 #include <memory>
+#include <vector>
+
+#include <grpc/impl/codegen/compression_types.h>
+#include <grpc/load_reporting.h>
 
 #include <grpc++/impl/codegen/config.h>
 #include <grpc++/impl/codegen/create_auth_context.h>
@@ -43,14 +47,12 @@
 #include <grpc++/impl/codegen/security/auth_context.h>
 #include <grpc++/impl/codegen/string_ref.h>
 #include <grpc++/impl/codegen/time.h>
-#include <grpc/impl/codegen/compression_types.h>
 
 struct grpc_metadata;
 struct grpc_call;
 struct census_context;
 
 namespace grpc {
-
 class ClientContext;
 template <class W, class R>
 class ServerAsyncReader;
@@ -143,6 +145,9 @@ class ServerContext {
   }
   void set_compression_algorithm(grpc_compression_algorithm algorithm);
 
+  // Set the load reporting costs in \a cost_data for the call.
+  void SetLoadReportingCosts(const std::vector<grpc::string>& cost_data);
+
   std::shared_ptr<const AuthContext> auth_context() const {
     if (auth_context_.get() == nullptr) {
       auth_context_ = CreateAuthContext(call_);
diff --git a/include/grpc++/test/server_context_test_spouse.h b/include/grpc++/test/server_context_test_spouse.h
index b2482854b5..4cb84e77c1 100644
--- a/include/grpc++/test/server_context_test_spouse.h
+++ b/include/grpc++/test/server_context_test_spouse.h
@@ -48,10 +48,22 @@ class ServerContextTestSpouse {
 
   // Inject client metadata to the ServerContext for the test. The test spouse
   // must be alive when ServerContext::client_metadata is called.
-  void AddClientMetadata(const grpc::string& key, const grpc::string& value);
+  void AddClientMetadata(const grpc::string& key, const grpc::string& value) {
+    client_metadata_storage_.insert(
+        std::pair<grpc::string, grpc::string>(key, value));
+    ctx_->client_metadata_.map()->clear();
+    for (auto iter = client_metadata_storage_.begin();
+         iter != client_metadata_storage_.end(); ++iter) {
+      ctx_->client_metadata_.map()->insert(
+          std::pair<grpc::string_ref, grpc::string_ref>(iter->first.c_str(),
+                                                        iter->second.c_str()));
+    }
+  }
+
   std::multimap<grpc::string, grpc::string> GetInitialMetadata() const {
     return ctx_->initial_metadata_;
   }
+
   std::multimap<grpc::string, grpc::string> GetTrailingMetadata() const {
     return ctx_->trailing_metadata_;
   }