diff --git a/include/grpc/slice.h b/include/grpc/slice.h index d65bde49ee6..039347f17e3 100644 --- a/include/grpc/slice.h +++ b/include/grpc/slice.h @@ -124,11 +124,17 @@ GPRAPI grpc_slice grpc_slice_sub_no_ref(grpc_slice s, size_t begin, size_t end); Requires s intialized, split <= s.length */ GPRAPI grpc_slice grpc_slice_split_tail(grpc_slice *s, size_t split); +typedef enum { + GRPC_SLICE_REF_TAIL = 1, + GRPC_SLICE_REF_HEAD = 2, + GRPC_SLICE_REF_BOTH = 1 + 2 +} grpc_slice_ref_whom; + /* The same as grpc_slice_split_tail, but with an option to skip altering * refcounts (grpc_slice_split_tail_maybe_ref(..., true) is equivalent to * grpc_slice_split_tail(...)) */ GPRAPI grpc_slice grpc_slice_split_tail_maybe_ref(grpc_slice *s, size_t split, - int inc_refs); + grpc_slice_ref_whom ref_whom); /* Splits s into two: modifies s to be s[split:s.length], and returns a new slice, sharing a refcount with s, that contains s[0:split]. diff --git a/src/core/lib/slice/slice.c b/src/core/lib/slice/slice.c index 219d1a4efcf..9e69b443fe1 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 5a4b434d28c..e8d41ca0f7d 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);