1 /* 2 * 3 * Copyright 2018 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19 #ifndef GRPCPP_IMPL_CODEGEN_PROTO_BUFFER_WRITER_H 20 #define GRPCPP_IMPL_CODEGEN_PROTO_BUFFER_WRITER_H 21 22 // IWYU pragma: private, include <grpcpp/support/proto_buffer_writer.h> 23 24 #include <type_traits> 25 26 #include <grpc/impl/codegen/grpc_types.h> 27 #include <grpc/impl/codegen/slice.h> 28 #include <grpcpp/impl/codegen/byte_buffer.h> 29 #include <grpcpp/impl/codegen/config_protobuf.h> 30 #include <grpcpp/impl/codegen/core_codegen_interface.h> 31 #include <grpcpp/impl/codegen/serialization_traits.h> 32 #include <grpcpp/impl/codegen/status.h> 33 34 /// This header provides an object that writes bytes directly into a 35 /// grpc::ByteBuffer, via the ZeroCopyOutputStream interface 36 37 namespace grpc { 38 39 extern CoreCodegenInterface* g_core_codegen_interface; 40 41 // Forward declaration for testing use only 42 namespace internal { 43 class ProtoBufferWriterPeer; 44 } // namespace internal 45 46 const int kProtoBufferWriterMaxBufferLength = 1024 * 1024; 47 48 /// This is a specialization of the protobuf class ZeroCopyOutputStream. 49 /// The principle is to give the proto layer one buffer of bytes at a time 50 /// that it can use to serialize the next portion of the message, with the 51 /// option to "backup" if more buffer is given than required at the last buffer. 52 /// 53 /// Read more about ZeroCopyOutputStream interface here: 54 /// https://developers.google.com/protocol-buffers/docs/reference/cpp/google.protobuf.io.zero_copy_stream#ZeroCopyOutputStream 55 class ProtoBufferWriter : public ::grpc::protobuf::io::ZeroCopyOutputStream { 56 public: 57 /// Constructor for this derived class 58 /// 59 /// \param[out] byte_buffer A pointer to the grpc::ByteBuffer created 60 /// \param block_size How big are the chunks to allocate at a time 61 /// \param total_size How many total bytes are required for this proto ProtoBufferWriter(ByteBuffer * byte_buffer,int block_size,int total_size)62 ProtoBufferWriter(ByteBuffer* byte_buffer, int block_size, int total_size) 63 : block_size_(block_size), 64 total_size_(total_size), 65 byte_count_(0), 66 have_backup_(false) { 67 GPR_CODEGEN_ASSERT(!byte_buffer->Valid()); 68 /// Create an empty raw byte buffer and look at its underlying slice buffer 69 grpc_byte_buffer* bp = 70 g_core_codegen_interface->grpc_raw_byte_buffer_create(nullptr, 0); 71 byte_buffer->set_buffer(bp); 72 slice_buffer_ = &bp->data.raw.slice_buffer; 73 } 74 ~ProtoBufferWriter()75 ~ProtoBufferWriter() override { 76 if (have_backup_) { 77 g_core_codegen_interface->grpc_slice_unref(backup_slice_); 78 } 79 } 80 81 /// Give the proto library the next buffer of bytes and its size. It is 82 /// safe for the caller to write from data[0, size - 1]. Next(void ** data,int * size)83 bool Next(void** data, int* size) override { 84 // Protobuf should not ask for more memory than total_size_. 85 GPR_CODEGEN_ASSERT(byte_count_ < total_size_); 86 // 1. Use the remaining backup slice if we have one 87 // 2. Otherwise allocate a slice, up to the remaining length needed 88 // or our maximum allocation size 89 // 3. Provide the slice start and size available 90 // 4. Add the slice being returned to the slice buffer 91 size_t remain = static_cast<size_t>(total_size_ - byte_count_); 92 if (have_backup_) { 93 /// If we have a backup slice, we should use it first 94 slice_ = backup_slice_; 95 have_backup_ = false; 96 if (GRPC_SLICE_LENGTH(slice_) > remain) { 97 GRPC_SLICE_SET_LENGTH(slice_, remain); 98 } 99 } else { 100 // When less than a whole block is needed, only allocate that much. 101 // But make sure the allocated slice is not inlined. 102 size_t allocate_length = 103 remain > static_cast<size_t>(block_size_) ? block_size_ : remain; 104 slice_ = g_core_codegen_interface->grpc_slice_malloc( 105 allocate_length > GRPC_SLICE_INLINED_SIZE 106 ? allocate_length 107 : GRPC_SLICE_INLINED_SIZE + 1); 108 } 109 *data = GRPC_SLICE_START_PTR(slice_); 110 // On win x64, int is only 32bit 111 GPR_CODEGEN_ASSERT(GRPC_SLICE_LENGTH(slice_) <= INT_MAX); 112 byte_count_ += * size = static_cast<int>(GRPC_SLICE_LENGTH(slice_)); 113 g_core_codegen_interface->grpc_slice_buffer_add(slice_buffer_, slice_); 114 return true; 115 } 116 117 /// Backup by \a count bytes because Next returned more bytes than needed 118 /// (only used in the last buffer). \a count must be less than or equal too 119 /// the last buffer returned from next. BackUp(int count)120 void BackUp(int count) override { 121 // count == 0 is invoked by ZeroCopyOutputStream users indicating that any 122 // potential buffer obtained through a previous call to Next() is final. 123 // ZeroCopyOutputStream implementations such as streaming output can use 124 // these calls to flush any temporary buffer and flush the output. The logic 125 // below is not robust against count == 0 invocations, so directly return. 126 if (count == 0) return; 127 128 /// 1. Remove the partially-used last slice from the slice buffer 129 /// 2. Split it into the needed (if any) and unneeded part 130 /// 3. Add the needed part back to the slice buffer 131 /// 4. Mark that we still have the remaining part (for later use/unref) 132 GPR_CODEGEN_ASSERT(count <= static_cast<int>(GRPC_SLICE_LENGTH(slice_))); 133 g_core_codegen_interface->grpc_slice_buffer_pop(slice_buffer_); 134 if (static_cast<size_t>(count) == GRPC_SLICE_LENGTH(slice_)) { 135 backup_slice_ = slice_; 136 } else { 137 backup_slice_ = g_core_codegen_interface->grpc_slice_split_tail( 138 &slice_, GRPC_SLICE_LENGTH(slice_) - count); 139 g_core_codegen_interface->grpc_slice_buffer_add(slice_buffer_, slice_); 140 } 141 // It's dangerous to keep an inlined grpc_slice as the backup slice, since 142 // on a following Next() call, a reference will be returned to this slice 143 // via GRPC_SLICE_START_PTR, which will not be an address held by 144 // slice_buffer_. 145 have_backup_ = backup_slice_.refcount != nullptr; 146 byte_count_ -= count; 147 } 148 149 /// Returns the total number of bytes written since this object was created. ByteCount()150 int64_t ByteCount() const override { return byte_count_; } 151 152 // These protected members are needed to support internal optimizations. 153 // they expose internal bits of grpc core that are NOT stable. If you have 154 // a use case needs to use one of these functions, please send an email to 155 // https://groups.google.com/forum/#!forum/grpc-io. 156 protected: slice_buffer()157 grpc_slice_buffer* slice_buffer() { return slice_buffer_; } set_byte_count(int64_t byte_count)158 void set_byte_count(int64_t byte_count) { byte_count_ = byte_count; } 159 160 private: 161 // friend for testing purposes only 162 friend class internal::ProtoBufferWriterPeer; 163 const int block_size_; ///< size to alloc for each new \a grpc_slice needed 164 const int total_size_; ///< byte size of proto being serialized 165 int64_t byte_count_; ///< bytes written since this object was created 166 grpc_slice_buffer* 167 slice_buffer_; ///< internal buffer of slices holding the serialized data 168 bool have_backup_; ///< if we are holding a backup slice or not 169 grpc_slice backup_slice_; ///< holds space we can still write to, if the 170 ///< caller has called BackUp 171 grpc_slice slice_; ///< current slice passed back to the caller 172 }; 173 174 } // namespace grpc 175 176 #endif // GRPCPP_IMPL_CODEGEN_PROTO_BUFFER_WRITER_H 177