include/flatbuffers/grpc.h

0001 /*
0002  * Copyright 2014 Google Inc. All rights reserved.
0003  *
0004  * Licensed under the Apache License, Version 2.0 (the "License");
0005  * you may not use this file except in compliance with the License.
0006  * You may obtain a copy of the License at
0007  *
0008  *     http://www.apache.org/licenses/LICENSE-2.0
0009  *
0010  * Unless required by applicable law or agreed to in writing, software
0011  * distributed under the License is distributed on an "AS IS" BASIS,
0012  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0013  * See the License for the specific language governing permissions and
0014  * limitations under the License.
0015  */
0016
0017 #ifndef FLATBUFFERS_GRPC_H_
0018 #define FLATBUFFERS_GRPC_H_
0019
0020 // Helper functionality to glue FlatBuffers and GRPC.
0021
0022 #include "flatbuffers/flatbuffers.h"
0023 #include "grpcpp/support/byte_buffer.h"
0024 #include "grpcpp/support/slice.h"
0025
0026 namespace flatbuffers {
0027 namespace grpc {
0028
0029 // Message is a typed wrapper around a buffer that manages the underlying
0030 // `grpc_slice` and also provides flatbuffers-specific helpers such as `Verify`
0031 // and `GetRoot`. Since it is backed by a `grpc_slice`, the underlying buffer
0032 // is refcounted and ownership is be managed automatically.
0033 template<class T> class Message {
0034  public:
0035   Message() {}
0036
0037   Message(::grpc::Slice slice) : slice_(slice) {}
0038
0039   Message &operator=(const Message &other) = delete;
0040
0041   Message(Message &&other) = default;
0042
0043   Message(const Message &other) = delete;
0044
0045   Message &operator=(Message &&other) = default;
0046
0047   const uint8_t *mutable_data() const { return slice_.begin(); }
0048
0049   const uint8_t *data() const { return slice_.begin(); }
0050
0051   size_t size() const { return slice_.size(); }
0052
0053   bool Verify() const {
0054     Verifier verifier(data(), size());
0055     return verifier.VerifyBuffer<T>(nullptr);
0056   }
0057
0058   T *GetMutableRoot() { return flatbuffers::GetMutableRoot<T>(mutable_data()); }
0059
0060   const T *GetRoot() const { return flatbuffers::GetRoot<T>(data()); }
0061
0062   // This is only intended for serializer use, or if you know what you're doing
0063   const ::grpc::Slice &BorrowSlice() const { return slice_; }
0064
0065  private:
0066   ::grpc::Slice slice_;
0067 };
0068
0069 class MessageBuilder;
0070
0071 // SliceAllocator is a gRPC-specific allocator that uses the `grpc_slice`
0072 // refcounted slices to manage memory ownership. This makes it easy and
0073 // efficient to transfer buffers to gRPC.
0074 class SliceAllocator : public Allocator {
0075  public:
0076   SliceAllocator() {}
0077
0078   SliceAllocator(const SliceAllocator &other) = delete;
0079   SliceAllocator &operator=(const SliceAllocator &other) = delete;
0080
0081   SliceAllocator(SliceAllocator &&other) {
0082     // default-construct and swap idiom
0083     swap(other);
0084   }
0085
0086   SliceAllocator &operator=(SliceAllocator &&other) {
0087     // move-construct and swap idiom
0088     SliceAllocator temp(std::move(other));
0089     swap(temp);
0090     return *this;
0091   }
0092
0093   void swap(SliceAllocator &other) {
0094     using std::swap;
0095     swap(slice_, other.slice_);
0096   }
0097
0098   virtual ~SliceAllocator() {}
0099
0100   virtual uint8_t *allocate(size_t size) override {
0101     FLATBUFFERS_ASSERT(slice_.size() == 0);
0102     slice_ = ::grpc::Slice(size);
0103     return const_cast<uint8_t *>(slice_.begin());
0104   }
0105
0106   virtual void deallocate(uint8_t *p, size_t size) override {
0107     FLATBUFFERS_ASSERT(p == slice_.begin());
0108     FLATBUFFERS_ASSERT(size == slice_.size());
0109     slice_ = ::grpc::Slice();
0110   }
0111
0112   virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size,
0113                                        size_t new_size, size_t in_use_back,
0114                                        size_t in_use_front) override {
0115     FLATBUFFERS_ASSERT(old_p == slice_.begin());
0116     FLATBUFFERS_ASSERT(old_size == slice_.size());
0117     FLATBUFFERS_ASSERT(new_size > old_size);
0118     ::grpc::Slice old_slice = slice_;
0119     ::grpc::Slice new_slice = ::grpc::Slice(new_size);
0120     uint8_t *new_p = const_cast<uint8_t *>(new_slice.begin());
0121     memcpy_downward(old_p, old_size, new_p, new_size, in_use_back,
0122                     in_use_front);
0123     slice_ = new_slice;
0124     return new_p;
0125   }
0126
0127  private:
0128   ::grpc::Slice &get_slice(uint8_t *p, size_t size) {
0129     FLATBUFFERS_ASSERT(p == slice_.begin());
0130     FLATBUFFERS_ASSERT(size == slice_.size());
0131     return slice_;
0132   }
0133
0134   ::grpc::Slice slice_;
0135
0136   friend class MessageBuilder;
0137 };
0138
0139 // SliceAllocatorMember is a hack to ensure that the MessageBuilder's
0140 // slice_allocator_ member is constructed before the FlatBufferBuilder, since
0141 // the allocator is used in the FlatBufferBuilder ctor.
0142 namespace detail {
0143 struct SliceAllocatorMember {
0144   SliceAllocator slice_allocator_;
0145 };
0146 }  // namespace detail
0147
0148 // MessageBuilder is a gRPC-specific FlatBufferBuilder that uses SliceAllocator
0149 // to allocate gRPC buffers.
0150 class MessageBuilder : private detail::SliceAllocatorMember,
0151                        public FlatBufferBuilder {
0152  public:
0153   explicit MessageBuilder(uoffset_t initial_size = 1024)
0154       : FlatBufferBuilder(initial_size, &slice_allocator_, false) {}
0155
0156   MessageBuilder(const MessageBuilder &other) = delete;
0157   MessageBuilder &operator=(const MessageBuilder &other) = delete;
0158
0159   MessageBuilder(MessageBuilder &&other)
0160       : FlatBufferBuilder(1024, &slice_allocator_, false) {
0161     // Default construct and swap idiom.
0162     Swap(other);
0163   }
0164
0165   /// Create a MessageBuilder from a FlatBufferBuilder.
0166   explicit MessageBuilder(FlatBufferBuilder &&src,
0167                           void (*dealloc)(void *,
0168                                           size_t) = &DefaultAllocator::dealloc)
0169       : FlatBufferBuilder(1024, &slice_allocator_, false) {
0170     src.Swap(*this);
0171     src.SwapBufAllocator(*this);
0172     if (buf_.capacity()) {
0173       uint8_t *buf = buf_.scratch_data();  // pointer to memory
0174       size_t capacity = buf_.capacity();   // size of memory
0175       slice_allocator_.slice_ = ::grpc::Slice(buf, capacity, dealloc);
0176     } else {
0177       slice_allocator_.slice_ = ::grpc::Slice();
0178     }
0179   }
0180
0181   /// Move-assign a FlatBufferBuilder to a MessageBuilder.
0182   /// Only FlatBufferBuilder with default allocator (basically, nullptr) is
0183   /// supported.
0184   MessageBuilder &operator=(FlatBufferBuilder &&src) {
0185     // Move construct a temporary and swap
0186     MessageBuilder temp(std::move(src));
0187     Swap(temp);
0188     return *this;
0189   }
0190
0191   MessageBuilder &operator=(MessageBuilder &&other) {
0192     // Move construct a temporary and swap
0193     MessageBuilder temp(std::move(other));
0194     Swap(temp);
0195     return *this;
0196   }
0197
0198   void Swap(MessageBuilder &other) {
0199     slice_allocator_.swap(other.slice_allocator_);
0200     FlatBufferBuilder::Swap(other);
0201     // After swapping the FlatBufferBuilder, we swap back the allocator, which
0202     // restores the original allocator back in place. This is necessary because
0203     // MessageBuilder's allocator is its own member (SliceAllocatorMember). The
0204     // allocator passed to FlatBufferBuilder::vector_downward must point to this
0205     // member.
0206     buf_.swap_allocator(other.buf_);
0207   }
0208
0209   // Releases the ownership of the buffer pointer.
0210   // Returns the size, offset, and the original grpc_slice that
0211   // allocated the buffer. Also see grpc_slice_unref().
0212   uint8_t *ReleaseRaw(size_t &size, size_t &offset, ::grpc::Slice &slice) {
0213     uint8_t *buf = FlatBufferBuilder::ReleaseRaw(size, offset);
0214     slice = slice_allocator_.slice_;
0215     slice_allocator_.slice_ = ::grpc::Slice();
0216     return buf;
0217   }
0218
0219   ~MessageBuilder() {}
0220
0221   // GetMessage extracts the subslice of the buffer corresponding to the
0222   // flatbuffers-encoded region and wraps it in a `Message<T>` to handle buffer
0223   // ownership.
0224   template<class T> Message<T> GetMessage() {
0225     auto buf_data = buf_.scratch_data();  // pointer to memory
0226     auto buf_size = buf_.capacity();      // size of memory
0227     auto msg_data = buf_.data();          // pointer to msg
0228     auto msg_size = buf_.size();          // size of msg
0229     // Do some sanity checks on data/size
0230     FLATBUFFERS_ASSERT(msg_data);
0231     FLATBUFFERS_ASSERT(msg_size);
0232     FLATBUFFERS_ASSERT(msg_data >= buf_data);
0233     FLATBUFFERS_ASSERT(msg_data + msg_size <= buf_data + buf_size);
0234     // Calculate offsets from the buffer start
0235     auto begin = msg_data - buf_data;
0236     auto end = begin + msg_size;
0237     // Get the slice we are working with (no refcount change)
0238     ::grpc::Slice slice = slice_allocator_.get_slice(buf_data, buf_size);
0239     // Extract a subslice of the existing slice (increment refcount)
0240     ::grpc::Slice subslice = slice.sub(begin, end);
0241     // Wrap the subslice in a `Message<T>`, but don't increment refcount
0242     Message<T> msg(subslice);
0243     return msg;
0244   }
0245
0246   template<class T> Message<T> ReleaseMessage() {
0247     Message<T> msg = GetMessage<T>();
0248     Reset();
0249     return msg;
0250   }
0251
0252  private:
0253   // SliceAllocator slice_allocator_;  // part of SliceAllocatorMember
0254 };
0255
0256 }  // namespace grpc
0257 }  // namespace flatbuffers
0258
0259 namespace grpc {
0260
0261 template<class T> class SerializationTraits<flatbuffers::grpc::Message<T>> {
0262  public:
0263   static grpc::Status Serialize(const flatbuffers::grpc::Message<T> &msg,
0264                                 ByteBuffer *buffer, bool *own_buffer) {
0265     // Package the single slice into a `ByteBuffer`,
0266     // incrementing the refcount in the process.
0267     *buffer = ByteBuffer(&msg.BorrowSlice(), 1);
0268     *own_buffer = true;
0269     return grpc::Status::OK;
0270   }
0271
0272   // Deserialize by pulling the
0273   static grpc::Status Deserialize(ByteBuffer *buf,
0274                                   flatbuffers::grpc::Message<T> *msg) {
0275     Slice slice;
0276     if (!buf->TrySingleSlice(&slice).ok()) {
0277       if (!buf->DumpToSingleSlice(&slice).ok()) {
0278         buf->Clear();
0279         return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload");
0280       }
0281     }
0282     *msg = flatbuffers::grpc::Message<T>(slice);
0283     buf->Clear();
0284 #if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION
0285     return ::grpc::Status::OK;
0286 #else
0287     if (msg->Verify()) {
0288       return ::grpc::Status::OK;
0289     } else {
0290       return ::grpc::Status(::grpc::StatusCode::INTERNAL,
0291                             "Message verification failed");
0292     }
0293 #endif
0294   }
0295 };
0296
0297 }  // namespace grpc
0298
0299 #endif  // FLATBUFFERS_GRPC_H_