buffer: improve read reservations to efficiently handle multiple slices

include/envoy/buffer/buffer.h

@@ -71,6 +71,8 @@ class SliceData {
 
 using SliceDataPtr = std::unique_ptr<SliceData>;
 
+class Reservation;
+
 /**
  * A basic buffer abstraction.
  */
@@ -201,6 +203,9 @@ class Instance {
    */
   virtual uint64_t reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iovecs) PURE;
 
+  virtual Reservation reserve(uint64_t preferred_length) PURE;
+  virtual void commit(Reservation& reservation, uint64_t length) PURE;
+
   /**
    * Search for an occurrence of data within the buffer.
    * @param data supplies the data to search for.
@@ -411,5 +416,36 @@ class WatermarkFactory {
 
 using WatermarkFactoryPtr = std::unique_ptr<WatermarkFactory>;
 
+/**
+ * Holds an in-progress addition to a buffer.
+ *
+ * @note For performance reasons, this class is passed by value to
+ * avoid an extra allocation, so it cannot have any virtual methods.
+ */
+class Reservation final {
+public:
+  Reservation(Reservation&&) = default;
+  ~Reservation() {
+    if (!slices_.empty()) {
+      buffer_.commit(*this, 0);
+    }
+  }
+  RawSlice* slices() { return slices_.data(); }
+  uint64_t numSlices() const { return slices_.size(); }
+  void commit(uint64_t length) {
+    buffer_.commit(*this, length);
+    slices_.clear();
+    owned_slices_.clear();
+  }
+
+  // private:
+  friend class Instance;
+  Reservation(Instance& buffer) : buffer_(buffer) {}
+  Instance& buffer_;
+  static constexpr uint32_t num_elements_ = 9;
+  absl::InlinedVector<RawSlice, num_elements_> slices_;
+  absl::InlinedVector<SliceDataPtr, num_elements_> owned_slices_;
+};
+
 } // namespace Buffer
 } // namespace Envoy

source/common/buffer/buffer_impl.cc

@@ -18,6 +18,9 @@ namespace {
 constexpr uint64_t CopyThreshold = 512;
 } // namespace
 
+thread_local absl::InlinedVector<std::unique_ptr<OwnedSlice>, OwnedSlice::free_list_max_>
+    OwnedSlice::free_list_;
+
 void OwnedImpl::addImpl(const void* data, uint64_t size) {
   const char* src = static_cast<const char*>(data);
   bool new_slice_needed = slices_.empty();
@@ -378,6 +381,69 @@ uint64_t OwnedImpl::reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iove
   return num_slices_used;
 }
 
+Reservation OwnedImpl::reserve(uint64_t length) {
+  Reservation reservation(*this);
+  if (length == 0) {
+    return reservation;
+  }
+
+  // Remove any empty slices at the end.
+  while (!slices_.empty() && slices_.back()->dataSize() == 0) {
+    slices_.pop_back();
+  }
+
+  uint64_t bytes_remaining = length;
+
+  // Check whether there are any empty slices with reservable space at the end of the buffer.
+  uint64_t reservable_size = slices_.empty() ? 0 : slices_.back()->reservableSize();
+  if (reservable_size >= length || reservable_size >= (OwnedSlice::default_slice_size_ / 8)) {
+    auto& last_slice = slices_.back();
+    const uint64_t reservation_size = std::min(last_slice->reservableSize(), bytes_remaining);
+    auto slice = last_slice->reserve(reservation_size);
+    reservation.slices_.push_back(slice);
+    reservation.owned_slices_.push_back(nullptr);
+    bytes_remaining -= slice.len_;
+  }
+
+  while (bytes_remaining != 0) {
+    const uint64_t size = OwnedSlice::default_slice_size_;
+    auto slice = OwnedSlice::create(size);
+    reservation.slices_.push_back(slice->reserve(size));
+    reservation.owned_slices_.emplace_back(std::move(slice));
+    bytes_remaining -= std::min(reservation.slices_.back().len_, bytes_remaining);
+  }
+
+  ASSERT(reservation.slices_.size() <= reservation.num_elements_,
+         "The absl::InlineVector should be sized so that out-of-line storage isn't needed.");
+  ASSERT(reservation.slices_.size() == reservation.owned_slices_.size());
+  return reservation;
+}
+
+void OwnedImpl::commit(Reservation& reservation, uint64_t length) {
+  ASSERT(reservation.slices_.size() == reservation.owned_slices_.size());
+  uint64_t bytes_remaining = length;
+  for (uint32_t i = 0; i < reservation.slices_.size(); i++) {
+    ASSERT((reservation.owned_slices_[i] != nullptr) ==
+           (dynamic_cast<OwnedSlice*>(reservation.owned_slices_[i].get()) != nullptr));
+    std::unique_ptr<OwnedSlice> owned_slice(
+        static_cast<OwnedSlice*>(reservation.owned_slices_[i].release()));
+
+    if (bytes_remaining > 0) {
+      if (owned_slice != nullptr) {
+        slices_.emplace_back(std::move(owned_slice));
+      }
+      reservation.slices_[i].len_ =
+          std::min<uint64_t>(reservation.slices_[i].len_, bytes_remaining);
+      bool success = slices_.back()->commit(reservation.slices_[i]);
+      ASSERT(success);
+      length_ += reservation.slices_[i].len_;
+      bytes_remaining -= reservation.slices_[i].len_;
+    } else {
+      OwnedSlice::free(std::move(owned_slice));
+    }
+  }
+}
+
 ssize_t OwnedImpl::search(const void* data, uint64_t size, size_t start, size_t length) const {
   // This implementation uses the same search algorithm as evbuffer_search(), a naive
   // scan that requires O(M*N) comparisons in the worst case.

source/common/buffer/buffer_impl.h

@@ -261,9 +261,18 @@ class OwnedSlice final : public Slice, public InlineStorage {
    * @param capacity number of bytes of space the slice should have.
    * @return an OwnedSlice with at least the specified capacity.
    */
-  static SlicePtr create(uint64_t capacity) {
+  static std::unique_ptr<OwnedSlice> create(uint64_t capacity) {
+    ASSERT(sliceSize(default_slice_size_) == default_slice_size_);
+
     uint64_t slice_capacity = sliceSize(capacity);
-    return SlicePtr(new (slice_capacity) OwnedSlice(slice_capacity));
+    std::unique_ptr<OwnedSlice> slice;
+    if (slice_capacity == default_slice_size_ && !free_list_.empty()) {
+      slice = std::move(free_list_.back());
+      free_list_.pop_back();
+    } else {
+      slice.reset(new (slice_capacity) OwnedSlice(slice_capacity));
+    }
+    return slice;
   }
 
   /**
@@ -274,13 +283,28 @@ class OwnedSlice final : public Slice, public InlineStorage {
    *         the internal implementation) have a nonzero amount of reservable space at the end.
    */
   static SlicePtr create(const void* data, uint64_t size) {
-    uint64_t slice_capacity = sliceSize(size);
-    std::unique_ptr<OwnedSlice> slice(new (slice_capacity) OwnedSlice(slice_capacity));
+    std::unique_ptr<OwnedSlice> slice(create(size));
     memcpy(slice->base_, data, size);
     slice->reservable_ = size;
     return slice;
   }
 
+  static constexpr uint32_t default_slice_size_ = 16384;
+
+  static void free(std::unique_ptr<OwnedSlice> slice) {
+    if (slice == nullptr) {
+      return;
+    }
+
+    if (slice->capacity_ == default_slice_size_ && slice->reservable_ == 0 &&
+        free_list_.size() < free_list_max_) {
+      free_list_.emplace_back(std::move(slice));
+      ASSERT(slice == nullptr);
+    } else {
+      slice.reset();
+    }
+  }
+
 private:
   OwnedSlice(uint64_t size) : Slice(0, 0, size) { base_ = storage_; }
 
@@ -291,12 +315,15 @@ class OwnedSlice final : public Slice, public InlineStorage {
    * @param data_size the minimum amount of data the slice must be able to store, in bytes.
    * @return a recommended slice size, in bytes.
    */
-  static uint64_t sliceSize(uint64_t data_size) {
-    static constexpr uint64_t PageSize = 4096;
-    const uint64_t num_pages = (sizeof(OwnedSlice) + data_size + PageSize - 1) / PageSize;
-    return num_pages * PageSize - sizeof(OwnedSlice);
+  constexpr static uint64_t sliceSize(uint64_t data_size) {
+    constexpr uint64_t PageSize = 4096;
+    const uint64_t num_pages = (data_size + PageSize - 1) / PageSize;
+    return num_pages * PageSize;
   }
 
+  static constexpr uint32_t free_list_max_ = 8;
+  static thread_local absl::InlinedVector<std::unique_ptr<OwnedSlice>, free_list_max_> free_list_;
+
   uint8_t storage_[];
 };
 
@@ -561,6 +588,7 @@ class OwnedImpl : public LibEventInstance {
   void prepend(absl::string_view data) override;
   void prepend(Instance& data) override;
   void commit(RawSlice* iovecs, uint64_t num_iovecs) override;
+  void commit(Reservation& reservation, uint64_t length) override;
   void copyOut(size_t start, uint64_t size, void* data) const override;
   void drain(uint64_t size) override;
   RawSliceVector getRawSlices(absl::optional<uint64_t> max_slices = absl::nullopt) const override;
@@ -570,6 +598,7 @@ class OwnedImpl : public LibEventInstance {
   void move(Instance& rhs) override;
   void move(Instance& rhs, uint64_t length) override;
   uint64_t reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iovecs) override;
+  Reservation reserve(uint64_t preferred_length) override;
   ssize_t search(const void* data, uint64_t size, size_t start, size_t length) const override;
   bool startsWith(absl::string_view data) const override;
   std::string toString() const override;

source/common/buffer/watermark_buffer.cc

@@ -36,6 +36,11 @@ void WatermarkBuffer::commit(RawSlice* iovecs, uint64_t num_iovecs) {
   checkHighAndOverflowWatermarks();
 }
 
+void WatermarkBuffer::commit(Reservation& reservation, uint64_t length) {
+  OwnedImpl::commit(reservation, length);
+  checkHighAndOverflowWatermarks();
+}
+
 void WatermarkBuffer::drain(uint64_t size) {
   OwnedImpl::drain(size);
   checkLowWatermark();

source/common/buffer/watermark_buffer.h

@@ -36,6 +36,7 @@ class WatermarkBuffer : public OwnedImpl {
   void move(Instance& rhs, uint64_t length) override;
   SliceDataPtr extractMutableFrontSlice() override;
   uint64_t reserve(uint64_t length, RawSlice* iovecs, uint64_t num_iovecs) override;
+  void commit(Reservation& reservation, uint64_t length) override;
   void postProcess() override { checkLowWatermark(); }
   void appendSliceForTest(const void* data, uint64_t size) override;
   void appendSliceForTest(absl::string_view data) override;

source/common/network/io_socket_handle_impl.cc

@@ -92,17 +92,11 @@ Api::IoCallUint64Result IoSocketHandleImpl::read(Buffer::Instance& buffer, uint6
   if (max_length == 0) {
     return Api::ioCallUint64ResultNoError();
   }
-  constexpr uint64_t MaxSlices = 2;
-  Buffer::RawSlice slices[MaxSlices];
-  const uint64_t num_slices = buffer.reserve(max_length, slices, MaxSlices);
-  Api::IoCallUint64Result result = readv(max_length, slices, num_slices);
+  Buffer::Reservation reservation = buffer.reserve(max_length);
+  Api::IoCallUint64Result result = readv(max_length, reservation.slices(), reservation.numSlices());
   uint64_t bytes_to_commit = result.ok() ? result.rc_ : 0;
   ASSERT(bytes_to_commit <= max_length);
-  for (uint64_t i = 0; i < num_slices; i++) {
-    slices[i].len_ = std::min(slices[i].len_, static_cast<size_t>(bytes_to_commit));
-    bytes_to_commit -= slices[i].len_;
-  }
-  buffer.commit(slices, num_slices);
+  reservation.commit(bytes_to_commit);
   return result;
 }
 

source/common/network/raw_buffer_socket.cc

@@ -19,7 +19,7 @@ IoResult RawBufferSocket::doRead(Buffer::Instance& buffer) {
   bool end_stream = false;
   do {
     // 16K read is arbitrary. TODO(mattklein123) PERF: Tune the read size.
-    Api::IoCallUint64Result result = callbacks_->ioHandle().read(buffer, 16384);
+    Api::IoCallUint64Result result = callbacks_->ioHandle().read(buffer, 131072);
 
     if (result.ok()) {
       ENVOY_CONN_LOG(trace, "read returns: {}", callbacks_->connection(), result.rc_);

source/extensions/transport_sockets/tls/ssl_socket.cc

@@ -93,16 +93,13 @@ SslSocket::ReadResult SslSocket::sslReadIntoSlice(Buffer::RawSlice& slice) {
       ASSERT(static_cast<size_t>(rc) <= remaining);
       mem += rc;
       remaining -= rc;
-      result.commit_slice_ = true;
+      result.bytes_read_ += rc;
     } else {
       result.error_ = absl::make_optional<int>(rc);
       break;
     }
   }
 
-  if (result.commit_slice_) {
-    slice.len_ -= remaining;
-  }
   return result;
 }
 
@@ -122,17 +119,16 @@ Network::IoResult SslSocket::doRead(Buffer::Instance& read_buffer) {
   PostIoAction action = PostIoAction::KeepOpen;
   uint64_t bytes_read = 0;
   while (keep_reading) {
+    uint64_t bytes_read_this_iteration = 0;
     // We use 2 slices here so that we can use the remainder of an existing buffer chain element
     // if there is extra space. 16K read is arbitrary and can be tuned later.
-    Buffer::RawSlice slices[2];
-    uint64_t slices_to_commit = 0;
-    uint64_t num_slices = read_buffer.reserve(16384, slices, 2);
-    for (uint64_t i = 0; i < num_slices; i++) {
-      auto result = sslReadIntoSlice(slices[i]);
-      if (result.commit_slice_) {
-        slices_to_commit++;
-        bytes_read += slices[i].len_;
-      }
+    Buffer::Reservation reservation = read_buffer.reserve(16384);
+    // Buffer::RawSlice slices[2];
+    // uint64_t slices_to_commit = 0;
+    // uint64_t num_slices = read_buffer.reserve(16384, slices, 2);
+    for (uint64_t i = 0; i < reservation.numSlices(); i++) {
+      auto result = sslReadIntoSlice(reservation.slices()[i]);
+      bytes_read_this_iteration += result.bytes_read_;
       if (result.error_.has_value()) {
         keep_reading = false;
         int err = SSL_get_error(rawSsl(), result.error_.value());
@@ -162,13 +158,13 @@ Network::IoResult SslSocket::doRead(Buffer::Instance& read_buffer) {
       }
     }
 
-    if (slices_to_commit > 0) {
-      read_buffer.commit(slices, slices_to_commit);
-      if (callbacks_->shouldDrainReadBuffer()) {
-        callbacks_->setReadBufferReady();
-        keep_reading = false;
-      }
+    reservation.commit(bytes_read_this_iteration);
+    if (bytes_read_this_iteration > 0 && callbacks_->shouldDrainReadBuffer()) {
+      callbacks_->setReadBufferReady();
+      keep_reading = false;
     }
+
+    bytes_read += bytes_read_this_iteration;
   }
 
   ENVOY_CONN_LOG(trace, "ssl read {} bytes", callbacks_->connection(), bytes_read);

source/extensions/transport_sockets/tls/ssl_socket.h

@@ -77,7 +77,7 @@ class SslSocket : public Network::TransportSocket,
 
 private:
   struct ReadResult {
-    bool commit_slice_{};
+    uint64_t bytes_read_{0};
     absl::optional<int> error_;
   };
   ReadResult sslReadIntoSlice(Buffer::RawSlice& slice);

test/common/buffer/buffer_fuzz.cc

@@ -110,6 +110,11 @@ class StringBuffer : public Buffer::Instance {
     size_ += iovecs[0].len_;
   }
 
+  void commit(Buffer::Reservation&, uint64_t length) override {
+    size_ += length;
+    FUZZ_ASSERT(start_ + size_ + length <= data_.size());
+  }
+
   void copyOut(size_t start, uint64_t size, void* data) const override {
     ::memcpy(data, this->start() + start, size);
   }
@@ -152,6 +157,15 @@ class StringBuffer : public Buffer::Instance {
     return 1;
   }
 
+  Buffer::Reservation reserve(uint64_t length) override {
+    Buffer::RawSlice slice;
+    reserve(length, &slice, 1);
+
+    Buffer::Reservation reservation(*this);
+    reservation.slices_.push_back(slice);
+    return reservation;
+  }
+
   ssize_t search(const void* data, uint64_t size, size_t start, size_t length) const override {
     UNREFERENCED_PARAMETER(length);
     return asStringView().find({static_cast<const char*>(data), size}, start);

test/common/network/connection_impl_test.cc

@@ -2618,7 +2618,7 @@ TEST_P(ReadBufferLimitTest, SomeLimit) {
   // Envoy has soft limits, so as long as the first read is <= read_buffer_limit - 1 it will do a
   // second read. The effective chunk size is then read_buffer_limit - 1 + MaxReadSize,
   // which is currently 16384.
-  readBufferLimitTest(read_buffer_limit, read_buffer_limit - 1 + 16384);
+  readBufferLimitTest(read_buffer_limit, read_buffer_limit - 1 + (8 * 16384));
 }
 
 class TcpClientConnectionImplTest : public testing::TestWithParam<Address::IpVersion> {

test/extensions/filters/network/postgres_proxy/postgres_decoder_test.cc

@@ -529,6 +529,8 @@ class FakeBuffer : public Buffer::Instance {
   MOCK_METHOD(void, move, (Instance&), (override));
   MOCK_METHOD(void, move, (Instance&, uint64_t), (override));
   MOCK_METHOD(uint64_t, reserve, (uint64_t, Buffer::RawSlice*, uint64_t), (override));
+  MOCK_METHOD(Buffer::Reservation, reserve, (uint64_t preferred_length), (override));
+  MOCK_METHOD(void, commit, (Buffer::Reservation & reservation, uint64_t length), (override));
   MOCK_METHOD(ssize_t, search, (const void*, uint64_t, size_t, size_t), (const, override));
   MOCK_METHOD(bool, startsWith, (absl::string_view), (const, override));
   MOCK_METHOD(std::string, toString, (), (const, override));