include/flatbuffers/grpc.h

   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:
  35   Message() : slice_(grpc_empty_slice()) {}
  36
  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
  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
  55   ~Message() { grpc_slice_unref(slice_); }
  56
  57   const uint8_t *mutable_data() const { return GRPC_SLICE_START_PTR(slice_); }
  58
  59   const uint8_t *data() const { return GRPC_SLICE_START_PTR(slice_); }
  60
  61   size_t size() const { return GRPC_SLICE_LENGTH(slice_); }
  62
  63   bool Verify() const {
  64     Verifier verifier(data(), size());
  65     return verifier.VerifyBuffer<T>(nullptr);
  66   }
  67
  68   T *GetMutableRoot() { return flatbuffers::GetMutableRoot<T>(mutable_data()); }
  69
  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
  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:
  86   SliceAllocator() : slice_(grpc_empty_slice()) {}
  87
  88   SliceAllocator(const SliceAllocator &other) = delete;
  89   SliceAllocator &operator=(const SliceAllocator &other) = delete;
  90
  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
 103   void swap(SliceAllocator &other) {
 104     using std::swap;
 105     swap(slice_, other.slice_);
 106   }
 107
 108   virtual ~SliceAllocator() { grpc_slice_unref(slice_); }
 109
 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
 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
 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:
 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
 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
 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().
 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
 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.
 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
 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:
 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
 290   static grpc::Status Deserialize(grpc_byte_buffer *buffer,
 291                                   flatbuffers::grpc::Message<T> *msg) {
 292     if (!buffer) {
 293       return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload");
 294     }
 295     // Check if this is a single uncompressed slice.
 296     if ((buffer->type == GRPC_BB_RAW) &&
 297         (buffer->data.raw.compression == GRPC_COMPRESS_NONE) &&
 298         (buffer->data.raw.slice_buffer.count == 1)) {
 299       // If it is, then we can reference the `grpc_slice` directly.
 300       grpc_slice slice = buffer->data.raw.slice_buffer.slices[0];
 301       // We wrap a `Message<T>` around the slice, incrementing the refcount.
 302       *msg = flatbuffers::grpc::Message<T>(slice, true);
 303     } else {
 304       // Otherwise, we need to use `grpc_byte_buffer_reader_readall` to read
 305       // `buffer` into a single contiguous `grpc_slice`. The gRPC reader gives
 306       // us back a new slice with the refcount already incremented.
 307       grpc_byte_buffer_reader reader;
 308       grpc_byte_buffer_reader_init(&reader, buffer);
 309       grpc_slice slice = grpc_byte_buffer_reader_readall(&reader);
 310       grpc_byte_buffer_reader_destroy(&reader);
 311       // We wrap a `Message<T>` around the slice, but don't increment refcount
 312       *msg = flatbuffers::grpc::Message<T>(slice, false);
 313     }
 314     grpc_byte_buffer_destroy(buffer);
 315 #if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION
 316     return ::grpc::Status::OK;
 317 #else
 318     if (msg->Verify()) {
 319       return ::grpc::Status::OK;
 320     } else {
 321       return ::grpc::Status(::grpc::StatusCode::INTERNAL,
 322                             "Message verification failed");
 323     }
 324 #endif
 325   }
 326 };
 327
 328 }  // namespace grpc
 329
 330 #endif  // FLATBUFFERS_GRPC_H_