chore: write path for TlsSocket::AsyncReadSome (#415)

This pr implements the asynchronous write path (when engine state NEED_WRITE) for TlsSocket::AsyncReadSome.

* add async write
* add test

Signed-off-by: kostas <[email protected]>
Signed-off-by: Roman Gershman <[email protected]>

Author: kostasrim

.github/workflows/mac-os.yml

@@ -45,4 +45,4 @@ jobs:
           ninja -k 5 base/all io/all strings/all util/all echo_server ping_iouring_server \
             https_client_cli s3_demo
           ./fibers_test --logtostderr --gtest_repeat=10
-          ctest -V -L CI
+          GLOG_logtostderr=1 ctest -V -L CI

io/io.cc

@@ -50,6 +50,7 @@ struct AsyncReadState {
 
   AsyncReadState(AsyncSource* source, const iovec* v, uint32_t length)
       : arr(length), owner(source) {
+    cur = arr.data();
     std::copy(v, v + length, arr.data());
   }
 

util/fibers/epoll_socket.cc

@@ -386,6 +386,7 @@ void EpollSocket::AsyncWriteSome(const iovec* v, uint32_t len, io::AsyncProgress
   async_write_pending_ = 1;
 }
 
+// TODO implement async functionality
 void EpollSocket::AsyncReadSome(const iovec* v, uint32_t len, io::AsyncProgressCb cb) {
   auto res = ReadSome(v, len);
   cb(res);

util/fibers/listener_interface.cc

@@ -98,7 +98,7 @@ void ListenerInterface::RunAcceptLoop() {
 
   if (!sock_->IsUDS()) {
     ep = sock_->LocalEndpoint();
-    VSOCK(0, *sock_) << "AcceptServer - listening on port " << ep.port();
+    VSOCK(0, *sock_) << "AcceptServer - listening on " << ep.address() << ":" << ep.port();
   }
 
   PreAcceptLoop(sock_->proactor());

util/tls/tls_socket.cc

@@ -516,9 +516,11 @@ void TlsSocket::SetProactor(ProactorBase* p) {
 
 void TlsSocket::AsyncReq::MaybeSendOutputAsyncWithRead() {
   if (owner->engine_->OutputPending() != 0) {
-    // sync interface, works because we are still executing within a fiber
-    // used for "mocking" and shall be replaced on the next PR with actual async op
-    owner->MaybeSendOutput();
+    // Once the networking socket completes the write, it will start the read path
+    // We use this bool to signal this.
+    should_read = true;
+    StartUpstreamWrite();
+    return;
   }
 
   // TODO handle WRITE_IN_PROGRESS here by adding pending_blocked_
@@ -535,16 +537,17 @@ void TlsSocket::AsyncReq::AsyncProgressCb(io::Result<size_t> read_result) {
       VLOG(1) << "sock[" << owner->native_handle() << "], state " << int(owner->state_)
               << ", write_total:" << owner->upstream_write_ << " "
               << " pending output: " << owner->engine_->OutputPending() << " "
-              << "StartUpstreamRead failed " << read_result.error();
+              << "AsyncProgressCb failed " << read_result.error();
     }
     // Erronous path. Apply the completion callback and exit.
     CompleteAsyncReq(read_result);
     return;
   }
 
-  DVLOG(1) << "HandleUpstreamRead " << *read_result << " bytes";
+  DVLOG(1) << "AsyncProgressCb " << *read_result << " bytes";
   owner->engine_->CommitInput(*read_result);
-  Engine::OpResult engine_read = owner->MaybeReadFromEngine(vec, len);
+  Engine::OpResult engine_read =
+      owner->engine_->Read(reinterpret_cast<uint8_t*>(vec->iov_base), vec->iov_len);
   if (engine_read > 0) {
     CompleteAsyncReq(engine_read);
     return;
@@ -579,42 +582,21 @@ void TlsSocket::AsyncReq::HandleOpAsync() {
     case Engine::NEED_READ_AND_MAYBE_WRITE:
       MaybeSendOutputAsyncWithRead();
       break;
-    // TODO handle NEED_WRITE
+    case Engine::NEED_WRITE:
+      MaybeSendOutputAsync();
+      break;
     default:
       // EOF_STREAM should be handled earlier
       LOG(DFATAL) << "Unsupported " << op_val;
   }
 }
 
-Engine::OpResult TlsSocket::MaybeReadFromEngine(const iovec* v, uint32_t len) {
-  size_t read_len = std::min(v->iov_len, size_t(INT_MAX));
-  Engine::OpResult op_val = engine_->Read(reinterpret_cast<uint8_t*>(v->iov_base), read_len);
-  DVLOG(2) << "Engine::Read " << read_len << " bytes, got " << op_val;
-  // if read_len == op_val we could try to read more. However, the next read might require
-  // an async operation on the underline socket because op_val < 0.
-  // The problem here is that SSL_read from engine_->Read is *not* idempotent and we might
-  // end up in a situation where we need to do two things at the same time:
-  // 1. Call the callers completion callback which will start another async op because
-  //    we read less bytes than what was requested, i.e, read_total < sum_of_all(v->len).
-  // 2. Start another async operation to satisfy the protocol because op_val < 0 and we
-  //    called engine_->Read which is *not* idempotent.
-  // For that, it's best to let it flow naturally. If there is some data in the engine read it
-  // and call the completion callback which will in turn try to read more from the engine.
-  // It will read everything or reach to a point that an async operation needs to be dispatched.
-  // That way, we get a linear view of the operations involved with the downside of a few more
-  // function calls (since we don't try to drain the whole engine as we don't know if the next
-  // read can be satisfied or dispatch as an async operation).
-  // Last but not least, it was advised here:
-  // https://github.com/romange/helio/pull/408#discussion_r2080998216
-  // That we should remove engine reads from the AsyncRequest all together and return
-  // to the caller if there was some data read.
-  return op_val;
-}
-
 void TlsSocket::AsyncReadSome(const iovec* v, uint32_t len, io::AsyncProgressCb cb) {
+  // Engine read
   CHECK(!async_read_req_);
 
-  Engine::OpResult op_val = MaybeReadFromEngine(v, len);
+  Engine::OpResult op_val = engine_->Read(reinterpret_cast<uint8_t*>(v->iov_base), v->iov_len);
+  DVLOG(2) << "Engine::Read tried to read " << v->iov_len << " bytes, got " << op_val;
   // We read some data from the engine. Satisfy the request and return.
   if (op_val > 0) {
     return cb(op_val);
@@ -632,5 +614,80 @@ void TlsSocket::AsyncReadSome(const iovec* v, uint32_t len, io::AsyncProgressCb
   async_read_req_->HandleOpAsync();
 }
 
+void TlsSocket::AsyncReq::CompleteAsyncWrite(io::Result<size_t> write_result) {
+  if (!write_result) {
+    owner->state_ &= ~WRITE_IN_PROGRESS;
+
+    // broken_pipe - happens when the other side closes the connection. do not log this.
+    if (write_result.error() != errc::broken_pipe) {
+      VLOG(1) << "sock[" << owner->native_handle() << "], state " << int(owner->state_)
+              << ", write_total:" << owner->upstream_write_ << " "
+              << " pending output: " << owner->engine_->OutputPending()
+              << " HandleUpstreamAsyncWrite failed " << write_result.error();
+    }
+
+    // We are done. Errornous exit.
+    CompleteAsyncReq(write_result);
+    return;
+  }
+
+  CHECK_GT(*write_result, 0u);
+  owner->upstream_write_ += *write_result;
+  owner->engine_->ConsumeOutputBuf(*write_result);
+  // We might have more data pending. Peek again.
+  Buffer buffer = owner->engine_->PeekOutputBuf();
+
+  // We are not done. Re-arm the async write until we drive it to completion or error.
+  // We would also like to avoid fragmented socket writes so we make sure we drain it here
+  if (!buffer.empty()) {
+    auto& scratch = scratch_iovec;
+    scratch.iov_base = const_cast<uint8_t*>(buffer.data());
+    scratch.iov_len = buffer.size();
+    owner->next_sock_->AsyncWriteSome(
+        &scratch, 1, [this](auto write_result) { CompleteAsyncWrite(write_result); });
+    return;
+  }
+
+  if (owner->engine_->OutputPending() > 0) {
+    LOG(DFATAL) << "ssl buffer is not empty with " << owner->engine_->OutputPending()
+                << " bytes. Async short write detected";
+  }
+
+  owner->state_ &= ~WRITE_IN_PROGRESS;
+
+  // We are done with the writes, check if we also need to read because we are
+  // in NEED_READ_AND_MAYBE_WRITE state
+  if (should_read) {
+    should_read = false;
+    StartUpstreamRead();
+  }
+}
+
+void TlsSocket::AsyncReq::StartUpstreamWrite() {
+  Engine::Buffer buffer = owner->engine_->PeekOutputBuf();
+  DCHECK(!buffer.empty());
+  DCHECK((owner->state_ & WRITE_IN_PROGRESS) == 0);
+
+  DVLOG(2) << "StartUpstreamWrite " << buffer.size();
+  // we do not allow concurrent writes from multiple fibers.
+  owner->state_ |= WRITE_IN_PROGRESS;
+
+  auto& scratch = scratch_iovec;
+  scratch.iov_base = const_cast<uint8_t*>(buffer.data());
+  scratch.iov_len = buffer.size();
+
+  owner->next_sock_->AsyncWriteSome(
+      &scratch, 1, [this](auto write_result) { CompleteAsyncWrite(write_result); });
+}
+
+void TlsSocket::AsyncReq::MaybeSendOutputAsync() {
+  if (owner->engine_->OutputPending() == 0) {
+    return;
+  }
+
+  // TODO handle WRITE_IN_PROGRESS to avoid deadlock
+  StartUpstreamWrite();
+}
+
 }  // namespace tls
 }  // namespace util

util/tls/tls_socket.h

@@ -116,6 +116,7 @@ class TlsSocket final : public FiberSocketBase {
   enum { WRITE_IN_PROGRESS = 1, READ_IN_PROGRESS = 2, SHUTDOWN_IN_PROGRESS = 4, SHUTDOWN_DONE = 8 };
   uint8_t state_{0};
 
+  // TODO turn this into a class with proper access specifiers
   struct AsyncReq {
     TlsSocket* owner;
     // Callback passed from the user.
@@ -127,26 +128,23 @@ class TlsSocket final : public FiberSocketBase {
 
     iovec scratch_iovec;
 
+    bool should_read = false;
+
     // Asynchronous helpers
     void MaybeSendOutputAsyncWithRead();
+    void MaybeSendOutputAsync();
 
     void HandleOpAsync();
 
     void StartUpstreamRead();
+    void StartUpstreamWrite();
 
     void CompleteAsyncReq(io::Result<size_t> result);
+    void CompleteAsyncWrite(io::Result<size_t> write_result);
 
     void AsyncProgressCb(io::Result<size_t> result);
   };
 
-  // Helper function that resets the internal async request, applies the
-  // user AsyncProgressCb and returns. We need this, because progress callbacks
-  // can start another async request and for that to work, we need to clean up
-  // the one we are running on.
-  void CompleteAsyncRequest(io::Result<size_t> result);
-
-  Engine::OpResult MaybeReadFromEngine(const iovec* v, uint32_t len);
-
   std::unique_ptr<AsyncReq> async_read_req_;
 };
 

util/tls/tls_socket_test.cc

@@ -227,15 +227,17 @@ class AsyncTlsSocketTest : public testing::TestWithParam<string_view> {
 
   using IoResult = int;
 
-  // TODO clean up
   virtual void HandleRequest() {
     tls_socket_ = std::make_unique<tls::TlsSocket>(conn_socket_.release());
     ssl_ctx_ = CreateSslCntx(SERVER);
     tls_socket_->InitSSL(ssl_ctx_);
     tls_socket_->Accept();
 
     uint8_t buf[16];
-    std::fill(std::begin(buf), std::end(buf), uint8_t(120));
+    auto res = tls_socket_->Recv(buf);
+    EXPECT_TRUE(res.has_value());
+    EXPECT_TRUE(res.value() == 16);
+
     auto write_res = tls_socket_->Write(buf);
     EXPECT_FALSE(write_res);
   }
@@ -345,11 +347,22 @@ TEST_P(AsyncTlsSocketTest, AsyncRW) {
   proactor_->Await([&] {
     ThisFiber::SetName("ConnectFb");
 
-    LOG(INFO) << "Connecting to " << listen_ep_;
     error_code ec = tls_sock->Connect(listen_ep_);
     EXPECT_FALSE(ec);
     uint8_t res[16];
     std::fill(std::begin(res), std::end(res), uint8_t(120));
+    {
+      Done done;
+      iovec v{.iov_base = &res, .iov_len = 16};
+
+      tls_sock->AsyncWriteSome(&v, 1, [done](auto result) mutable {
+        EXPECT_TRUE(result.has_value());
+        EXPECT_EQ(*result, 16);
+        done.Notify();
+      });
+
+      done.Wait();
+    }
     {
       uint8_t buf[16];
       Done done;