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/byte_buffer_reader.h" 24 #include "grpcpp/support/byte_buffer.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) : slice_(grpc_empty_slice()) { 92 // default-construct and swap idiom 93 swap(other); 94 } 95 96 SliceAllocator &operator=(SliceAllocator &&other) { 97 // move-construct and swap idiom 98 SliceAllocator temp(std::move(other)); 99 swap(temp); 100 return *this; 101 } 102 swap(SliceAllocator & other)103 void swap(SliceAllocator &other) { 104 using std::swap; 105 swap(slice_, other.slice_); 106 } 107 ~SliceAllocator()108 virtual ~SliceAllocator() { grpc_slice_unref(slice_); } 109 allocate(size_t size)110 virtual uint8_t *allocate(size_t size) override { 111 FLATBUFFERS_ASSERT(GRPC_SLICE_IS_EMPTY(slice_)); 112 slice_ = grpc_slice_malloc(size); 113 return GRPC_SLICE_START_PTR(slice_); 114 } 115 deallocate(uint8_t * p,size_t size)116 virtual void deallocate(uint8_t *p, size_t size) override { 117 FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_)); 118 FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_)); 119 grpc_slice_unref(slice_); 120 slice_ = grpc_empty_slice(); 121 } 122 reallocate_downward(uint8_t * old_p,size_t old_size,size_t new_size,size_t in_use_back,size_t in_use_front)123 virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size, 124 size_t new_size, size_t in_use_back, 125 size_t in_use_front) override { 126 FLATBUFFERS_ASSERT(old_p == GRPC_SLICE_START_PTR(slice_)); 127 FLATBUFFERS_ASSERT(old_size == GRPC_SLICE_LENGTH(slice_)); 128 FLATBUFFERS_ASSERT(new_size > old_size); 129 grpc_slice old_slice = slice_; 130 grpc_slice new_slice = grpc_slice_malloc(new_size); 131 uint8_t *new_p = GRPC_SLICE_START_PTR(new_slice); 132 memcpy_downward(old_p, old_size, new_p, new_size, in_use_back, 133 in_use_front); 134 slice_ = new_slice; 135 grpc_slice_unref(old_slice); 136 return new_p; 137 } 138 139 private: get_slice(uint8_t * p,size_t size)140 grpc_slice &get_slice(uint8_t *p, size_t size) { 141 FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_)); 142 FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_)); 143 return slice_; 144 } 145 146 grpc_slice slice_; 147 148 friend class MessageBuilder; 149 }; 150 151 // SliceAllocatorMember is a hack to ensure that the MessageBuilder's 152 // slice_allocator_ member is constructed before the FlatBufferBuilder, since 153 // the allocator is used in the FlatBufferBuilder ctor. 154 namespace detail { 155 struct SliceAllocatorMember { 156 SliceAllocator slice_allocator_; 157 }; 158 } // namespace detail 159 160 // MessageBuilder is a gRPC-specific FlatBufferBuilder that uses SliceAllocator 161 // to allocate gRPC buffers. 162 class MessageBuilder : private detail::SliceAllocatorMember, 163 public FlatBufferBuilder { 164 public: 165 explicit MessageBuilder(uoffset_t initial_size = 1024) 166 : FlatBufferBuilder(initial_size, &slice_allocator_, false) {} 167 168 MessageBuilder(const MessageBuilder &other) = delete; 169 MessageBuilder &operator=(const MessageBuilder &other) = delete; 170 MessageBuilder(MessageBuilder && other)171 MessageBuilder(MessageBuilder &&other) 172 : FlatBufferBuilder(1024, &slice_allocator_, false) { 173 // Default construct and swap idiom. 174 Swap(other); 175 } 176 177 /// Create a MessageBuilder from a FlatBufferBuilder. 178 explicit MessageBuilder(FlatBufferBuilder &&src, 179 void (*dealloc)(void *, 180 size_t) = &DefaultAllocator::dealloc) 181 : FlatBufferBuilder(1024, &slice_allocator_, false) { 182 src.Swap(*this); 183 src.SwapBufAllocator(*this); 184 if (buf_.capacity()) { 185 uint8_t *buf = buf_.scratch_data(); // pointer to memory 186 size_t capacity = buf_.capacity(); // size of memory 187 slice_allocator_.slice_ = grpc_slice_new_with_len(buf, capacity, dealloc); 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 195 /// supported. 196 MessageBuilder &operator=(FlatBufferBuilder &&src) { 197 // Move construct a temporary and swap 198 MessageBuilder temp(std::move(src)); 199 Swap(temp); 200 return *this; 201 } 202 203 MessageBuilder &operator=(MessageBuilder &&other) { 204 // Move construct a temporary and swap 205 MessageBuilder temp(std::move(other)); 206 Swap(temp); 207 return *this; 208 } 209 Swap(MessageBuilder & other)210 void Swap(MessageBuilder &other) { 211 slice_allocator_.swap(other.slice_allocator_); 212 FlatBufferBuilder::Swap(other); 213 // After swapping the FlatBufferBuilder, we swap back the allocator, which 214 // restores the original allocator back in place. This is necessary because 215 // MessageBuilder's allocator is its own member (SliceAllocatorMember). The 216 // allocator passed to FlatBufferBuilder::vector_downward must point to this 217 // member. 218 buf_.swap_allocator(other.buf_); 219 } 220 221 // Releases the ownership of the buffer pointer. 222 // Returns the size, offset, and the original grpc_slice that 223 // allocated the buffer. Also see grpc_slice_unref(). ReleaseRaw(size_t & size,size_t & offset,grpc_slice & slice)224 uint8_t *ReleaseRaw(size_t &size, size_t &offset, grpc_slice &slice) { 225 uint8_t *buf = FlatBufferBuilder::ReleaseRaw(size, offset); 226 slice = slice_allocator_.slice_; 227 slice_allocator_.slice_ = grpc_empty_slice(); 228 return buf; 229 } 230 ~MessageBuilder()231 ~MessageBuilder() {} 232 233 // GetMessage extracts the subslice of the buffer corresponding to the 234 // flatbuffers-encoded region and wraps it in a `Message<T>` to handle buffer 235 // ownership. GetMessage()236 template<class T> Message<T> GetMessage() { 237 auto buf_data = buf_.scratch_data(); // pointer to memory 238 auto buf_size = buf_.capacity(); // size of memory 239 auto msg_data = buf_.data(); // pointer to msg 240 auto msg_size = buf_.size(); // size of msg 241 // Do some sanity checks on data/size 242 FLATBUFFERS_ASSERT(msg_data); 243 FLATBUFFERS_ASSERT(msg_size); 244 FLATBUFFERS_ASSERT(msg_data >= buf_data); 245 FLATBUFFERS_ASSERT(msg_data + msg_size <= buf_data + buf_size); 246 // Calculate offsets from the buffer start 247 auto begin = msg_data - buf_data; 248 auto end = begin + msg_size; 249 // Get the slice we are working with (no refcount change) 250 grpc_slice slice = slice_allocator_.get_slice(buf_data, buf_size); 251 // Extract a subslice of the existing slice (increment refcount) 252 grpc_slice subslice = grpc_slice_sub(slice, begin, end); 253 // Wrap the subslice in a `Message<T>`, but don't increment refcount 254 Message<T> msg(subslice, false); 255 return msg; 256 } 257 ReleaseMessage()258 template<class T> Message<T> ReleaseMessage() { 259 Message<T> msg = GetMessage<T>(); 260 Reset(); 261 return msg; 262 } 263 264 private: 265 // SliceAllocator slice_allocator_; // part of SliceAllocatorMember 266 }; 267 268 } // namespace grpc 269 } // namespace flatbuffers 270 271 namespace grpc { 272 273 template<class T> class SerializationTraits<flatbuffers::grpc::Message<T>> { 274 public: Serialize(const flatbuffers::grpc::Message<T> & msg,grpc_byte_buffer ** buffer,bool * own_buffer)275 static grpc::Status Serialize(const flatbuffers::grpc::Message<T> &msg, 276 grpc_byte_buffer **buffer, bool *own_buffer) { 277 // We are passed in a `Message<T>`, which is a wrapper around a 278 // `grpc_slice`. We extract it here using `BorrowSlice()`. The const cast 279 // is necessary because the `grpc_raw_byte_buffer_create` func expects 280 // non-const slices in order to increment their refcounts. 281 grpc_slice *slice = const_cast<grpc_slice *>(&msg.BorrowSlice()); 282 // Now use `grpc_raw_byte_buffer_create` to package the single slice into a 283 // `grpc_byte_buffer`, incrementing the refcount in the process. 284 *buffer = grpc_raw_byte_buffer_create(slice, 1); 285 *own_buffer = true; 286 return grpc::Status::OK; 287 } 288 289 // Deserialize by pulling the Deserialize(ByteBuffer * buf,flatbuffers::grpc::Message<T> * msg)290 static grpc::Status Deserialize(ByteBuffer *buf, 291 flatbuffers::grpc::Message<T> *msg) { 292 grpc_byte_buffer *buffer = *reinterpret_cast<grpc_byte_buffer **>(buf); 293 if (!buffer) { 294 return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload"); 295 } 296 // Check if this is a single uncompressed slice. 297 if ((buffer->type == GRPC_BB_RAW) && 298 (buffer->data.raw.compression == GRPC_COMPRESS_NONE) && 299 (buffer->data.raw.slice_buffer.count == 1)) { 300 // If it is, then we can reference the `grpc_slice` directly. 301 grpc_slice slice = buffer->data.raw.slice_buffer.slices[0]; 302 // We wrap a `Message<T>` around the slice, incrementing the refcount. 303 *msg = flatbuffers::grpc::Message<T>(slice, true); 304 } else { 305 // Otherwise, we need to use `grpc_byte_buffer_reader_readall` to read 306 // `buffer` into a single contiguous `grpc_slice`. The gRPC reader gives 307 // us back a new slice with the refcount already incremented. 308 grpc_byte_buffer_reader reader; 309 grpc_byte_buffer_reader_init(&reader, buffer); 310 grpc_slice slice = grpc_byte_buffer_reader_readall(&reader); 311 grpc_byte_buffer_reader_destroy(&reader); 312 // We wrap a `Message<T>` around the slice, but don't increment refcount 313 *msg = flatbuffers::grpc::Message<T>(slice, false); 314 } 315 grpc_byte_buffer_destroy(buffer); 316 #if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION 317 return ::grpc::Status::OK; 318 #else 319 if (msg->Verify()) { 320 return ::grpc::Status::OK; 321 } else { 322 return ::grpc::Status(::grpc::StatusCode::INTERNAL, 323 "Message verification failed"); 324 } 325 #endif 326 } 327 }; 328 329 } // namespace grpc 330 331 #endif // FLATBUFFERS_GRPC_H_ 332