• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright 2014 Google Inc. All rights reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef FLATBUFFERS_GRPC_H_
18 #define FLATBUFFERS_GRPC_H_
19 
20 // Helper functionality to glue FlatBuffers and GRPC.
21 
22 #include "flatbuffers/flatbuffers.h"
23 #include "grpc++/support/byte_buffer.h"
24 #include "grpc/byte_buffer_reader.h"
25 
26 namespace flatbuffers {
27 namespace grpc {
28 
29 // Message is a typed wrapper around a buffer that manages the underlying
30 // `grpc_slice` and also provides flatbuffers-specific helpers such as `Verify`
31 // and `GetRoot`. Since it is backed by a `grpc_slice`, the underlying buffer
32 // is refcounted and ownership is be managed automatically.
33 template<class T> class Message {
34  public:
Message()35   Message() : slice_(grpc_empty_slice()) {}
36 
Message(grpc_slice slice,bool add_ref)37   Message(grpc_slice slice, bool add_ref)
38       : slice_(add_ref ? grpc_slice_ref(slice) : slice) {}
39 
40   Message &operator=(const Message &other) = delete;
41 
Message(Message && other)42   Message(Message &&other) : slice_(other.slice_) {
43     other.slice_ = grpc_empty_slice();
44   }
45 
46   Message(const Message &other) = delete;
47 
48   Message &operator=(Message &&other) {
49     grpc_slice_unref(slice_);
50     slice_ = other.slice_;
51     other.slice_ = grpc_empty_slice();
52     return *this;
53   }
54 
~Message()55   ~Message() { grpc_slice_unref(slice_); }
56 
mutable_data()57   const uint8_t *mutable_data() const { return GRPC_SLICE_START_PTR(slice_); }
58 
data()59   const uint8_t *data() const { return GRPC_SLICE_START_PTR(slice_); }
60 
size()61   size_t size() const { return GRPC_SLICE_LENGTH(slice_); }
62 
Verify()63   bool Verify() const {
64     Verifier verifier(data(), size());
65     return verifier.VerifyBuffer<T>(nullptr);
66   }
67 
GetMutableRoot()68   T *GetMutableRoot() { return flatbuffers::GetMutableRoot<T>(mutable_data()); }
69 
GetRoot()70   const T *GetRoot() const { return flatbuffers::GetRoot<T>(data()); }
71 
72   // This is only intended for serializer use, or if you know what you're doing
BorrowSlice()73   const grpc_slice &BorrowSlice() const { return slice_; }
74 
75  private:
76   grpc_slice slice_;
77 };
78 
79 class MessageBuilder;
80 
81 // SliceAllocator is a gRPC-specific allocator that uses the `grpc_slice`
82 // refcounted slices to manage memory ownership. This makes it easy and
83 // efficient to transfer buffers to gRPC.
84 class SliceAllocator : public Allocator {
85  public:
SliceAllocator()86   SliceAllocator() : slice_(grpc_empty_slice()) {}
87 
88   SliceAllocator(const SliceAllocator &other) = delete;
89   SliceAllocator &operator=(const SliceAllocator &other) = delete;
90 
SliceAllocator(SliceAllocator && other)91   SliceAllocator(SliceAllocator &&other)
92     : slice_(grpc_empty_slice()) {
93     // default-construct and swap idiom
94     swap(other);
95   }
96 
97   SliceAllocator &operator=(SliceAllocator &&other) {
98     // move-construct and swap idiom
99     SliceAllocator temp(std::move(other));
100     swap(temp);
101     return *this;
102   }
103 
swap(SliceAllocator & other)104   void swap(SliceAllocator &other) {
105     using std::swap;
106     swap(slice_, other.slice_);
107   }
108 
~SliceAllocator()109   virtual ~SliceAllocator() { grpc_slice_unref(slice_); }
110 
allocate(size_t size)111   virtual uint8_t *allocate(size_t size) override {
112     FLATBUFFERS_ASSERT(GRPC_SLICE_IS_EMPTY(slice_));
113     slice_ = grpc_slice_malloc(size);
114     return GRPC_SLICE_START_PTR(slice_);
115   }
116 
deallocate(uint8_t * p,size_t size)117   virtual void deallocate(uint8_t *p, size_t size) override {
118     FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_));
119     FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_));
120     grpc_slice_unref(slice_);
121     slice_ = grpc_empty_slice();
122   }
123 
reallocate_downward(uint8_t * old_p,size_t old_size,size_t new_size,size_t in_use_back,size_t in_use_front)124   virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size,
125                                        size_t new_size, size_t in_use_back,
126                                        size_t in_use_front) override {
127     FLATBUFFERS_ASSERT(old_p == GRPC_SLICE_START_PTR(slice_));
128     FLATBUFFERS_ASSERT(old_size == GRPC_SLICE_LENGTH(slice_));
129     FLATBUFFERS_ASSERT(new_size > old_size);
130     grpc_slice old_slice = slice_;
131     grpc_slice new_slice = grpc_slice_malloc(new_size);
132     uint8_t *new_p = GRPC_SLICE_START_PTR(new_slice);
133     memcpy_downward(old_p, old_size, new_p, new_size, in_use_back,
134                     in_use_front);
135     slice_ = new_slice;
136     grpc_slice_unref(old_slice);
137     return new_p;
138   }
139 
140  private:
get_slice(uint8_t * p,size_t size)141   grpc_slice &get_slice(uint8_t *p, size_t size) {
142     FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_));
143     FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_));
144     return slice_;
145   }
146 
147   grpc_slice slice_;
148 
149   friend class MessageBuilder;
150 };
151 
152 // SliceAllocatorMember is a hack to ensure that the MessageBuilder's
153 // slice_allocator_ member is constructed before the FlatBufferBuilder, since
154 // the allocator is used in the FlatBufferBuilder ctor.
155 namespace detail {
156 struct SliceAllocatorMember {
157   SliceAllocator slice_allocator_;
158 };
159 }  // namespace detail
160 
161 // MessageBuilder is a gRPC-specific FlatBufferBuilder that uses SliceAllocator
162 // to allocate gRPC buffers.
163 class MessageBuilder : private detail::SliceAllocatorMember,
164                        public FlatBufferBuilder {
165  public:
166   explicit MessageBuilder(uoffset_t initial_size = 1024)
167     : FlatBufferBuilder(initial_size, &slice_allocator_, false) {}
168 
169   MessageBuilder(const MessageBuilder &other) = delete;
170   MessageBuilder &operator=(const MessageBuilder &other) = delete;
171 
MessageBuilder(MessageBuilder && other)172   MessageBuilder(MessageBuilder &&other)
173     : FlatBufferBuilder(1024, &slice_allocator_, false) {
174     // Default construct and swap idiom.
175     Swap(other);
176   }
177 
178   /// Create a MessageBuilder from a FlatBufferBuilder.
179   explicit MessageBuilder(FlatBufferBuilder &&src, void (*dealloc)(void*, size_t) = &DefaultAllocator::dealloc)
180     : FlatBufferBuilder(1024, &slice_allocator_, false) {
181     src.Swap(*this);
182     src.SwapBufAllocator(*this);
183     if (buf_.capacity()) {
184       uint8_t *buf = buf_.scratch_data();       // pointer to memory
185       size_t capacity = buf_.capacity();        // size of memory
186       slice_allocator_.slice_ = grpc_slice_new_with_len(buf, capacity, dealloc);
187     }
188     else {
189       slice_allocator_.slice_ = grpc_empty_slice();
190     }
191   }
192 
193   /// Move-assign a FlatBufferBuilder to a MessageBuilder.
194   /// Only FlatBufferBuilder with default allocator (basically, nullptr) is supported.
195   MessageBuilder &operator=(FlatBufferBuilder &&src) {
196     // Move construct a temporary and swap
197     MessageBuilder temp(std::move(src));
198     Swap(temp);
199     return *this;
200   }
201 
202   MessageBuilder &operator=(MessageBuilder &&other) {
203     // Move construct a temporary and swap
204     MessageBuilder temp(std::move(other));
205     Swap(temp);
206     return *this;
207   }
208 
Swap(MessageBuilder & other)209   void Swap(MessageBuilder &other) {
210     slice_allocator_.swap(other.slice_allocator_);
211     FlatBufferBuilder::Swap(other);
212     // After swapping the FlatBufferBuilder, we swap back the allocator, which restores
213     // the original allocator back in place. This is necessary because MessageBuilder's
214     // allocator is its own member (SliceAllocatorMember). The allocator passed to
215     // FlatBufferBuilder::vector_downward must point to this member.
216     buf_.swap_allocator(other.buf_);
217   }
218 
219   // Releases the ownership of the buffer pointer.
220   // Returns the size, offset, and the original grpc_slice that
221   // allocated the buffer. Also see grpc_slice_unref().
ReleaseRaw(size_t & size,size_t & offset,grpc_slice & slice)222   uint8_t *ReleaseRaw(size_t &size, size_t &offset, grpc_slice &slice) {
223     uint8_t *buf = FlatBufferBuilder::ReleaseRaw(size, offset);
224     slice = slice_allocator_.slice_;
225     slice_allocator_.slice_ = grpc_empty_slice();
226     return buf;
227   }
228 
~MessageBuilder()229   ~MessageBuilder() {}
230 
231   // GetMessage extracts the subslice of the buffer corresponding to the
232   // flatbuffers-encoded region and wraps it in a `Message<T>` to handle buffer
233   // ownership.
GetMessage()234   template<class T> Message<T> GetMessage() {
235     auto buf_data = buf_.scratch_data();       // pointer to memory
236     auto buf_size = buf_.capacity();  // size of memory
237     auto msg_data = buf_.data();      // pointer to msg
238     auto msg_size = buf_.size();      // size of msg
239     // Do some sanity checks on data/size
240     FLATBUFFERS_ASSERT(msg_data);
241     FLATBUFFERS_ASSERT(msg_size);
242     FLATBUFFERS_ASSERT(msg_data >= buf_data);
243     FLATBUFFERS_ASSERT(msg_data + msg_size <= buf_data + buf_size);
244     // Calculate offsets from the buffer start
245     auto begin = msg_data - buf_data;
246     auto end = begin + msg_size;
247     // Get the slice we are working with (no refcount change)
248     grpc_slice slice = slice_allocator_.get_slice(buf_data, buf_size);
249     // Extract a subslice of the existing slice (increment refcount)
250     grpc_slice subslice = grpc_slice_sub(slice, begin, end);
251     // Wrap the subslice in a `Message<T>`, but don't increment refcount
252     Message<T> msg(subslice, false);
253     return msg;
254   }
255 
ReleaseMessage()256   template<class T> Message<T> ReleaseMessage() {
257     Message<T> msg = GetMessage<T>();
258     Reset();
259     return msg;
260   }
261 
262  private:
263   // SliceAllocator slice_allocator_;  // part of SliceAllocatorMember
264 };
265 
266 }  // namespace grpc
267 }  // namespace flatbuffers
268 
269 namespace grpc {
270 
271 template<class T> class SerializationTraits<flatbuffers::grpc::Message<T>> {
272  public:
Serialize(const flatbuffers::grpc::Message<T> & msg,grpc_byte_buffer ** buffer,bool * own_buffer)273   static grpc::Status Serialize(const flatbuffers::grpc::Message<T> &msg,
274                                 grpc_byte_buffer **buffer, bool *own_buffer) {
275     // We are passed in a `Message<T>`, which is a wrapper around a
276     // `grpc_slice`. We extract it here using `BorrowSlice()`. The const cast
277     // is necesary because the `grpc_raw_byte_buffer_create` func expects
278     // non-const slices in order to increment their refcounts.
279     grpc_slice *slice = const_cast<grpc_slice *>(&msg.BorrowSlice());
280     // Now use `grpc_raw_byte_buffer_create` to package the single slice into a
281     // `grpc_byte_buffer`, incrementing the refcount in the process.
282     *buffer = grpc_raw_byte_buffer_create(slice, 1);
283     *own_buffer = true;
284     return grpc::Status::OK;
285   }
286 
287   // Deserialize by pulling the
Deserialize(grpc_byte_buffer * buffer,flatbuffers::grpc::Message<T> * msg)288   static grpc::Status Deserialize(grpc_byte_buffer *buffer,
289                                   flatbuffers::grpc::Message<T> *msg) {
290     if (!buffer) {
291       return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload");
292     }
293     // Check if this is a single uncompressed slice.
294     if ((buffer->type == GRPC_BB_RAW) &&
295         (buffer->data.raw.compression == GRPC_COMPRESS_NONE) &&
296         (buffer->data.raw.slice_buffer.count == 1)) {
297       // If it is, then we can reference the `grpc_slice` directly.
298       grpc_slice slice = buffer->data.raw.slice_buffer.slices[0];
299       // We wrap a `Message<T>` around the slice, incrementing the refcount.
300       *msg = flatbuffers::grpc::Message<T>(slice, true);
301     } else {
302       // Otherwise, we need to use `grpc_byte_buffer_reader_readall` to read
303       // `buffer` into a single contiguous `grpc_slice`. The gRPC reader gives
304       // us back a new slice with the refcount already incremented.
305       grpc_byte_buffer_reader reader;
306       grpc_byte_buffer_reader_init(&reader, buffer);
307       grpc_slice slice = grpc_byte_buffer_reader_readall(&reader);
308       grpc_byte_buffer_reader_destroy(&reader);
309       // We wrap a `Message<T>` around the slice, but dont increment refcount
310       *msg = flatbuffers::grpc::Message<T>(slice, false);
311     }
312     grpc_byte_buffer_destroy(buffer);
313 #if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION
314     return ::grpc::Status::OK;
315 #else
316     if (msg->Verify()) {
317       return ::grpc::Status::OK;
318     } else {
319       return ::grpc::Status(::grpc::StatusCode::INTERNAL,
320                             "Message verification failed");
321     }
322 #endif
323   }
324 };
325 
326 }  // namespace grpc
327 
328 #endif  // FLATBUFFERS_GRPC_H_
329