diff options
Diffstat (limited to 'include/grpc++/impl/codegen/proto_utils.h')
-rw-r--r-- | include/grpc++/impl/codegen/proto_utils.h | 133 |
1 files changed, 86 insertions, 47 deletions
diff --git a/include/grpc++/impl/codegen/proto_utils.h b/include/grpc++/impl/codegen/proto_utils.h index 67e8f71a89..b7636034d4 100644 --- a/include/grpc++/impl/codegen/proto_utils.h +++ b/include/grpc++/impl/codegen/proto_utils.h @@ -39,11 +39,13 @@ class GrpcBufferWriterPeer; const int kGrpcBufferWriterMaxBufferLength = 1024 * 1024; -class GrpcBufferWriter final - : public ::grpc::protobuf::io::ZeroCopyOutputStream { +class GrpcBufferWriter : public ::grpc::protobuf::io::ZeroCopyOutputStream { public: - explicit GrpcBufferWriter(grpc_byte_buffer** bp, int block_size) - : block_size_(block_size), byte_count_(0), have_backup_(false) { + GrpcBufferWriter(grpc_byte_buffer** bp, int block_size, int total_size) + : block_size_(block_size), + total_size_(total_size), + byte_count_(0), + have_backup_(false) { *bp = g_core_codegen_interface->grpc_raw_byte_buffer_create(NULL, 0); slice_buffer_ = &(*bp)->data.raw.slice_buffer; } @@ -55,11 +57,20 @@ class GrpcBufferWriter final } bool Next(void** data, int* size) override { + // Protobuf should not ask for more memory than total_size_. + GPR_CODEGEN_ASSERT(byte_count_ < total_size_); if (have_backup_) { slice_ = backup_slice_; have_backup_ = false; } else { - slice_ = g_core_codegen_interface->grpc_slice_malloc(block_size_); + // When less than a whole block is needed, only allocate that much. + // But make sure the allocated slice is not inlined. + size_t remain = total_size_ - byte_count_ > block_size_ + ? block_size_ + : total_size_ - byte_count_; + slice_ = g_core_codegen_interface->grpc_slice_malloc( + remain > GRPC_SLICE_INLINED_SIZE ? remain + : GRPC_SLICE_INLINED_SIZE + 1); } *data = GRPC_SLICE_START_PTR(slice_); // On win x64, int is only 32bit @@ -71,7 +82,7 @@ class GrpcBufferWriter final void BackUp(int count) override { g_core_codegen_interface->grpc_slice_buffer_pop(slice_buffer_); - if (count == block_size_) { + if ((size_t)count == GRPC_SLICE_LENGTH(slice_)) { backup_slice_ = slice_; } else { backup_slice_ = g_core_codegen_interface->grpc_slice_split_tail( @@ -88,9 +99,10 @@ class GrpcBufferWriter final grpc::protobuf::int64 ByteCount() const override { return byte_count_; } - private: + protected: friend class GrpcBufferWriterPeer; const int block_size_; + const int total_size_; int64_t byte_count_; grpc_slice_buffer* slice_buffer_; bool have_backup_; @@ -98,8 +110,7 @@ class GrpcBufferWriter final grpc_slice slice_; }; -class GrpcBufferReader final - : public ::grpc::protobuf::io::ZeroCopyInputStream { +class GrpcBufferReader : public ::grpc::protobuf::io::ZeroCopyInputStream { public: explicit GrpcBufferReader(grpc_byte_buffer* buffer) : byte_count_(0), backup_count_(0), status_() { @@ -160,7 +171,7 @@ class GrpcBufferReader final return byte_count_ - backup_count_; } - private: + protected: int64_t byte_count_; int64_t backup_count_; grpc_byte_buffer_reader reader_; @@ -168,57 +179,85 @@ class GrpcBufferReader final Status status_; }; +// BufferWriter must be a subclass of io::ZeroCopyOutputStream. +template <class BufferWriter, class T> +Status GenericSerialize(const grpc::protobuf::Message& msg, + grpc_byte_buffer** bp, bool* own_buffer) { + static_assert( + std::is_base_of<protobuf::io::ZeroCopyOutputStream, BufferWriter>::value, + "BufferWriter must be a subclass of io::ZeroCopyOutputStream"); + *own_buffer = true; + int byte_size = msg.ByteSize(); + if ((size_t)byte_size <= GRPC_SLICE_INLINED_SIZE) { + grpc_slice slice = g_core_codegen_interface->grpc_slice_malloc(byte_size); + GPR_CODEGEN_ASSERT( + GRPC_SLICE_END_PTR(slice) == + msg.SerializeWithCachedSizesToArray(GRPC_SLICE_START_PTR(slice))); + *bp = g_core_codegen_interface->grpc_raw_byte_buffer_create(&slice, 1); + g_core_codegen_interface->grpc_slice_unref(slice); + + return g_core_codegen_interface->ok(); + } + BufferWriter writer(bp, kGrpcBufferWriterMaxBufferLength, byte_size); + return msg.SerializeToZeroCopyStream(&writer) + ? g_core_codegen_interface->ok() + : Status(StatusCode::INTERNAL, "Failed to serialize message"); +} + +// BufferReader must be a subclass of io::ZeroCopyInputStream. +template <class BufferReader, class T> +Status GenericDeserialize(grpc_byte_buffer* buffer, + grpc::protobuf::Message* msg) { + static_assert( + std::is_base_of<protobuf::io::ZeroCopyInputStream, BufferReader>::value, + "BufferReader must be a subclass of io::ZeroCopyInputStream"); + if (buffer == nullptr) { + return Status(StatusCode::INTERNAL, "No payload"); + } + Status result = g_core_codegen_interface->ok(); + { + BufferReader reader(buffer); + if (!reader.status().ok()) { + return reader.status(); + } + ::grpc::protobuf::io::CodedInputStream decoder(&reader); + decoder.SetTotalBytesLimit(INT_MAX, INT_MAX); + if (!msg->ParseFromCodedStream(&decoder)) { + result = Status(StatusCode::INTERNAL, msg->InitializationErrorString()); + } + if (!decoder.ConsumedEntireMessage()) { + result = Status(StatusCode::INTERNAL, "Did not read entire message"); + } + } + g_core_codegen_interface->grpc_byte_buffer_destroy(buffer); + return result; +} + } // namespace internal +// this is needed so the following class does not conflict with protobuf +// serializers that utilize internal-only tools. +#ifdef GRPC_OPEN_SOURCE_PROTO +// This class provides a protobuf serializer. It translates between protobuf +// objects and grpc_byte_buffers. More information about SerializationTraits can +// be found in include/grpc++/impl/codegen/serialization_traits.h. template <class T> class SerializationTraits<T, typename std::enable_if<std::is_base_of< grpc::protobuf::Message, T>::value>::type> { public: static Status Serialize(const grpc::protobuf::Message& msg, grpc_byte_buffer** bp, bool* own_buffer) { - *own_buffer = true; - int byte_size = msg.ByteSize(); - if (byte_size <= internal::kGrpcBufferWriterMaxBufferLength) { - grpc_slice slice = g_core_codegen_interface->grpc_slice_malloc(byte_size); - GPR_CODEGEN_ASSERT( - GRPC_SLICE_END_PTR(slice) == - msg.SerializeWithCachedSizesToArray(GRPC_SLICE_START_PTR(slice))); - *bp = g_core_codegen_interface->grpc_raw_byte_buffer_create(&slice, 1); - g_core_codegen_interface->grpc_slice_unref(slice); - return g_core_codegen_interface->ok(); - } else { - internal::GrpcBufferWriter writer( - bp, internal::kGrpcBufferWriterMaxBufferLength); - return msg.SerializeToZeroCopyStream(&writer) - ? g_core_codegen_interface->ok() - : Status(StatusCode::INTERNAL, "Failed to serialize message"); - } + return internal::GenericSerialize<internal::GrpcBufferWriter, T>( + msg, bp, own_buffer); } static Status Deserialize(grpc_byte_buffer* buffer, grpc::protobuf::Message* msg) { - if (buffer == nullptr) { - return Status(StatusCode::INTERNAL, "No payload"); - } - Status result = g_core_codegen_interface->ok(); - { - internal::GrpcBufferReader reader(buffer); - if (!reader.status().ok()) { - return reader.status(); - } - ::grpc::protobuf::io::CodedInputStream decoder(&reader); - decoder.SetTotalBytesLimit(INT_MAX, INT_MAX); - if (!msg->ParseFromCodedStream(&decoder)) { - result = Status(StatusCode::INTERNAL, msg->InitializationErrorString()); - } - if (!decoder.ConsumedEntireMessage()) { - result = Status(StatusCode::INTERNAL, "Did not read entire message"); - } - } - g_core_codegen_interface->grpc_byte_buffer_destroy(buffer); - return result; + return internal::GenericDeserialize<internal::GrpcBufferReader, T>(buffer, + msg); } }; +#endif } // namespace grpc |