Zero copy for receive (#222)

Signed-off-by: ice1000 <[email protected]>

Author: ice1000

.gitignore

@@ -3,3 +3,5 @@ Cargo.lock
 *.rs.bk
 *.rs.fmt
 .vscode
+.idea
+*.o

benchmark/src/bench.rs

@@ -14,13 +14,15 @@
 // TODO: Remove it once Rust's tool_lints is stabilized.
 #![cfg_attr(feature = "cargo-clippy", allow(renamed_and_removed_lints))]
 
+use std::io::Read;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use futures::{Future, Sink, Stream};
 use grpc::{
-    self, ClientStreamingSink, DuplexSink, Method, MethodType, RequestStream, RpcContext,
-    RpcStatus, RpcStatusCode, ServerStreamingSink, ServiceBuilder, UnarySink, WriteFlags,
+    self, ClientStreamingSink, DuplexSink, MessageReader, Method, MethodType, RequestStream,
+    RpcContext, RpcStatus, RpcStatusCode, ServerStreamingSink, ServiceBuilder, UnarySink,
+    WriteFlags,
 };
 use grpc_proto::testing::messages::{SimpleRequest, SimpleResponse};
 use grpc_proto::testing::services_grpc::BenchmarkService;
@@ -115,8 +117,10 @@ pub fn bin_ser(t: &Vec<u8>, buf: &mut Vec<u8>) {
 }
 
 #[inline]
-pub fn bin_de(buf: &[u8]) -> grpc::Result<Vec<u8>> {
-    Ok(buf.to_vec())
+pub fn bin_de(mut reader: MessageReader) -> grpc::Result<Vec<u8>> {
+    let mut buf = vec![];
+    reader.read_to_end(&mut buf).unwrap();
+    Ok(buf)
 }
 
 pub const METHOD_BENCHMARK_SERVICE_GENERIC_CALL: Method<Vec<u8>, Vec<u8>> = Method {

benchmark/src/client.rs

@@ -325,11 +325,13 @@ fn execute<B: Backoff + Send + 'static>(
                 }
                 RequestExecutor::new(ctx, ch, cfg).execute_unary_async()
             }
-            RpcType::STREAMING => if cfg.get_payload_config().has_bytebuf_params() {
-                GenericExecutor::new(ctx, ch, cfg).execute_stream()
-            } else {
-                RequestExecutor::new(ctx, ch, cfg).execute_stream_ping_pong()
-            },
+            RpcType::STREAMING => {
+                if cfg.get_payload_config().has_bytebuf_params() {
+                    GenericExecutor::new(ctx, ch, cfg).execute_stream()
+                } else {
+                    RequestExecutor::new(ctx, ch, cfg).execute_stream_ping_pong()
+                }
+            }
             _ => unimplemented!(),
         },
         _ => unimplemented!(),

grpc-sys/grpc_wrap.cc

@@ -297,42 +297,25 @@ grpcwrap_batch_context_recv_initial_metadata(
   return &(ctx->recv_initial_metadata);
 }
 
-GPR_EXPORT size_t GPR_CALLTYPE
-grpcwrap_batch_context_recv_message_length(const grpcwrap_batch_context* ctx) {
-  grpc_byte_buffer_reader reader;
-  if (!ctx->recv_message) {
-    return (size_t)-1;
-  }
-
-  GPR_ASSERT(grpc_byte_buffer_reader_init(&reader, ctx->recv_message));
-  size_t result = grpc_byte_buffer_length(reader.buffer_out);
-  grpc_byte_buffer_reader_destroy(&reader);
+GPR_EXPORT const char *GPR_CALLTYPE
+grpcwrap_slice_raw_offset(const grpc_slice *slice, size_t offset, size_t *len) {
+  *len = GRPC_SLICE_LENGTH(*slice) - offset;
+  return (const char *)(GRPC_SLICE_START_PTR(*slice)) + offset;
+}
 
-  return result;
+GPR_EXPORT size_t GPR_CALLTYPE
+grpcwrap_slice_length(const grpc_slice *slice) {
+  return GRPC_SLICE_LENGTH(*slice);
 }
 
-/*
- * Copies data from recv_message to a buffer. Fatal error occurs if
- * buffer is too small.
- */
-GPR_EXPORT void GPR_CALLTYPE grpcwrap_batch_context_recv_message_to_buffer(
-    const grpcwrap_batch_context* ctx, char* buffer, size_t buffer_len) {
-  grpc_byte_buffer_reader reader;
-  grpc_slice slice;
-  size_t offset = 0;
-
-  GPR_ASSERT(grpc_byte_buffer_reader_init(&reader, ctx->recv_message));
-
-  while (grpc_byte_buffer_reader_next(&reader, &slice)) {
-    size_t len = GRPC_SLICE_LENGTH(slice);
-    GPR_ASSERT(offset + len <= buffer_len);
-    memcpy(buffer + offset, GRPC_SLICE_START_PTR(slice),
-           GRPC_SLICE_LENGTH(slice));
-    offset += len;
-    grpc_slice_unref(slice);
+GPR_EXPORT grpc_byte_buffer *GPR_CALLTYPE
+grpcwrap_batch_context_take_recv_message(grpcwrap_batch_context *ctx) {
+  grpc_byte_buffer *buf = NULL;
+  if (ctx->recv_message) {
+    buf = ctx->recv_message;
+    ctx->recv_message = NULL;
   }
-
-  grpc_byte_buffer_reader_destroy(&reader);
+  return buf;
 }
 
 GPR_EXPORT grpc_status_code GPR_CALLTYPE
@@ -858,3 +841,17 @@ grpcwrap_ssl_server_credentials_create(
 }
 
 #endif
+
+/* Sanity check for complicated types */
+
+GPR_EXPORT void GPR_CALLTYPE grpcwrap_sanity_check_slice(size_t size,
+                                                         size_t align) {
+  GPR_ASSERT(sizeof(grpc_slice) == size);
+  GPR_ASSERT(alignof(grpc_slice) == align);
+}
+
+GPR_EXPORT void GPR_CALLTYPE
+grpcwrap_sanity_check_byte_buffer_reader(size_t size, size_t align) {
+  GPR_ASSERT(sizeof(grpc_byte_buffer_reader) == size);
+  GPR_ASSERT(alignof(grpc_byte_buffer_reader) == align);
+}

grpc-sys/src/lib.rs

@@ -12,10 +12,13 @@
 // limitations under the License.
 
 #![allow(unknown_lints)]
+#![allow(renamed_and_removed_lints)]
+// remove this after Rust's tool_lints is stabilized
 
 extern crate libc;
 
-use libc::{c_char, c_int, c_uint, c_void, int32_t, int64_t, size_t, uint32_t};
+use libc::{c_char, c_int, c_uint, c_void, int32_t, int64_t, size_t, uint32_t, uint8_t};
+use std::mem;
 use std::time::Duration;
 
 /// The clocks gRPC supports.
@@ -348,6 +351,52 @@ pub struct GrpcMetadataArray {
     pub metadata: *mut GrpcMetadata,
 }
 
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct GrpcSliceRefCounted {
+    bytes: *mut uint8_t,
+    length: size_t,
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct GrpcSliceInlined {
+    length: uint8_t,
+    // TODO: use size_of when it becomes a const function.
+    #[cfg(target_pointer_width = "64")]
+    bytes: [uint8_t; 23],
+    #[cfg(target_pointer_width = "32")]
+    bytes: [uint8_t; 11],
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub union GrpcSliceData {
+    ref_counted: GrpcSliceRefCounted,
+    inlined: GrpcSliceInlined,
+}
+
+pub enum GrpcSliceRefCount {}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct GrpcSlice {
+    ref_count: *mut GrpcSliceRefCount,
+    data: GrpcSliceData,
+}
+
+#[repr(C)]
+pub union GrpcByteBufferReaderCurrent {
+    index: c_uint,
+}
+
+#[repr(C)]
+pub struct GrpcByteBufferReader {
+    pub buffer_in: *mut GrpcByteBuffer,
+    pub buffer_out: *mut GrpcByteBuffer,
+    current: GrpcByteBufferReaderCurrent,
+}
+
 pub const GRPC_INITIAL_METADATA_IDEMPOTENT_REQUEST: uint32_t = 0x0000_0010;
 pub const GRPC_INITIAL_METADATA_WAIT_FOR_READY: uint32_t = 0x0000_0020;
 pub const GRPC_INITIAL_METADATA_CACHEABLE_REQUEST: uint32_t = 0x0000_0040;
@@ -356,7 +405,6 @@ pub const GRPC_WRITE_BUFFER_HINT: uint32_t = 0x0000_0001;
 pub const GRPC_WRITE_NO_COMPRESS: uint32_t = 0x0000_0002;
 
 pub enum GrpcMetadata {}
-pub enum GrpcSlice {}
 pub enum GrpcCallDetails {}
 pub enum GrpcCompletionQueue {}
 pub enum GrpcChannel {}
@@ -443,17 +491,33 @@ extern "C" {
     ) -> *mut GrpcChannel;
     pub fn grpc_channel_destroy(channel: *mut GrpcChannel);
 
+    pub fn grpc_slice_unref(slice: GrpcSlice);
+    pub fn grpc_byte_buffer_length(buf: *const GrpcByteBuffer) -> size_t;
+    pub fn grpcwrap_slice_length(slice: *const GrpcSlice) -> size_t;
+    pub fn grpcwrap_slice_raw_offset(
+        slice: *const GrpcSlice,
+        offset: size_t,
+        len: *mut size_t,
+    ) -> *const c_char;
+    pub fn grpc_byte_buffer_reader_init(
+        reader: *mut GrpcByteBufferReader,
+        buf: *mut GrpcByteBuffer,
+    ) -> c_int;
+    pub fn grpc_byte_buffer_reader_next(
+        reader: *mut GrpcByteBufferReader,
+        buf: *mut GrpcSlice,
+    ) -> c_int;
+    pub fn grpc_byte_buffer_reader_destroy(reader: *mut GrpcByteBufferReader);
+    pub fn grpc_byte_buffer_destroy(buf: *mut GrpcByteBuffer);
+
     pub fn grpcwrap_batch_context_create() -> *mut GrpcBatchContext;
     pub fn grpcwrap_batch_context_destroy(ctx: *mut GrpcBatchContext);
     pub fn grpcwrap_batch_context_recv_initial_metadata(
         ctx: *mut GrpcBatchContext,
     ) -> *const GrpcMetadataArray;
-    pub fn grpcwrap_batch_context_recv_message_length(ctx: *mut GrpcBatchContext) -> size_t;
-    pub fn grpcwrap_batch_context_recv_message_to_buffer(
+    pub fn grpcwrap_batch_context_take_recv_message(
         ctx: *mut GrpcBatchContext,
-        buffer: *mut c_char,
-        buffer_len: size_t,
-    );
+    ) -> *mut GrpcByteBuffer;
     pub fn grpcwrap_batch_context_recv_status_on_client_status(
         ctx: *mut GrpcBatchContext,
     ) -> GrpcStatusCode;
@@ -476,7 +540,7 @@ extern "C" {
     pub fn grpcwrap_call_start_unary(
         call: *mut GrpcCall,
         ctx: *mut GrpcBatchContext,
-        send_bufer: *const c_char,
+        send_buffer: *const c_char,
         send_buffer_len: size_t,
         write_flags: uint32_t,
         initial_metadata: *mut GrpcMetadataArray,
@@ -493,7 +557,7 @@ extern "C" {
     pub fn grpcwrap_call_start_server_streaming(
         call: *mut GrpcCall,
         ctx: *mut GrpcBatchContext,
-        send_bufer: *const c_char,
+        send_buffer: *const c_char,
         send_buffer_len: size_t,
         write_flags: uint32_t,
         initial_metadata: *mut GrpcMetadataArray,
@@ -515,7 +579,7 @@ extern "C" {
     pub fn grpcwrap_call_send_message(
         call: *mut GrpcCall,
         ctx: *mut GrpcBatchContext,
-        send_bufer: *const c_char,
+        send_buffer: *const c_char,
         send_buffer_len: size_t,
         write_flags: uint32_t,
         send_empty_initial_metadata: uint32_t,
@@ -570,7 +634,7 @@ extern "C" {
         server: *mut GrpcServer,
         method: *const c_char,
         host: *const c_char,
-        paylod_handling: GrpcServerRegisterMethodPayloadHandling,
+        payload_handling: GrpcServerRegisterMethodPayloadHandling,
         flags: uint32_t,
     ) -> *mut c_void;
     pub fn grpc_server_create(
@@ -646,6 +710,19 @@ extern "C" {
     pub fn grpcwrap_metadata_array_cleanup(array: *mut GrpcMetadataArray);
 
     pub fn gpr_free(p: *mut c_void);
+
+    pub fn grpcwrap_sanity_check_slice(size: size_t, align: size_t);
+    pub fn grpcwrap_sanity_check_byte_buffer_reader(size: size_t, align: size_t);
+}
+
+/// Make sure the complicated struct written in rust is the same with
+/// its C one.
+pub unsafe fn sanity_check() {
+    grpcwrap_sanity_check_slice(mem::size_of::<GrpcSlice>(), mem::align_of::<GrpcSlice>());
+    grpcwrap_sanity_check_byte_buffer_reader(
+        mem::size_of::<GrpcByteBufferReader>(),
+        mem::align_of::<GrpcByteBufferReader>(),
+    );
 }
 
 #[cfg(feature = "secure")]
@@ -655,6 +732,7 @@ mod secure_component {
     use super::{GrpcChannel, GrpcChannelArgs, GrpcServer};
 
     pub enum GrpcChannelCredentials {}
+
     pub enum GrpcServerCredentials {}
 
     extern "C" {
@@ -703,6 +781,7 @@ mod tests {
     fn smoke() {
         unsafe {
             super::grpc_init();
+            super::sanity_check();
             let cq = super::grpc_completion_queue_create_for_next(ptr::null_mut());
             super::grpc_completion_queue_destroy(cq);
             super::grpc_shutdown();

src/async/callback.rs

@@ -69,9 +69,8 @@ impl UnaryRequest {
             return;
         }
 
-        let data = self.ctx.batch_ctx().recv_message();
-        self.ctx
-            .handle(&mut rc, cq, data.as_ref().map(|v| v.as_slice()));
+        let reader = self.ctx.batch_ctx_mut().recv_message();
+        self.ctx.handle(&mut rc, cq, reader);
         server::request_call(rc, cq);
     }
 }

src/async/mod.rs

@@ -26,7 +26,7 @@ use self::callback::{Abort, Request as RequestCallback, UnaryRequest as UnaryReq
 use self::executor::SpawnNotify;
 use self::promise::{Batch as BatchPromise, Shutdown as ShutdownPromise};
 use call::server::RequestContext;
-use call::{BatchContext, Call};
+use call::{BatchContext, Call, MessageReader};
 use cq::CompletionQueue;
 use error::{Error, Result};
 use server::RequestCallContext;
@@ -115,9 +115,8 @@ impl<T> Future for CqFuture<T> {
     }
 }
 
-pub type BatchMessage = Option<Vec<u8>>;
 /// Future object for batch jobs.
-pub type BatchFuture = CqFuture<BatchMessage>;
+pub type BatchFuture = CqFuture<Option<MessageReader>>;
 
 /// A result holder for asynchronous execution.
 // This enum is going to be passed to FFI, so don't use trait or generic here.

src/async/promise.rs

@@ -14,8 +14,8 @@
 use std::fmt::{self, Debug, Formatter};
 use std::sync::Arc;
 
-use super::{BatchMessage, Inner};
-use call::{BatchContext, RpcStatusCode};
+use super::Inner;
+use call::{BatchContext, MessageReader, RpcStatusCode};
 use error::Error;
 
 /// Batch job type.
@@ -33,11 +33,11 @@ pub enum BatchType {
 pub struct Batch {
     ty: BatchType,
     ctx: BatchContext,
-    inner: Arc<Inner<BatchMessage>>,
+    inner: Arc<Inner<Option<MessageReader>>>,
 }
 
 impl Batch {
-    pub fn new(ty: BatchType, inner: Arc<Inner<BatchMessage>>) -> Batch {
+    pub fn new(ty: BatchType, inner: Arc<Inner<Option<MessageReader>>>) -> Batch {
         Batch {
             ty,
             ctx: BatchContext::new(),