Modifying slice buffer add to merge two contiguous slices sharing the same refcount object (#29466)

* Modifying slice buffer add to merge two contiguous slices sharing the same refcount object

* sanity checks

* updating test

* fixing grpc_slice_buffer_add API misuse in proto utils

* fix sanity checks

* minor fix

include/grpcpp/impl/codegen/core_codegen.h

@@ -107,6 +107,8 @@ class CoreCodegen final : public CoreCodegenInterface {
   grpc_slice grpc_slice_sub(grpc_slice s, size_t begin, size_t end) override;
   void grpc_slice_buffer_add(grpc_slice_buffer* sb, grpc_slice slice) override;
   void grpc_slice_buffer_pop(grpc_slice_buffer* sb) override;
+  void grpc_slice_buffer_add_indexed(grpc_slice_buffer* sb,
+                                     grpc_slice slice) override;
   grpc_slice grpc_slice_from_static_buffer(const void* buffer,
                                            size_t length) override;
   grpc_slice grpc_slice_from_copied_buffer(const void* buffer,

include/grpcpp/impl/codegen/core_codegen_interface.h

@@ -127,6 +127,8 @@ class CoreCodegenInterface {
   virtual grpc_slice grpc_slice_sub(grpc_slice s, size_t begin, size_t end) = 0;
   virtual void grpc_slice_buffer_add(grpc_slice_buffer* sb,
                                      grpc_slice slice) = 0;
+  virtual void grpc_slice_buffer_add_indexed(grpc_slice_buffer* sb,
+                                             grpc_slice slice) = 0;
   virtual void grpc_slice_buffer_pop(grpc_slice_buffer* sb) = 0;
   virtual grpc_slice grpc_slice_from_static_buffer(const void* buffer,
                                                    size_t length) = 0;

include/grpcpp/impl/codegen/proto_buffer_writer.h

@@ -110,7 +110,12 @@ class ProtoBufferWriter : public grpc::protobuf::io::ZeroCopyOutputStream {
     // On win x64, int is only 32bit
     GPR_CODEGEN_ASSERT(GRPC_SLICE_LENGTH(slice_) <= INT_MAX);
     byte_count_ += * size = static_cast<int>(GRPC_SLICE_LENGTH(slice_));
-    g_core_codegen_interface->grpc_slice_buffer_add(slice_buffer_, slice_);
+    // Using grpc_slice_buffer_add could modify slice_ and merge it with the
+    // previous slice. Therefore, use grpc_slice_buffer_add_indexed method to
+    // ensure the slice gets added at a separate index. It can then be kept
+    // around and popped later in the BackUp function.
+    g_core_codegen_interface->grpc_slice_buffer_add_indexed(slice_buffer_,
+                                                            slice_);
     return true;
   }
 

src/core/lib/slice/slice_buffer.cc

@@ -138,13 +138,29 @@ size_t grpc_slice_buffer_add_indexed(grpc_slice_buffer* sb, grpc_slice s) {
 
 void grpc_slice_buffer_add(grpc_slice_buffer* sb, grpc_slice s) {
   size_t n = sb->count;
-  /* if both the last slice in the slice buffer and the slice being added
+  grpc_slice* back = nullptr;
+  if (n != 0) {
+    back = &sb->slices[n - 1];
+  }
+  if (s.refcount != nullptr && back != nullptr &&
+      s.refcount == back->refcount &&
+      GRPC_SLICE_START_PTR(s) == GRPC_SLICE_END_PTR(*back)) {
+    // Merge the two slices into one because they are contiguous and share the
+    // same refcount object.
+    back->data.refcounted.length += GRPC_SLICE_LENGTH(s);
+    sb->length += GRPC_SLICE_LENGTH(s);
+    // Unref the merged slice.
+    grpc_slice_unref_internal(s);
+    // early out
+    return;
+  }
+
+  if (!s.refcount && n) {
+    /* if both the last slice in the slice buffer and the slice being added
      are inlined (that is, that they carry their data inside the slice data
      structure), and the back slice is not full, then concatenate directly
      into the back slice, preventing many small slices being passed into
      writes */
-  if (!s.refcount && n) {
-    grpc_slice* back = &sb->slices[n - 1];
     if (!back->refcount &&
         back->data.inlined.length < GRPC_SLICE_INLINED_SIZE) {
       if (s.data.inlined.length + back->data.inlined.length <=

src/cpp/common/core_codegen.cc

@@ -210,6 +210,11 @@ void CoreCodegen::grpc_slice_buffer_add(grpc_slice_buffer* sb,
   ::grpc_slice_buffer_add(sb, slice);
 }
 
+void CoreCodegen::grpc_slice_buffer_add_indexed(grpc_slice_buffer* sb,
+                                                grpc_slice slice) {
+  ::grpc_slice_buffer_add_indexed(sb, slice);
+}
+
 void CoreCodegen::grpc_slice_buffer_pop(grpc_slice_buffer* sb) {
   ::grpc_slice_buffer_pop(sb);
 }

test/core/slice/slice_buffer_test.cc

@@ -23,6 +23,8 @@
 #include "src/core/lib/slice/slice_internal.h"
 #include "test/core/util/test_config.h"
 
+static constexpr size_t kTotalDataLength = 4096;
+
 void test_slice_buffer_add() {
   grpc_slice_buffer buf;
   grpc_slice aaa = grpc_slice_from_copied_string("aaa");
@@ -59,6 +61,35 @@ void test_slice_buffer_add() {
   grpc_slice_buffer_destroy(&buf);
 }
 
+static void free_data(void* data, size_t len) {
+  GPR_ASSERT(len == kTotalDataLength);
+  gpr_free(data);
+}
+
+void test_slice_buffer_add_contiguous_slices() {
+  grpc_slice_buffer buf;
+  grpc_slice_buffer_init(&buf);
+  char* data = reinterpret_cast<char*>(gpr_malloc(kTotalDataLength));
+  GPR_ASSERT(data != nullptr);
+  grpc_slice a = grpc_slice_new_with_len(data, kTotalDataLength, free_data);
+  grpc_slice s1 = grpc_slice_split_head(&a, kTotalDataLength / 4);
+  grpc_slice s2 = grpc_slice_split_head(&a, kTotalDataLength / 4);
+  grpc_slice s3 = grpc_slice_split_head(&a, kTotalDataLength / 4);
+  grpc_slice_buffer_add(&buf, s1);
+  GPR_ASSERT(buf.count == 1);
+  GPR_ASSERT(buf.length == kTotalDataLength / 4);
+  grpc_slice_buffer_add(&buf, s2);
+  GPR_ASSERT(buf.count == 1);
+  GPR_ASSERT(buf.length == kTotalDataLength / 2);
+  grpc_slice_buffer_add(&buf, s3);
+  GPR_ASSERT(buf.count == 1);
+  GPR_ASSERT(buf.length == 3 * kTotalDataLength / 4);
+  grpc_slice_buffer_add(&buf, a);
+  GPR_ASSERT(buf.count == 1);
+  GPR_ASSERT(buf.length == kTotalDataLength);
+  grpc_slice_buffer_destroy(&buf);
+}
+
 void test_slice_buffer_move_first() {
   grpc_slice slices[3];
   grpc_slice_buffer src;
@@ -153,6 +184,7 @@ int main(int argc, char** argv) {
   grpc_init();
 
   test_slice_buffer_add();
+  test_slice_buffer_add_contiguous_slices();
   test_slice_buffer_move_first();
   test_slice_buffer_first();