proto_buffer_writer.h

Go to the documentation of this file.

//
//
// Copyright 2018 gRPC authors.
//
// 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 GRPCPP_SUPPORT_PROTO_BUFFER_WRITER_H
#define GRPCPP_SUPPORT_PROTO_BUFFER_WRITER_H
 
#include <grpc/byte_buffer.h>
#include <grpc/event_engine/memory_allocator.h>
#include <grpc/impl/grpc_types.h>
#include <grpc/slice.h>
#include <grpc/slice_buffer.h>
#include <grpcpp/impl/codegen/config_protobuf.h>
#include <grpcpp/impl/serialization_traits.h>
#include <grpcpp/support/byte_buffer.h>
#include <grpcpp/support/status.h>
 
#include <type_traits>
 
#include "absl/log/absl_check.h"
#include "absl/strings/cord.h"
 
 
namespace grpc {
 
// Forward declaration for testing use only
namespace internal {
class ProtoBufferWriterPeer;
}  // namespace internal
 
const int kProtoBufferWriterMaxBufferLength = 1024 * 1024;
 
class ProtoBufferWriter : public grpc::protobuf::io::ZeroCopyOutputStream {
 public:
  ProtoBufferWriter(
      ByteBuffer* byte_buffer, int block_size, int total_size,
      grpc_event_engine::experimental::MemoryAllocator* allocator = nullptr)
      : block_size_(block_size),
        total_size_(total_size),
        byte_count_(0),
        allocator_(allocator),
        have_backup_(false) {
    ABSL_CHECK(!byte_buffer->Valid());
    grpc_byte_buffer* bp = grpc_raw_byte_buffer_create(nullptr, 0);
    byte_buffer->set_buffer(bp);
    slice_buffer_ = &bp->data.raw.slice_buffer;
  }
 
  ~ProtoBufferWriter() override {
    if (have_backup_) {
      grpc_slice_unref(backup_slice_);
    }
  }
 
  bool Next(void** data, int* size) override {
    // Protobuf should not ask for more memory than total_size_.
    ABSL_CHECK_LT(byte_count_, total_size_);
    // 1. Use the remaining backup slice if we have one
    // 2. Otherwise allocate a slice, up to the remaining length needed
    //    or our maximum allocation size
    // 3. Provide the slice start and size available
    // 4. Add the slice being returned to the slice buffer
    size_t remain = static_cast<size_t>(total_size_ - byte_count_);
    if (have_backup_) {
      slice_ = backup_slice_;
      have_backup_ = false;
      if (GRPC_SLICE_LENGTH(slice_) > remain) {
        GRPC_SLICE_SET_LENGTH(slice_, remain);
      }
    } else {
      // When less than a whole block is needed, only allocate that much.
      // But make sure the allocated slice is not inlined.
      const size_t allocate_length =
          remain > static_cast<size_t>(block_size_) ? block_size_ : remain;
      const size_t slice_size = allocate_length > GRPC_SLICE_INLINED_SIZE
                                    ? allocate_length
                                    : GRPC_SLICE_INLINED_SIZE + 1;
      if (allocator_ == nullptr) {
        slice_ = grpc_slice_malloc(slice_size);
      } else {
        slice_ = allocator_->MakeSlice(slice_size);
      }
    }
    *data = GRPC_SLICE_START_PTR(slice_);
    // On win x64, int is only 32bit
    ABSL_CHECK(GRPC_SLICE_LENGTH(slice_) <= static_cast<size_t>(INT_MAX));
    byte_count_ += * size = static_cast<int>(GRPC_SLICE_LENGTH(slice_));
    // Using grpc_slice_buffer_add could modify slice_ and merge it with the
    // previous slice. Therefore, use grpc_slice_buffer_add_indexed method to
    // ensure the slice gets added at a separate index. It can then be kept
    // around and popped later in the BackUp function.
    grpc_slice_buffer_add_indexed(slice_buffer_, slice_);
    return true;
  }
 
  void BackUp(int count) override {
    // count == 0 is invoked by ZeroCopyOutputStream users indicating that any
    // potential buffer obtained through a previous call to Next() is final.
    // ZeroCopyOutputStream implementations such as streaming output can use
    // these calls to flush any temporary buffer and flush the output. The logic
    // below is not robust against count == 0 invocations, so directly return.
    if (count == 0) return;
 
    ABSL_CHECK_LE(count, static_cast<int>(GRPC_SLICE_LENGTH(slice_)));
    grpc_slice_buffer_pop(slice_buffer_);
    if (static_cast<size_t>(count) == GRPC_SLICE_LENGTH(slice_)) {
      backup_slice_ = slice_;
    } else {
      backup_slice_ =
          grpc_slice_split_tail(&slice_, GRPC_SLICE_LENGTH(slice_) - count);
      grpc_slice_buffer_add(slice_buffer_, slice_);
    }
    // It's dangerous to keep an inlined grpc_slice as the backup slice, since
    // on a following Next() call, a reference will be returned to this slice
    // via GRPC_SLICE_START_PTR, which will not be an address held by
    // slice_buffer_.
    have_backup_ = backup_slice_.refcount != nullptr;
    byte_count_ -= count;
  }
 
  int64_t ByteCount() const override { return byte_count_; }
 
#ifdef GRPC_PROTOBUF_CORD_SUPPORT_ENABLED
  // (override is conditionally omitted here to support old Protobuf which
  //  doesn't have ReadCord method)
  // NOLINTBEGIN(modernize-use-override,
  // clang-diagnostic-inconsistent-missing-override)
  virtual bool WriteCord(const absl::Cord& cord)
#if GOOGLE_PROTOBUF_VERSION >= 4022000
      override
#endif
  // NOLINTEND(modernize-use-override,
  // clang-diagnostic-inconsistent-missing-override)
  {
    grpc_slice_buffer* buffer = slice_buffer();
    size_t cur = 0;
    for (absl::string_view chunk : cord.Chunks()) {
      // TODO(veblush): Revisit this 512 threadhold which could be smaller.
      if (chunk.size() < 512) {
        // If chunk is small enough, just copy it.
        grpc_slice slice =
            grpc_slice_from_copied_buffer(chunk.data(), chunk.size());
        grpc_slice_buffer_add(buffer, slice);
      } else {
        // If chunk is large, just use the pointer instead of copying.
        // To make sure it's alive while being used, a subcord for chunk is
        // created and attached to a grpc_slice instance.
        absl::Cord* subcord = new absl::Cord(cord.Subcord(cur, chunk.size()));
        grpc_slice slice = grpc_slice_new_with_user_data(
            const_cast<uint8_t*>(
                reinterpret_cast<const uint8_t*>(chunk.data())),
            chunk.size(), [](void* p) { delete static_cast<absl::Cord*>(p); },
            subcord);
        grpc_slice_buffer_add(buffer, slice);
      }
      cur += chunk.size();
    }
    set_byte_count(ByteCount() + cur);
    return true;
  }
#endif  // GRPC_PROTOBUF_CORD_SUPPORT_ENABLED
 
  // These protected members are needed to support internal optimizations.
  // they expose internal bits of grpc core that are NOT stable. If you have
  // a use case needs to use one of these functions, please send an email to
  // https://groups.google.com/forum/#!forum/grpc-io.
 protected:
  grpc_slice_buffer* slice_buffer() { return slice_buffer_; }
  void set_byte_count(int64_t byte_count) { byte_count_ = byte_count; }
 
 private:
  // friend for testing purposes only
  friend class internal::ProtoBufferWriterPeer;
  const int block_size_;  
  const int total_size_;  
  int64_t byte_count_;    
  grpc_event_engine::experimental::MemoryAllocator* allocator_ = nullptr;
  grpc_slice_buffer*
      slice_buffer_;  
  bool have_backup_;  
  grpc_slice backup_slice_;  
  grpc_slice slice_;         
};
 
}  // namespace grpc
 
#endif  // GRPCPP_SUPPORT_PROTO_BUFFER_WRITER_H