Add flexibility on *which* slice gets reffed

Use it to ensure that sb_move_first acts predictably

2 files changed, 28 insertions, 12 deletions

diff --git a/src/core/lib/slice/slice.c b/src/core/lib/slice/slice.c
index 219d1a4efc..9e69b443fe 100644
--- a/src/core/lib/slice/slice.c
+++ b/src/core/lib/slice/slice.c
@@ -314,7 +314,7 @@ grpc_slice grpc_slice_sub(grpc_slice source, size_t begin, size_t end) {
 }
 
 grpc_slice grpc_slice_split_tail_maybe_ref(grpc_slice *source, size_t split,
-                                           int incref) {
+                                           grpc_slice_ref_whom ref_whom) {
   grpc_slice tail;
 
   if (source->refcount == NULL) {
@@ -328,26 +328,36 @@ grpc_slice grpc_slice_split_tail_maybe_ref(grpc_slice *source, size_t split,
   } else {
     size_t tail_length = source->data.refcounted.length - split;
     GPR_ASSERT(source->data.refcounted.length >= split);
-    if (tail_length < sizeof(tail.data.inlined.bytes)) {
+    if (tail_length < sizeof(tail.data.inlined.bytes) &&
+        ref_whom != GRPC_SLICE_REF_TAIL) {
       /* Copy out the bytes - it'll be cheaper than refcounting */
       tail.refcount = NULL;
       tail.data.inlined.length = (uint8_t)tail_length;
       memcpy(tail.data.inlined.bytes, source->data.refcounted.bytes + split,
              tail_length);
+      source->refcount = source->refcount->sub_refcount;
     } else {
       /* Build the result */
-      if (incref) {
-        tail.refcount = source->refcount->sub_refcount;
-        /* Bump the refcount */
-        tail.refcount->vtable->ref(tail.refcount);
-      } else {
-        tail.refcount = &noop_refcount;
+      switch (ref_whom) {
+        case GRPC_SLICE_REF_TAIL:
+          tail.refcount = source->refcount->sub_refcount;
+          source->refcount = &noop_refcount;
+          break;
+        case GRPC_SLICE_REF_HEAD:
+          tail.refcount = &noop_refcount;
+          source->refcount = source->refcount->sub_refcount;
+          break;
+        case GRPC_SLICE_REF_BOTH:
+          tail.refcount = source->refcount->sub_refcount;
+          source->refcount = source->refcount->sub_refcount;
+          /* Bump the refcount */
+          tail.refcount->vtable->ref(tail.refcount);
+          break;
       }
       /* Point into the source array */
       tail.data.refcounted.bytes = source->data.refcounted.bytes + split;
       tail.data.refcounted.length = tail_length;
     }
-    source->refcount = source->refcount->sub_refcount;
     source->data.refcounted.length = split;
   }
 
@@ -355,7 +365,7 @@ grpc_slice grpc_slice_split_tail_maybe_ref(grpc_slice *source, size_t split,
 }
 
 grpc_slice grpc_slice_split_tail(grpc_slice *source, size_t split) {
-  return grpc_slice_split_tail_maybe_ref(source, split, true);
+  return grpc_slice_split_tail_maybe_ref(source, split, GRPC_SLICE_REF_BOTH);
 }
 
 grpc_slice grpc_slice_split_head(grpc_slice *source, size_t split) {
diff --git a/src/core/lib/slice/slice_buffer.c b/src/core/lib/slice/slice_buffer.c
index 5a4b434d28..e8d41ca0f7 100644
--- a/src/core/lib/slice/slice_buffer.c
+++ b/src/core/lib/slice/slice_buffer.c
@@ -274,12 +274,18 @@ static void slice_buffer_move_first_maybe_ref(grpc_slice_buffer *src, size_t n,
     } else if (n == slice_len) {
       grpc_slice_buffer_add(dst, slice);
       break;
-    } else { /* n < slice_len */
+    } else if (incref) { /* n < slice_len */
       grpc_slice_buffer_undo_take_first(
-          src, grpc_slice_split_tail_maybe_ref(&slice, n, incref));
+          src, grpc_slice_split_tail_maybe_ref(&slice, n, GRPC_SLICE_REF_BOTH));
       GPR_ASSERT(GRPC_SLICE_LENGTH(slice) == n);
       grpc_slice_buffer_add(dst, slice);
       break;
+    } else { /* n < slice_len */
+      grpc_slice_buffer_undo_take_first(
+          src, grpc_slice_split_tail_maybe_ref(&slice, n, GRPC_SLICE_REF_TAIL));
+      GPR_ASSERT(GRPC_SLICE_LENGTH(slice) == n);
+      grpc_slice_buffer_add_indexed(dst, slice);
+      break;
     }
   }
   GPR_ASSERT(dst->length == output_len);