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