tls: improve write performance by reducing copying

include/envoy/buffer/buffer.h

@@ -29,6 +29,7 @@ struct RawSlice {
   size_t len_ = 0;
 
   bool operator==(const RawSlice& rhs) const { return mem_ == rhs.mem_ && len_ == rhs.len_; }
+  bool operator!=(const RawSlice& rhs) const { return !(*this == rhs); }
 };
 
 using RawSliceVector = absl::InlinedVector<RawSlice, 16>;
@@ -187,6 +188,16 @@ class Instance {
    */
   virtual void* linearize(uint32_t size) PURE;
 
+  /**
+   * Get a pointer to a linear chunk of this buffer. The chunk may be smaller than max_size, even if
+   * the length of the buffer is larger. The function will heuristically determine how much data to
+   * copy based on `desired_min_size`, in order to avoid patterns in which all the data is copied
+   * when it doesn't need to be. For example, if the buffer contains a slice containing 1 byte,
+   * followed by 100 slices containing ``max_size``, repeatedly calling this function would avoid
+   * repeatedly copying ``max_size - 1`` bytes to make chunks of ``max_size``.
+   */
+  virtual RawSlice maybeLinearize(uint32_t max_size, uint32_t desired_min_size) PURE;
+
   /**
    * Move a buffer into this buffer. As little copying is done as possible.
    * @param rhs supplies the buffer to move.

source/common/buffer/buffer_impl.cc

@@ -277,6 +277,30 @@ void* OwnedImpl::linearize(uint32_t size) {
   return slices_.front()->data();
 }
 
+RawSlice OwnedImpl::maybeLinearize(uint32_t max_size, uint32_t desired_min_size) {
+  while (!slices_.empty() && slices_[0]->dataSize() == 0) {
+    slices_.pop_front();
+  }
+
+  if (slices_.empty()) {
+    return {nullptr, 0};
+  }
+
+  const uint64_t slice_size = std::min<uint64_t>(slices_[0]->dataSize(), max_size);
+  if (slice_size >= desired_min_size) {
+    return {slices_[0]->data(), slice_size};
+  }
+
+  // The next slice will already be of the desired size, so don't copy and
+  // return the front slice.
+  if (slices_.size() >= 2 && slices_[1]->dataSize() >= max_size) {
+    return {slices_[0]->data(), slice_size};
+  }
+
+  auto size = std::min<size_t>(max_size, length_);
+  return {linearize(size), size};
+}
+
 void OwnedImpl::coalesceOrAddSlice(SlicePtr&& other_slice) {
   const uint64_t slice_size = other_slice->dataSize();
   // The `other_slice` content can be coalesced into the existing slice IFF:

source/common/buffer/buffer_impl.h

@@ -568,6 +568,7 @@ class OwnedImpl : public LibEventInstance {
   SliceDataPtr extractMutableFrontSlice() override;
   uint64_t length() const override;
   void* linearize(uint32_t size) override;
+  RawSlice maybeLinearize(uint32_t max_size, uint32_t desired_min_size) override;
   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;

source/extensions/transport_sockets/tls/context_impl.cc

@@ -84,6 +84,10 @@ ContextImpl::ContextImpl(Stats::Scope& scope, const Envoy::Ssl::ContextConfig& c
     int rc = SSL_CTX_set_app_data(ctx.ssl_ctx_.get(), this);
     RELEASE_ASSERT(rc == 1, Utility::getLastCryptoError().value_or(""));
 
+    constexpr uint32_t mode = SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER;
+    rc = SSL_CTX_set_mode(ctx.ssl_ctx_.get(), mode);
+    RELEASE_ASSERT((rc & mode) == mode, Utility::getLastCryptoError().value_or(""));
+
     rc = SSL_CTX_set_min_proto_version(ctx.ssl_ctx_.get(), config.minProtocolVersion());
     RELEASE_ASSERT(rc == 1, Utility::getLastCryptoError().value_or(""));
 

source/extensions/transport_sockets/tls/context_impl.h

@@ -45,7 +45,8 @@ namespace Tls {
   COUNTER(ocsp_staple_failed)                                                                      \
   COUNTER(ocsp_staple_omitted)                                                                     \
   COUNTER(ocsp_staple_responses)                                                                   \
-  COUNTER(ocsp_staple_requests)
+  COUNTER(ocsp_staple_requests)                                                                    \
+  HISTOGRAM(write_size, Bytes)
 
 /**
  * Wrapper struct for SSL stats. @see stats_macros.h

source/extensions/transport_sockets/tls/ssl_socket.cc

@@ -236,35 +236,27 @@ Network::IoResult SslSocket::doWrite(Buffer::Instance& write_buffer, bool end_st
     }
   }
 
-  uint64_t bytes_to_write;
-  if (bytes_to_retry_) {
-    bytes_to_write = bytes_to_retry_;
-    bytes_to_retry_ = 0;
-  } else {
-    bytes_to_write = std::min(write_buffer.length(), static_cast<uint64_t>(16384));
-  }
-
   uint64_t total_bytes_written = 0;
-  while (bytes_to_write > 0) {
+  while (true) {
     // TODO(mattklein123): As it relates to our fairness efforts, we might want to limit the number
     // of iterations of this loop, either by pure iterations, bytes written, etc.
+    const auto slice = write_buffer.maybeLinearize(16384, 4096);
+    if (slice.len_ == 0) {
+      break;
+    }
 
-    // SSL_write() requires that if a previous call returns SSL_ERROR_WANT_WRITE, we need to call
-    // it again with the same parameters. This is done by tracking last write size, but not write
-    // data, since linearize() will return the same undrained data anyway.
-    ASSERT(bytes_to_write <= write_buffer.length());
-    int rc = SSL_write(rawSsl(), write_buffer.linearize(bytes_to_write), bytes_to_write);
+    ASSERT(slice.mem_ != nullptr);
+    int rc = SSL_write(rawSsl(), slice.mem_, slice.len_);
     ENVOY_CONN_LOG(trace, "ssl write returns: {}", callbacks_->connection(), rc);
     if (rc > 0) {
-      ASSERT(rc == static_cast<int>(bytes_to_write));
+      ASSERT(rc == static_cast<int>(slice.len_));
+      ctx_->stats().write_size_.recordValue(rc);
       total_bytes_written += rc;
       write_buffer.drain(rc);
-      bytes_to_write = std::min(write_buffer.length(), static_cast<uint64_t>(16384));
     } else {
       int err = SSL_get_error(rawSsl(), rc);
       switch (err) {
       case SSL_ERROR_WANT_WRITE:
-        bytes_to_retry_ = bytes_to_write;
         break;
       case SSL_ERROR_WANT_READ:
       // Renegotiation has started. We don't handle renegotiation so just fall through.

source/extensions/transport_sockets/tls/ssl_socket.h

@@ -91,7 +91,6 @@ class SslSocket : public Network::TransportSocket,
   const Network::TransportSocketOptionsSharedPtr transport_socket_options_;
   Network::TransportSocketCallbacks* callbacks_{};
   ContextImplSharedPtr ctx_;
-  uint64_t bytes_to_retry_{};
   std::string failure_reason_;
 
   SslHandshakerImplSharedPtr info_;

test/common/buffer/buffer_fuzz.cc

@@ -135,6 +135,10 @@ class StringBuffer : public Buffer::Instance {
     return mutableStart();
   }
 
+  Buffer::RawSlice maybeLinearize(uint32_t max_size, uint32_t /*desired_min_size*/) override {
+    return {mutableStart(), std::min(size_, max_size)};
+  }
+
   Buffer::SliceDataPtr extractMutableFrontSlice() override { NOT_IMPLEMENTED_GCOVR_EXCL_LINE; }
 
   void move(Buffer::Instance& rhs) override { move(rhs, rhs.length()); }

test/common/buffer/owned_impl_test.cc

@@ -786,6 +786,60 @@ TEST_F(OwnedImplTest, LinearizeDrainTracking) {
   expectSlices({}, buffer);
 }
 
+TEST_F(OwnedImplTest, MaybeLinearizeEmpty) {
+  Buffer::OwnedImpl empty;
+  EXPECT_EQ(0, empty.maybeLinearize(1024, 1024).len_);
+}
+
+// Test that the correct value is returned in both the case where
+// the slice has a larger and a smaller length than `desired_min_size`.
+TEST_F(OwnedImplTest, MaybeLinearizeSingleSlice) {
+  Buffer::OwnedImpl buffer;
+  buffer.add(std::string(100, 'a'));
+  EXPECT_EQ(100, buffer.maybeLinearize(1024, 512).len_);
+  EXPECT_EQ(100, buffer.maybeLinearize(1024, 1).len_);
+}
+
+TEST_F(OwnedImplTest, MaybeLinearizeDesiredMinSize) {
+  Buffer::OwnedImpl buffer;
+  buffer.add(std::string(10000, 'a'));
+  Buffer::OwnedImpl other;
+  other.add(std::string(10000, 'b'));
+  buffer.move(other);
+
+  // Verify test slices are as expected
+  const auto slices = buffer.getRawSlices();
+  ASSERT_EQ(2, slices.size());
+  ASSERT_EQ(10000, slices[0].len_);
+  ASSERT_EQ(10000, slices[1].len_);
+
+  // Ask for the entire buffer size. This should return only the first slice because
+  // `desired_min_size` is less than the size of that slice.
+  EXPECT_EQ(slices[0], buffer.maybeLinearize(20000, 9999));
+
+  // Ask for the entire buffer size, but with a desired_min_size greater than the first
+  // slice. This should get fully linearized into a single slice.
+  EXPECT_EQ(20000, buffer.maybeLinearize(20000, 10001).len_);
+}
+
+// Test that a smaller slice than `desired_min_size` is returned if the next slice
+// after it is full-sized.
+TEST_F(OwnedImplTest, MaybeLinearizePreferNextSlice) {
+  Buffer::OwnedImpl buffer;
+  buffer.add("a");
+  Buffer::OwnedImpl other;
+  other.add(std::string(10000, 'b'));
+  buffer.move(other);
+
+  // Verify test slices are as expected
+  const auto slices = buffer.getRawSlices();
+  ASSERT_EQ(2, slices.size());
+  ASSERT_EQ(1, slices[0].len_);
+  ASSERT_EQ(10000, slices[1].len_);
+
+  EXPECT_EQ(1, buffer.maybeLinearize(10000, 1024).len_);
+}
+
 TEST_F(OwnedImplTest, ReserveCommit) {
   // This fragment will later be added to the buffer. It is declared in an enclosing scope to
   // ensure it is not destructed until after the buffer is.

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

@@ -527,6 +527,8 @@ class FakeBuffer : public Buffer::Instance {
   MOCK_METHOD(Buffer::SliceDataPtr, extractMutableFrontSlice, (), (override));
   MOCK_METHOD(uint64_t, length, (), (const, override));
   MOCK_METHOD(void*, linearize, (uint32_t), (override));
+  MOCK_METHOD(Buffer::RawSlice, maybeLinearize, (uint32_t max_size, uint32_t desired_min_size),
+              (override));
   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));

test/extensions/transport_sockets/tls/BUILD

@@ -1,5 +1,7 @@
 load(
     "//bazel:envoy_build_system.bzl",
+    "envoy_benchmark_test",
+    "envoy_cc_benchmark_binary",
     "envoy_cc_test",
     "envoy_cc_test_library",
     "envoy_package",
@@ -179,3 +181,23 @@ envoy_cc_test(
         "//test/mocks/stats:stats_mocks",
     ],
 )
+
+envoy_cc_benchmark_binary(
+    name = "tls_throughput_benchmark",
+    srcs = ["tls_throughput_test.cc"],
+    data = [
+        "//test/extensions/transport_sockets/tls/test_data:certs",
+    ],
+    external_deps = [
+        "benchmark",
+        "ssl",
+    ],
+    deps = [
+        "//source/common/buffer:buffer_lib",
+    ],
+)
+
+envoy_benchmark_test(
+    name = "tls_throughput_benchmark_test",
+    benchmark_binary = "tls_throughput_benchmark",
+)