From f64befd27fb08104fb5d1a7006d4f8ac08510f8f Mon Sep 17 00:00:00 2001
From: yang-g <yangg@google.com>
Date: Mon, 21 Mar 2016 14:04:10 -0700
Subject: [PATCH] Make a copy of ByteBuffer when writing

---
 include/grpc++/support/byte_buffer.h     |  4 ++--
 test/cpp/end2end/generic_end2end_test.cc |  5 +++++
 test/cpp/util/byte_buffer_test.cc        | 20 +++++++++++++++++++-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/include/grpc++/support/byte_buffer.h b/include/grpc++/support/byte_buffer.h
index 82591a88ef2..27307f8fcdb 100644
--- a/include/grpc++/support/byte_buffer.h
+++ b/include/grpc++/support/byte_buffer.h
@@ -99,8 +99,8 @@ class SerializationTraits<ByteBuffer, void> {
   }
   static Status Serialize(const ByteBuffer& source, grpc_byte_buffer** buffer,
                           bool* own_buffer) {
-    *buffer = source.buffer();
-    *own_buffer = false;
+    *buffer = grpc_byte_buffer_copy(source.buffer());
+    *own_buffer = true;
     return Status::OK;
   }
 };
diff --git a/test/cpp/end2end/generic_end2end_test.cc b/test/cpp/end2end/generic_end2end_test.cc
index 4e6d50ea80f..8dad1c2005f 100644
--- a/test/cpp/end2end/generic_end2end_test.cc
+++ b/test/cpp/end2end/generic_end2end_test.cc
@@ -135,6 +135,8 @@ class GenericEnd2endTest : public ::testing::Test {
       std::unique_ptr<ByteBuffer> send_buffer =
           SerializeToByteBuffer(&send_request);
       call->Write(*send_buffer, tag(2));
+      // Send ByteBuffer can be destroyed after calling Write.
+      send_buffer.reset();
       client_ok(2);
       call->WritesDone(tag(3));
       client_ok(3);
@@ -154,6 +156,7 @@ class GenericEnd2endTest : public ::testing::Test {
       send_response.set_message(recv_request.message());
       send_buffer = SerializeToByteBuffer(&send_response);
       stream.Write(*send_buffer, tag(6));
+      send_buffer.reset();
       server_ok(6);
 
       stream.Finish(Status::OK, tag(7));
@@ -223,6 +226,7 @@ TEST_F(GenericEnd2endTest, SimpleBidiStreaming) {
   std::unique_ptr<ByteBuffer> send_buffer =
       SerializeToByteBuffer(&send_request);
   cli_stream->Write(*send_buffer, tag(3));
+  send_buffer.reset();
   client_ok(3);
 
   ByteBuffer recv_buffer;
@@ -234,6 +238,7 @@ TEST_F(GenericEnd2endTest, SimpleBidiStreaming) {
   send_response.set_message(recv_request.message());
   send_buffer = SerializeToByteBuffer(&send_response);
   srv_stream.Write(*send_buffer, tag(5));
+  send_buffer.reset();
   server_ok(5);
 
   cli_stream->Read(&recv_buffer, tag(6));
diff --git a/test/cpp/util/byte_buffer_test.cc b/test/cpp/util/byte_buffer_test.cc
index f36c32cac5b..eb9dabcc2a7 100644
--- a/test/cpp/util/byte_buffer_test.cc
+++ b/test/cpp/util/byte_buffer_test.cc
@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2015, Google Inc.
+ * Copyright 2015-2016, Google Inc.
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -105,6 +105,24 @@ TEST_F(ByteBufferTest, Dump) {
   EXPECT_TRUE(SliceEqual(slices[1], world));
 }
 
+TEST_F(ByteBufferTest, SerializationMakesCopy) {
+  gpr_slice hello = gpr_slice_from_copied_string(kContent1);
+  gpr_slice world = gpr_slice_from_copied_string(kContent2);
+  std::vector<Slice> slices;
+  slices.push_back(Slice(hello, Slice::STEAL_REF));
+  slices.push_back(Slice(world, Slice::STEAL_REF));
+  grpc_byte_buffer* send_buffer = nullptr;
+  bool owned = false;
+  ByteBuffer buffer(&slices[0], 2);
+  slices.clear();
+  auto status = SerializationTraits<ByteBuffer, void>::Serialize(
+      buffer, &send_buffer, &owned);
+  EXPECT_TRUE(status.ok());
+  EXPECT_TRUE(owned);
+  EXPECT_TRUE(send_buffer != nullptr);
+  grpc_byte_buffer_destroy(send_buffer);
+}
+
 }  // namespace
 }  // namespace grpc