Start resolving memory issues in C++ metadata

5 files changed, 97 insertions, 41 deletions

diff --git a/include/grpc++/impl/codegen/call.h b/include/grpc++/impl/codegen/call.h
index 1d51179257..9165f4fd6c 100644
--- a/include/grpc++/impl/codegen/call.h
+++ b/include/grpc++/impl/codegen/call.h
@@ -63,19 +63,6 @@ class CallHook;
 class CompletionQueue;
 extern CoreCodegenInterface* g_core_codegen_interface;
 
-inline void FillMetadataMap(
-    grpc_metadata_array* arr,
-    std::multimap<grpc::string_ref, grpc::string_ref>* metadata) {
-  for (size_t i = 0; i < arr->count; i++) {
-    // TODO(yangg) handle duplicates?
-    metadata->insert(std::pair<grpc::string_ref, grpc::string_ref>(
-        StringRefFromSlice(arr->metadata[i].key),
-        StringRefFromSlice(arr->metadata[i].value)));
-  }
-  g_core_codegen_interface->grpc_metadata_array_destroy(arr);
-  g_core_codegen_interface->grpc_metadata_array_init(arr);
-}
-
 // TODO(yangg) if the map is changed before we send, the pointers will be a
 // mess. Make sure it does not happen.
 inline grpc_metadata* FillMetadataArray(
@@ -474,32 +461,30 @@ class CallOpServerSendStatus {
 
 class CallOpRecvInitialMetadata {
  public:
-  CallOpRecvInitialMetadata() : recv_initial_metadata_(nullptr) {}
+  CallOpRecvInitialMetadata() : metadata_map_(nullptr) {}
 
   void RecvInitialMetadata(ClientContext* context) {
     context->initial_metadata_received_ = true;
-    recv_initial_metadata_ = &context->recv_initial_metadata_;
+    metadata_map_ = &context->recv_initial_metadata_;
   }
 
  protected:
   void AddOp(grpc_op* ops, size_t* nops) {
-    if (!recv_initial_metadata_) return;
-    memset(&recv_initial_metadata_arr_, 0, sizeof(recv_initial_metadata_arr_));
+    if (metadata_map_ == nullptr) return;
     grpc_op* op = &ops[(*nops)++];
     op->op = GRPC_OP_RECV_INITIAL_METADATA;
-    op->data.recv_initial_metadata = &recv_initial_metadata_arr_;
+    op->data.recv_initial_metadata = metadata_map_->arr();
     op->flags = 0;
     op->reserved = NULL;
   }
   void FinishOp(bool* status, int max_receive_message_size) {
-    if (recv_initial_metadata_ == nullptr) return;
-    FillMetadataMap(&recv_initial_metadata_arr_, recv_initial_metadata_);
-    recv_initial_metadata_ = nullptr;
+    if (metadata_map_ == nullptr) return;
+    metadata_map_->FillMap();
+    metadata_map_ = nullptr;
   }
 
  private:
-  std::multimap<grpc::string_ref, grpc::string_ref>* recv_initial_metadata_;
-  grpc_metadata_array recv_initial_metadata_arr_;
+  MetadataMap* metadata_map_;
 };
 
 class CallOpClientRecvStatus {
@@ -507,19 +492,16 @@ class CallOpClientRecvStatus {
   CallOpClientRecvStatus() : recv_status_(nullptr) {}
 
   void ClientRecvStatus(ClientContext* context, Status* status) {
-    recv_trailing_metadata_ = &context->trailing_metadata_;
+    metadata_map_ = &context->trailing_metadata_;
     recv_status_ = status;
   }
 
  protected:
   void AddOp(grpc_op* ops, size_t* nops) {
     if (recv_status_ == nullptr) return;
-    memset(&recv_trailing_metadata_arr_, 0,
-           sizeof(recv_trailing_metadata_arr_));
     grpc_op* op = &ops[(*nops)++];
     op->op = GRPC_OP_RECV_STATUS_ON_CLIENT;
-    op->data.recv_status_on_client.trailing_metadata =
-        &recv_trailing_metadata_arr_;
+    op->data.recv_status_on_client.trailing_metadata = metadata_map_->arr();
     op->data.recv_status_on_client.status = &status_code_;
     op->data.recv_status_on_client.status_details = &status_details_;
     op->flags = 0;
@@ -528,7 +510,7 @@ class CallOpClientRecvStatus {
 
   void FinishOp(bool* status, int max_receive_message_size) {
     if (recv_status_ == nullptr) return;
-    FillMetadataMap(&recv_trailing_metadata_arr_, recv_trailing_metadata_);
+    metadata_map_->FillMap();
     *recv_status_ = Status(static_cast<StatusCode>(status_code_),
                            grpc::string(GRPC_SLICE_START_PTR(status_details_),
                                         GRPC_SLICE_END_PTR(status_details_)));
@@ -537,9 +519,8 @@ class CallOpClientRecvStatus {
   }
 
  private:
-  std::multimap<grpc::string_ref, grpc::string_ref>* recv_trailing_metadata_;
+  MetadataMap* metadata_map_;
   Status* recv_status_;
-  grpc_metadata_array recv_trailing_metadata_arr_;
   grpc_status_code status_code_;
   grpc_slice status_details_;
 };
diff --git a/include/grpc++/impl/codegen/client_context.h b/include/grpc++/impl/codegen/client_context.h
index 777b2f8847..b91c7f65d4 100644
--- a/include/grpc++/impl/codegen/client_context.h
+++ b/include/grpc++/impl/codegen/client_context.h
@@ -57,7 +57,9 @@
 #include <grpc++/impl/codegen/config.h>
 #include <grpc++/impl/codegen/core_codegen_interface.h>
 #include <grpc++/impl/codegen/create_auth_context.h>
+#include <grpc++/impl/codegen/metadata_map.h>
 #include <grpc++/impl/codegen/security/auth_context.h>
+#include <grpc++/impl/codegen/slice.h>
 #include <grpc++/impl/codegen/status.h>
 #include <grpc++/impl/codegen/string_ref.h>
 #include <grpc++/impl/codegen/time.h>
@@ -193,7 +195,7 @@ class ClientContext {
   const std::multimap<grpc::string_ref, grpc::string_ref>&
   GetServerInitialMetadata() const {
     GPR_CODEGEN_ASSERT(initial_metadata_received_);
-    return recv_initial_metadata_;
+    return *recv_initial_metadata_.map();
   }
 
   /// Return a collection of trailing metadata key-value pairs. Note that keys
@@ -205,7 +207,7 @@ class ClientContext {
   const std::multimap<grpc::string_ref, grpc::string_ref>&
   GetServerTrailingMetadata() const {
     // TODO(yangg) check finished
-    return trailing_metadata_;
+    return *trailing_metadata_.map();
   }
 
   /// Set the deadline for the client call.
@@ -375,8 +377,8 @@ class ClientContext {
   mutable std::shared_ptr<const AuthContext> auth_context_;
   struct census_context* census_context_;
   std::multimap<grpc::string, grpc::string> send_initial_metadata_;
-  std::multimap<grpc::string_ref, grpc::string_ref> recv_initial_metadata_;
-  std::multimap<grpc::string_ref, grpc::string_ref> trailing_metadata_;
+  MetadataMap recv_initial_metadata_;
+  MetadataMap trailing_metadata_;
 
   grpc_call* propagate_from_call_;
   PropagationOptions propagation_options_;
diff --git a/include/grpc++/impl/codegen/metadata_map.h b/include/grpc++/impl/codegen/metadata_map.h
new file mode 100644
index 0000000000..fc2ed7d9e1
--- /dev/null
+++ b/include/grpc++/impl/codegen/metadata_map.h
@@ -0,0 +1,71 @@
+/*
+*
+* Copyright 2015, Google Inc.
+* All rights reserved.
+*
+* Redistribution and use in source and binary forms, with or without
+* modification, are permitted provided that the following conditions are
+* met:
+*
+*     * Redistributions of source code must retain the above copyright
+* notice, this list of conditions and the following disclaimer.
+*     * Redistributions in binary form must reproduce the above
+* copyright notice, this list of conditions and the following disclaimer
+* in the documentation and/or other materials provided with the
+* distribution.
+*     * Neither the name of Google Inc. nor the names of its
+* contributors may be used to endorse or promote products derived from
+* this software without specific prior written permission.
+*
+* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*
+*/
+
+#ifndef META_H
+#define META_H
+
+#include <grpc++/impl/codegen/slice.h>
+
+namespace grpc {
+
+class MetadataMap {
+ public:
+  MetadataMap() { memset(&arr_, 0, sizeof(arr_)); }
+
+  ~MetadataMap() {
+    g_core_codegen_interface->grpc_metadata_array_destroy(&arr_);
+  }
+
+  void FillMap() {
+    for (size_t i = 0; i < arr_.count; i++) {
+      // TODO(yangg) handle duplicates?
+      map_.insert(std::pair<grpc::string_ref, grpc::string_ref>(
+          StringRefFromSlice(&arr_.metadata[i].key),
+          StringRefFromSlice(&arr_.metadata[i].value)));
+    }
+  }
+
+  std::multimap<grpc::string_ref, grpc::string_ref> *map() { return &map_; }
+  const std::multimap<grpc::string_ref, grpc::string_ref> *map() const {
+    return &map_;
+  }
+  grpc_metadata_array *arr() { return &arr_; }
+
+ private:
+  grpc_metadata_array arr_;
+  std::multimap<grpc::string_ref, grpc::string_ref> map_;
+};
+
+}  // namespace grpc
+
+#endif
diff --git a/include/grpc++/impl/codegen/server_context.h b/include/grpc++/impl/codegen/server_context.h
index dd30576379..8c7fe0809e 100644
--- a/include/grpc++/impl/codegen/server_context.h
+++ b/include/grpc++/impl/codegen/server_context.h
@@ -39,6 +39,7 @@
 
 #include <grpc++/impl/codegen/config.h>
 #include <grpc++/impl/codegen/create_auth_context.h>
+#include <grpc++/impl/codegen/metadata_map.h>
 #include <grpc++/impl/codegen/security/auth_context.h>
 #include <grpc++/impl/codegen/string_ref.h>
 #include <grpc++/impl/codegen/time.h>
@@ -123,7 +124,7 @@ class ServerContext {
 
   const std::multimap<grpc::string_ref, grpc::string_ref>& client_metadata()
       const {
-    return client_metadata_;
+    return *client_metadata_.map();
   }
 
   grpc_compression_level compression_level() const {
@@ -223,7 +224,7 @@ class ServerContext {
   CompletionQueue* cq_;
   bool sent_initial_metadata_;
   mutable std::shared_ptr<const AuthContext> auth_context_;
-  std::multimap<grpc::string_ref, grpc::string_ref> client_metadata_;
+  MetadataMap client_metadata_;
   std::multimap<grpc::string, grpc::string> initial_metadata_;
   std::multimap<grpc::string, grpc::string> trailing_metadata_;
 
diff --git a/include/grpc++/impl/codegen/slice.h b/include/grpc++/impl/codegen/slice.h
index e79754f943..04b2f9af01 100644
--- a/include/grpc++/impl/codegen/slice.h
+++ b/include/grpc++/impl/codegen/slice.h
@@ -39,9 +39,10 @@
 
 namespace grpc {
 
-inline grpc::string_ref StringRefFromSlice(grpc_slice slice) {
-  return grpc::string_ref(reinterpret_cast<char*>(GRPC_SLICE_START_PTR(slice)),
-                          GRPC_SLICE_LENGTH(slice));
+inline grpc::string_ref StringRefFromSlice(const grpc_slice* slice) {
+  return grpc::string_ref(
+      reinterpret_cast<const char*>(GRPC_SLICE_START_PTR(*slice)),
+      GRPC_SLICE_LENGTH(*slice));
 }
 
 inline grpc::string StringFromCopiedSlice(grpc_slice slice) {
@@ -61,4 +62,4 @@ inline grpc_slice SliceFromCopiedString(const grpc::string& str) {
 
 }  // namespace grpc
 
-#endif
+#endif  // GRPCXX_IMPL_CODEGEN_SLICE_H