grpc: utilities for inter-converting grpc::ByteBuffer and Buffer::Instance.

source/common/grpc/BUILD

@@ -28,6 +28,7 @@ envoy_cc_library(
     hdrs = ["async_client_manager_impl.h"],
     deps = [
         ":async_client_lib",
+        ":common_lib",
         "//include/envoy/grpc:async_client_manager_interface",
         "//include/envoy/singleton:manager_interface",
         "//include/envoy/thread_local:thread_local_interface",
@@ -61,7 +62,10 @@ envoy_cc_library(
     name = "common_lib",
     srcs = ["common.cc"],
     hdrs = ["common.h"],
-    external_deps = ["abseil_optional"],
+    external_deps = [
+        "abseil_optional",
+        "grpc",
+    ],
     deps = [
         "//include/envoy/http:header_map_interface",
         "//include/envoy/http:message_interface",
@@ -91,6 +95,8 @@ envoy_cc_library(
         "grpc",
     ],
     deps = [
+        ":async_client_lib",
+        ":common_lib",
         ":google_grpc_creds_lib",
         "//include/envoy/api:api_interface",
         "//include/envoy/grpc:google_grpc_creds_interface",

source/common/grpc/async_client_impl.cc

@@ -177,7 +177,7 @@ void AsyncStreamImpl::onReset() {
 }
 
 void AsyncStreamImpl::sendMessage(const Protobuf::Message& request, bool end_stream) {
-  stream_->sendData(*Common::serializeBody(request), end_stream);
+  stream_->sendData(*Common::serializeToGrpcFrame(request), end_stream);
 }
 
 void AsyncStreamImpl::closeStream() {

source/common/grpc/common.cc

@@ -2,6 +2,7 @@
 
 #include <arpa/inet.h>
 
+#include <atomic>
 #include <cstdint>
 #include <cstring>
 #include <string>
@@ -12,6 +13,7 @@
 #include "common/common/enum_to_int.h"
 #include "common/common/fmt.h"
 #include "common/common/macros.h"
+#include "common/common/stack_array.h"
 #include "common/common/utility.h"
 #include "common/http/headers.h"
 #include "common/http/message_impl.h"
@@ -111,7 +113,7 @@ bool Common::resolveServiceAndMethod(const Http::HeaderEntry* path, std::string*
   return true;
 }
 
-Buffer::InstancePtr Common::serializeBody(const Protobuf::Message& message) {
+Buffer::InstancePtr Common::serializeToGrpcFrame(const Protobuf::Message& message) {
   // http://www.grpc.io/docs/guides/wire.html
   // Reserve enough space for the entire message and the 5 byte header.
   Buffer::InstancePtr body(new Buffer::OwnedImpl());
@@ -133,6 +135,21 @@ Buffer::InstancePtr Common::serializeBody(const Protobuf::Message& message) {
   return body;
 }
 
+Buffer::InstancePtr Common::serializeMessage(const Protobuf::Message& message) {
+  Buffer::InstancePtr body(new Buffer::OwnedImpl());
+  const uint32_t size = message.ByteSize();
+  Buffer::RawSlice iovec;
+  body->reserve(size, &iovec, 1);
+  ASSERT(iovec.len_ >= size);
+  iovec.len_ = size;
+  uint8_t* current = reinterpret_cast<uint8_t*>(iovec.mem_);
+  Protobuf::io::ArrayOutputStream stream(current, size, -1);
+  Protobuf::io::CodedOutputStream codec_stream(&stream);
+  message.SerializeWithCachedSizes(&codec_stream);
+  body->commit(&iovec, 1);
+  return body;
+}
+
 std::chrono::milliseconds Common::getGrpcTimeout(Http::HeaderMap& request_headers) {
   std::chrono::milliseconds timeout(0);
   Http::HeaderEntry* header_grpc_timeout_entry = request_headers.GrpcTimeout();
@@ -269,5 +286,98 @@ std::string Common::typeUrl(const std::string& qualified_name) {
   return typeUrlPrefix() + "/" + qualified_name;
 }
 
+struct BufferInstanceContainer {
+  BufferInstanceContainer(int ref_count, Buffer::InstancePtr buffer)
+      : ref_count_(ref_count), buffer_(std::move(buffer)) {}
+  std::atomic<int> ref_count_;
+  Buffer::InstancePtr buffer_;
+};
+
+static void derefBufferInstanceContainer(void* container_ptr) {
+  auto container = reinterpret_cast<BufferInstanceContainer*>(container_ptr);
+  container->ref_count_--;
+  if (container->ref_count_ <= 0) {
+    delete container;
+  }
+}
+
+grpc::ByteBuffer Common::makeByteBuffer(Buffer::InstancePtr bufferInstance) {
+  if (!bufferInstance) {
+    return {};
+  }
+  Buffer::RawSlice oneRawSlice;
+  // NB: we need to pass in >= 1 in order to get the real "n" (see Buffer::Instance for details).
+  int nSlices = bufferInstance->getRawSlices(&oneRawSlice, 1);
+  if (nSlices <= 0) {
+    return {};
+  }
+  auto container = new BufferInstanceContainer{nSlices, std::move(bufferInstance)};
+  if (nSlices == 1) {
+    grpc::Slice oneSlice(oneRawSlice.mem_, oneRawSlice.len_, &derefBufferInstanceContainer,
+                         container);
+    return {&oneSlice, 1};
+  }
+  STACK_ARRAY(manyRawSlices, Buffer::RawSlice, nSlices);
+  container->buffer_->getRawSlices(manyRawSlices.begin(), nSlices);
+  std::vector<grpc::Slice> slices;
+  slices.reserve(nSlices);
+  for (int i = 0; i < nSlices; i++) {
+    slices.emplace_back(manyRawSlices[i].mem_, manyRawSlices[i].len_, &derefBufferInstanceContainer,
+                        container);
+  }
+  return {&slices[0], slices.size()};
+}
+
+struct ByteBufferContainer {
+  ByteBufferContainer(int ref_count) : ref_count_(ref_count) {}
+  ~ByteBufferContainer() { ::free(fragments); }
+  std::atomic<int> ref_count_;
+  Buffer::BufferFragmentImpl* fragments = nullptr;
+  std::vector<grpc::Slice> slices_;
+};
+
+Buffer::InstancePtr Common::makeBufferInstance(const grpc::ByteBuffer& byteBuffer) {
+  auto buffer = std::make_unique<Buffer::OwnedImpl>();
+  if (byteBuffer.Length() == 0) {
+    return buffer;
+  }
+  // NB: ByteBuffer::Dump moves the data out of the ByteBuffer so we need to ensure that the
+  // lifetime of the Slice(s) exceeds our Buffer::Instance.
+  std::vector<grpc::Slice> slices;
+  byteBuffer.Dump(&slices);
+  if (slices.size() == 0) {
+    return buffer;
+  }
+  auto container = new ByteBufferContainer(static_cast<int>(slices.size()));
+  std::function<void(const void*, size_t, const Buffer::BufferFragmentImpl*)> releaser =
+      [container](const void*, size_t, const Buffer::BufferFragmentImpl*) {
+        container->ref_count_--;
+        if (container->ref_count_ <= 0) {
+          delete container;
+        }
+      };
+  // NB: addBufferFragment takes a pointer alias to the BufferFragmentImpl which is passed in so we
+  // need to ensure that the lifetime of those objects exceeds that of the Buffer::Instance.
+  container->fragments = static_cast<Buffer::BufferFragmentImpl*>(
+      ::malloc(sizeof(Buffer::BufferFragmentImpl) * slices.size()));
+  for (size_t i = 0; i < slices.size(); i++) {
+    new (&container->fragments[i])
+        Buffer::BufferFragmentImpl(slices[i].begin(), slices[i].size(), releaser);
+  }
+  for (size_t i = 0; i < slices.size(); i++) {
+    buffer->addBufferFragment(container->fragments[i]);
+  }
+  container->slices_ = std::move(slices);
+  return buffer;
+}
+
+void Common::PrependGrpcFrameHeader(Buffer::Instance* buffer) {
+  char header[5];
+  header[0] = 0; // flags
+  const uint32_t nsize = htonl(buffer->length());
+  std::memcpy(&header[1], reinterpret_cast<const void*>(&nsize), sizeof(uint32_t));
+  buffer->prepend(absl::string_view(header, 5));
+}
+
 } // namespace Grpc
 } // namespace Envoy

source/common/grpc/common.h

@@ -13,6 +13,7 @@
 #include "common/protobuf/protobuf.h"
 
 #include "absl/types/optional.h"
+#include "grpcpp/grpcpp.h"
 
 namespace Envoy {
 namespace Grpc {
@@ -119,9 +120,14 @@ class Common {
                                       std::string* method);
 
   /**
-   * Serialize protobuf message.
+   * Serialize protobuf message. With grpc header.
    */
-  static Buffer::InstancePtr serializeBody(const Protobuf::Message& message);
+  static Buffer::InstancePtr serializeToGrpcFrame(const Protobuf::Message& message);
+
+  /**
+   * Serialize protobuf message. Without grpc header.
+   */
+  static Buffer::InstancePtr serializeMessage(const Protobuf::Message& message);
 
   /**
    * Prepare headers for protobuf service.
@@ -148,6 +154,27 @@ class Common {
    */
   static std::string typeUrl(const std::string& qualified_name);
 
+  /**
+   * Build grpc::ByteBuffer which aliases the data in a Buffer::InstancePtr.
+   * @param bufferInstance source data container.
+   * @return byteBuffer target container aliased to the data in Buffer::Instance and owning the
+   * Buffer::Instance.
+   */
+  static grpc::ByteBuffer makeByteBuffer(Buffer::InstancePtr bufferInstance);
+
+  /**
+   * Build Buffer::Instance which aliases the data in a grpc::ByteBuffer.
+   * @param byteBuffer source data container.
+   * @param Buffer::InstancePtr target container aliased to the data in grpc::ByteBuffer.
+   */
+  static Buffer::InstancePtr makeBufferInstance(const grpc::ByteBuffer& byteBuffer);
+
+  /**
+   * Prepend a gRPC frame header to a Buffer::Instance containing a single gRPC frame.
+   * @param bufferInstance containing the frame data which will be modified.
+   */
+  static void PrependGrpcFrameHeader(Buffer::Instance* buffer);
+
 private:
   static void checkForHeaderOnlyError(Http::Message& http_response);
 };

source/common/upstream/health_checker_impl.cc

@@ -567,7 +567,7 @@ void GrpcHealthCheckerImpl::GrpcActiveHealthCheckSession::onInterval() {
     request.set_service(parent_.service_name_.value());
   }
 
-  request_encoder_->encodeData(*Grpc::Common::serializeBody(request), true);
+  request_encoder_->encodeData(*Grpc::Common::serializeToGrpcFrame(request), true);
 }
 
 void GrpcHealthCheckerImpl::GrpcActiveHealthCheckSession::onResetStream(Http::StreamResetReason,

source/extensions/tracers/lightstep/lightstep_tracer_impl.cc

@@ -64,7 +64,7 @@ void LightStepDriver::LightStepTransporter::Send(const Protobuf::Message& reques
   Http::MessagePtr message = Grpc::Common::prepareHeaders(
       driver_.cluster()->name(), lightstep::CollectorServiceFullName(),
       lightstep::CollectorMethodName(), absl::optional<std::chrono::milliseconds>(timeout));
-  message->body() = Grpc::Common::serializeBody(request);
+  message->body() = Grpc::Common::serializeToGrpcFrame(request);
 
   active_request_ =
       driver_.clusterManager()

test/common/grpc/common_test.cc

@@ -1,3 +1,5 @@
+#include <arpa/inet.h>
+
 #include "common/grpc/common.h"
 #include "common/http/headers.h"
 #include "common/http/message_impl.h"
@@ -349,5 +351,78 @@ TEST(GrpcCommonTest, ValidateResponse) {
   }
 }
 
+TEST(GrpcCommonTest, MakeBufferInstanceEmpty) {
+  grpc::ByteBuffer byteBuffer;
+  Common::makeBufferInstance(byteBuffer);
+}
+
+TEST(GrpcCommonTest, MakeByteBufferEmpty) {
+  auto buffer = std::make_unique<Buffer::OwnedImpl>();
+  Common::makeByteBuffer(std::move(buffer));
+}
+
+TEST(GrpcCommonTest, MakeBufferInstance1) {
+  grpc::Slice slice("test");
+  grpc::ByteBuffer byteBuffer(&slice, 1);
+  auto bufferInstance = Common::makeBufferInstance(byteBuffer);
+  EXPECT_EQ(bufferInstance->toString(), "test");
+}
+
+TEST(GrpcCommonTest, MakeBufferInstance3) {
+  grpc::Slice slices[3] = {{"test"}, {" "}, {"this"}};
+  grpc::ByteBuffer byteBuffer(slices, 3);
+  auto bufferInstance = Common::makeBufferInstance(byteBuffer);
+  EXPECT_EQ(bufferInstance->toString(), "test this");
+}
+
+TEST(GrpcCommonTest, MakeByteBuffer1) {
+  auto buffer = std::make_unique<Buffer::OwnedImpl>();
+  buffer->add("test", 4);
+  auto byteBuffer = Common::makeByteBuffer(std::move(buffer));
+  std::vector<grpc::Slice> slices;
+  byteBuffer.Dump(&slices);
+  std::string str;
+  for (auto& s : slices) {
+    str.append(std::string(reinterpret_cast<const char*>(s.begin()), s.size()));
+  }
+  EXPECT_EQ(str, "test");
+}
+
+TEST(GrpcCommonTest, MakeByteBuffer3) {
+  auto buffer = std::make_unique<Buffer::OwnedImpl>();
+  buffer->add("test", 4);
+  buffer->add(" ", 1);
+  buffer->add("this", 4);
+  auto byteBuffer = Common::makeByteBuffer(std::move(buffer));
+  std::vector<grpc::Slice> slices;
+  byteBuffer.Dump(&slices);
+  std::string str;
+  for (auto& s : slices) {
+    str.append(std::string(reinterpret_cast<const char*>(s.begin()), s.size()));
+  }
+  EXPECT_EQ(str, "test this");
+}
+
+TEST(GrpcCommonTest, ByteBufferInstanceRoundTrip) {
+  grpc::Slice slices[3] = {{"test"}, {" "}, {"this"}};
+  grpc::ByteBuffer byteBuffer1(slices, 3);
+  auto bufferInstance1 = Common::makeBufferInstance(byteBuffer1);
+  auto byteBuffer2 = Common::makeByteBuffer(std::move(bufferInstance1));
+  auto bufferInstance2 = Common::makeBufferInstance(byteBuffer2);
+  EXPECT_EQ(bufferInstance2->toString(), "test this");
+}
+
+TEST(GrpcCommonTest, PreependGrpcFrameHeader) {
+  auto buffer = std::make_unique<Buffer::OwnedImpl>();
+  buffer->add("test", 4);
+  char expected_header[5];
+  expected_header[0] = 0; // flags
+  const uint32_t nsize = htonl(4);
+  std::memcpy(&expected_header[1], reinterpret_cast<const void*>(&nsize), sizeof(uint32_t));
+  std::string header_string(expected_header, 5);
+  Common::PrependGrpcFrameHeader(buffer.get());
+  EXPECT_EQ(buffer->toString(), header_string + "test");
+}
+
 } // namespace Grpc
 } // namespace Envoy

test/common/grpc/grpc_client_integration_test.cc

@@ -188,7 +188,7 @@ TEST_P(GrpcClientIntegrationTest, ReplyNoTrailers) {
   dispatcher_helper_.setStreamEventPending();
   stream->expectTrailingMetadata(empty_metadata_);
   stream->expectGrpcStatus(Status::GrpcStatus::InvalidCode);
-  auto serialized_response = Grpc::Common::serializeBody(reply);
+  auto serialized_response = Grpc::Common::serializeToGrpcFrame(reply);
   stream->fake_stream_->encodeData(*serialized_response, true);
   stream->fake_stream_->encodeResetStream();
   dispatcher_helper_.runDispatcher();

test/common/upstream/health_checker_impl_test.cc

@@ -2608,7 +2608,7 @@ class GrpcHealthCheckerImplTestBase {
     serializeResponse(grpc::health::v1::HealthCheckResponse::ServingStatus status) {
       grpc::health::v1::HealthCheckResponse response;
       response.set_status(status);
-      const auto data = Grpc::Common::serializeBody(response);
+      const auto data = Grpc::Common::serializeToGrpcFrame(response);
       auto ret = std::vector<uint8_t>(data->length(), 0);
       data->copyOut(0, data->length(), &ret[0]);
       return ret;
@@ -2925,7 +2925,7 @@ TEST_F(GrpcHealthCheckerImplTest, SuccessResponseSplitBetweenChunks) {
 
   grpc::health::v1::HealthCheckResponse response;
   response.set_status(grpc::health::v1::HealthCheckResponse::SERVING);
-  auto data = Grpc::Common::serializeBody(response);
+  auto data = Grpc::Common::serializeToGrpcFrame(response);
 
   const char* raw_data = static_cast<char*>(data->linearize(data->length()));
   const uint64_t chunk_size = data->length() / 5;

test/extensions/filters/http/grpc_json_transcoder/grpc_json_transcoder_integration_test.cc

@@ -109,7 +109,7 @@ class GrpcJsonTranscoderIntegrationTest
         for (const auto& response_message_str : grpc_response_messages) {
           ResponseType response_message;
           EXPECT_TRUE(TextFormat::ParseFromString(response_message_str, &response_message));
-          auto buffer = Grpc::Common::serializeBody(response_message);
+          auto buffer = Grpc::Common::serializeToGrpcFrame(response_message);
           upstream_request_->encodeData(*buffer, false);
         }
         Http::TestHeaderMapImpl response_trailers;