grpc.h

Go to the documentation of this file.

/*
 * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
 * Copyright 2014 Google Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#ifndef FLATBUFFERS_GRPC_H_
#define FLATBUFFERS_GRPC_H_
 
// Helper functionality to glue FlatBuffers and GRPC.
 
#include "flatbuffers/flatbuffers.h"
#include "grpc/byte_buffer_reader.h"
#include "grpcpp/support/byte_buffer.h"
 
namespace flatbuffers
{
namespace grpc
{
 
// Message is a typed wrapper around a buffer that manages the underlying
// `grpc_slice` and also provides flatbuffers-specific helpers such as `Verify`
// and `GetRoot`. Since it is backed by a `grpc_slice`, the underlying buffer
// is refcounted and ownership is be managed automatically.
template <class T> class Message
{
public:
  Message() : slice_(grpc_empty_slice()) {}
 
  Message(grpc_slice slice, bool add_ref) : slice_(add_ref ? grpc_slice_ref(slice) : slice) {}
 
  Message &operator=(const Message &other) = delete;
 
  Message(Message &&other) : slice_(other.slice_) { other.slice_ = grpc_empty_slice(); }
 
  Message(const Message &other) = delete;
 
  Message &operator=(Message &&other)
  {
    grpc_slice_unref(slice_);
    slice_ = other.slice_;
    other.slice_ = grpc_empty_slice();
    return *this;
  }
 
  ~Message() { grpc_slice_unref(slice_); }
 
  const uint8_t *mutable_data() const { return GRPC_SLICE_START_PTR(slice_); }
 
  const uint8_t *data() const { return GRPC_SLICE_START_PTR(slice_); }
 
  size_t size() const { return GRPC_SLICE_LENGTH(slice_); }
 
  bool Verify() const
  {
    Verifier verifier(data(), size());
    return verifier.VerifyBuffer<T>(nullptr);
  }
 
  T *GetMutableRoot() { return flatbuffers::GetMutableRoot<T>(mutable_data()); }
 
  const T *GetRoot() const { return flatbuffers::GetRoot<T>(data()); }
 
  // This is only intended for serializer use, or if you know what you're doing
  const grpc_slice &BorrowSlice() const { return slice_; }
 
private:
  grpc_slice slice_;
};
 
class MessageBuilder;
 
// SliceAllocator is a gRPC-specific allocator that uses the `grpc_slice`
// refcounted slices to manage memory ownership. This makes it easy and
// efficient to transfer buffers to gRPC.
class SliceAllocator : public Allocator
{
public:
  SliceAllocator() : slice_(grpc_empty_slice()) {}
 
  SliceAllocator(const SliceAllocator &other) = delete;
  SliceAllocator &operator=(const SliceAllocator &other) = delete;
 
  SliceAllocator(SliceAllocator &&other) : slice_(grpc_empty_slice())
  {
    // default-construct and swap idiom
    swap(other);
  }
 
  SliceAllocator &operator=(SliceAllocator &&other)
  {
    // move-construct and swap idiom
    SliceAllocator temp(std::move(other));
    swap(temp);
    return *this;
  }
 
  void swap(SliceAllocator &other)
  {
    using std::swap;
    swap(slice_, other.slice_);
  }
 
  virtual ~SliceAllocator() { grpc_slice_unref(slice_); }
 
  virtual uint8_t *allocate(size_t size) override
  {
    FLATBUFFERS_ASSERT(GRPC_SLICE_IS_EMPTY(slice_));
    slice_ = grpc_slice_malloc(size);
    return GRPC_SLICE_START_PTR(slice_);
  }
 
  virtual void deallocate(uint8_t *p, size_t size) override
  {
    FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_));
    FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_));
    grpc_slice_unref(slice_);
    slice_ = grpc_empty_slice();
  }
 
  virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size, size_t new_size,
                                       size_t in_use_back, size_t in_use_front) override
  {
    FLATBUFFERS_ASSERT(old_p == GRPC_SLICE_START_PTR(slice_));
    FLATBUFFERS_ASSERT(old_size == GRPC_SLICE_LENGTH(slice_));
    FLATBUFFERS_ASSERT(new_size > old_size);
    grpc_slice old_slice = slice_;
    grpc_slice new_slice = grpc_slice_malloc(new_size);
    uint8_t *new_p = GRPC_SLICE_START_PTR(new_slice);
    memcpy_downward(old_p, old_size, new_p, new_size, in_use_back, in_use_front);
    slice_ = new_slice;
    grpc_slice_unref(old_slice);
    return new_p;
  }
 
private:
  grpc_slice &get_slice(uint8_t *p, size_t size)
  {
    FLATBUFFERS_ASSERT(p == GRPC_SLICE_START_PTR(slice_));
    FLATBUFFERS_ASSERT(size == GRPC_SLICE_LENGTH(slice_));
    return slice_;
  }
 
  grpc_slice slice_;
 
  friend class MessageBuilder;
};
 
// SliceAllocatorMember is a hack to ensure that the MessageBuilder's
// slice_allocator_ member is constructed before the FlatBufferBuilder, since
// the allocator is used in the FlatBufferBuilder ctor.
namespace detail
{
struct SliceAllocatorMember
{
  SliceAllocator slice_allocator_;
};
} // namespace detail
 
// MessageBuilder is a gRPC-specific FlatBufferBuilder that uses SliceAllocator
// to allocate gRPC buffers.
class MessageBuilder : private detail::SliceAllocatorMember, public FlatBufferBuilder
{
public:
  explicit MessageBuilder(uoffset_t initial_size = 1024)
    : FlatBufferBuilder(initial_size, &slice_allocator_, false)
  {
  }
 
  MessageBuilder(const MessageBuilder &other) = delete;
  MessageBuilder &operator=(const MessageBuilder &other) = delete;
 
  MessageBuilder(MessageBuilder &&other) : FlatBufferBuilder(1024, &slice_allocator_, false)
  {
    // Default construct and swap idiom.
    Swap(other);
  }
 
  explicit MessageBuilder(FlatBufferBuilder &&src,
                          void (*dealloc)(void *, size_t) = &DefaultAllocator::dealloc)
    : FlatBufferBuilder(1024, &slice_allocator_, false)
  {
    src.Swap(*this);
    src.SwapBufAllocator(*this);
    if (buf_.capacity())
    {
      uint8_t *buf = buf_.scratch_data(); // pointer to memory
      size_t capacity = buf_.capacity();  // size of memory
      slice_allocator_.slice_ = grpc_slice_new_with_len(buf, capacity, dealloc);
    }
    else
    {
      slice_allocator_.slice_ = grpc_empty_slice();
    }
  }
 
  MessageBuilder &operator=(FlatBufferBuilder &&src)
  {
    // Move construct a temporary and swap
    MessageBuilder temp(std::move(src));
    Swap(temp);
    return *this;
  }
 
  MessageBuilder &operator=(MessageBuilder &&other)
  {
    // Move construct a temporary and swap
    MessageBuilder temp(std::move(other));
    Swap(temp);
    return *this;
  }
 
  void Swap(MessageBuilder &other)
  {
    slice_allocator_.swap(other.slice_allocator_);
    FlatBufferBuilder::Swap(other);
    // After swapping the FlatBufferBuilder, we swap back the allocator, which
    // restores the original allocator back in place. This is necessary because
    // MessageBuilder's allocator is its own member (SliceAllocatorMember). The
    // allocator passed to FlatBufferBuilder::vector_downward must point to this
    // member.
    buf_.swap_allocator(other.buf_);
  }
 
  // Releases the ownership of the buffer pointer.
  // Returns the size, offset, and the original grpc_slice that
  // allocated the buffer. Also see grpc_slice_unref().
  uint8_t *ReleaseRaw(size_t &size, size_t &offset, grpc_slice &slice)
  {
    uint8_t *buf = FlatBufferBuilder::ReleaseRaw(size, offset);
    slice = slice_allocator_.slice_;
    slice_allocator_.slice_ = grpc_empty_slice();
    return buf;
  }
 
  ~MessageBuilder() {}
 
  // GetMessage extracts the subslice of the buffer corresponding to the
  // flatbuffers-encoded region and wraps it in a `Message<T>` to handle buffer
  // ownership.
  template <class T> Message<T> GetMessage()
  {
    auto buf_data = buf_.scratch_data(); // pointer to memory
    auto buf_size = buf_.capacity();     // size of memory
    auto msg_data = buf_.data();         // pointer to msg
    auto msg_size = buf_.size();         // size of msg
    // Do some sanity checks on data/size
    FLATBUFFERS_ASSERT(msg_data);
    FLATBUFFERS_ASSERT(msg_size);
    FLATBUFFERS_ASSERT(msg_data >= buf_data);
    FLATBUFFERS_ASSERT(msg_data + msg_size <= buf_data + buf_size);
    // Calculate offsets from the buffer start
    auto begin = msg_data - buf_data;
    auto end = begin + msg_size;
    // Get the slice we are working with (no refcount change)
    grpc_slice slice = slice_allocator_.get_slice(buf_data, buf_size);
    // Extract a subslice of the existing slice (increment refcount)
    grpc_slice subslice = grpc_slice_sub(slice, begin, end);
    // Wrap the subslice in a `Message<T>`, but don't increment refcount
    Message<T> msg(subslice, false);
    return msg;
  }
 
  template <class T> Message<T> ReleaseMessage()
  {
    Message<T> msg = GetMessage<T>();
    Reset();
    return msg;
  }
 
private:
  // SliceAllocator slice_allocator_;  // part of SliceAllocatorMember
};
 
} // namespace grpc
} // namespace flatbuffers
 
namespace grpc
{
 
template <class T> class SerializationTraits<flatbuffers::grpc::Message<T>>
{
public:
  static grpc::Status Serialize(const flatbuffers::grpc::Message<T> &msg, grpc_byte_buffer **buffer,
                                bool *own_buffer)
  {
    // We are passed in a `Message<T>`, which is a wrapper around a
    // `grpc_slice`. We extract it here using `BorrowSlice()`. The const cast
    // is necessary because the `grpc_raw_byte_buffer_create` func expects
    // non-const slices in order to increment their refcounts.
    grpc_slice *slice = const_cast<grpc_slice *>(&msg.BorrowSlice());
    // Now use `grpc_raw_byte_buffer_create` to package the single slice into a
    // `grpc_byte_buffer`, incrementing the refcount in the process.
    *buffer = grpc_raw_byte_buffer_create(slice, 1);
    *own_buffer = true;
    return grpc::Status::OK;
  }
 
  // Deserialize by pulling the
  static grpc::Status Deserialize(ByteBuffer *buf, flatbuffers::grpc::Message<T> *msg)
  {
    grpc_byte_buffer *buffer = *reinterpret_cast<grpc_byte_buffer **>(buf);
    if (!buffer)
    {
      return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload");
    }
    // Check if this is a single uncompressed slice.
    if ((buffer->type == GRPC_BB_RAW) && (buffer->data.raw.compression == GRPC_COMPRESS_NONE) &&
        (buffer->data.raw.slice_buffer.count == 1))
    {
      // If it is, then we can reference the `grpc_slice` directly.
      grpc_slice slice = buffer->data.raw.slice_buffer.slices[0];
      // We wrap a `Message<T>` around the slice, incrementing the refcount.
      *msg = flatbuffers::grpc::Message<T>(slice, true);
    }
    else
    {
      // Otherwise, we need to use `grpc_byte_buffer_reader_readall` to read
      // `buffer` into a single contiguous `grpc_slice`. The gRPC reader gives
      // us back a new slice with the refcount already incremented.
      grpc_byte_buffer_reader reader;
      grpc_byte_buffer_reader_init(&reader, buffer);
      grpc_slice slice = grpc_byte_buffer_reader_readall(&reader);
      grpc_byte_buffer_reader_destroy(&reader);
      // We wrap a `Message<T>` around the slice, but don't increment refcount
      *msg = flatbuffers::grpc::Message<T>(slice, false);
    }
    grpc_byte_buffer_destroy(buffer);
#if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION
    return ::grpc::Status::OK;
#else
    if (msg->Verify())
    {
      return ::grpc::Status::OK;
    }
    else
    {
      return ::grpc::Status(::grpc::StatusCode::INTERNAL, "Message verification failed");
    }
#endif
  }
};
 
} // namespace grpc
 
#endif // FLATBUFFERS_GRPC_H_