chore: Introduce multishot-recv functionality.

util/fiber_socket_base.h

@@ -61,8 +61,14 @@ class FiberSocketBase : public io::Sink, public io::AsyncSink, public io::Source
   struct ProvidedBuffer {
     io::Bytes buffer;
     uint32_t allocated;
-    uint16_t err_no;
-    uint8_t cookie;  // Used by the socket to identify the buffer.
+    uint16_t err_no;   // Relevant only if buffer is empty.
+    uint8_t cookie;    // Used by the socket to identify the buffer source.
+
+    void SetError(uint16_t err) {
+      err_no = err;
+      allocated = 0;
+      buffer = {};
+    }
   };
 
   // Unlike Recv/ReadSome, this method returns buffers managed by the socket.

util/fibers/fiber_socket_test.cc

@@ -365,7 +365,7 @@ TEST_P(FiberSocketTest, RecvProvided) {
   });
   ASSERT_FALSE(ec);
 
-  io::Result<unsigned> res;
+  unsigned res;
   FiberSocketBase::ProvidedBuffer pbuf[8];
 
   auto recv_fb = proactor_->LaunchFiber([&] {
@@ -388,24 +388,91 @@ TEST_P(FiberSocketTest, RecvProvided) {
 
   recv_fb.Join();
   proactor_->Await([&] { std::ignore = sock->Close(); });
-  ASSERT_TRUE(res);
 
-  ASSERT_TRUE(*res > 0 && *res < 8);
+  ASSERT_TRUE(res > 0 && res < 8);
   size_t total_size = 0;
-  for (unsigned i = 0; i < *res; ++i) {
+  for (unsigned i = 0; i < res; ++i) {
+    ASSERT_FALSE(pbuf[i].buffer.empty());
     total_size += pbuf[i].buffer.size();
   }
 
   ASSERT_LE(total_size, sizeof(buf));
 
   proactor_->Await([&] {
-    for (unsigned i = 0; i < *res; ++i) {
+    for (unsigned i = 0; i < res; ++i) {
       conn_socket_->ReturnProvided(pbuf[i]);
     }
   });
 }
 
 #ifdef __linux__
+TEST_P(FiberSocketTest, RecvMultiShot) {
+  constexpr unsigned kBufLen = 40;
+  bool use_uring = GetParam() == "uring";
+  if (!use_uring) {
+    GTEST_SKIP() << "RecvMultiShot is supported only on uring";
+    return;
+  }
+
+  UringProactor* up = static_cast<UringProactor*>(proactor_.get());
+  up->Await([up] {
+      UringSocket::InitProvidedBuffers(4, kBufLen, up); });
+
+  unique_ptr<FiberSocketBase> sock;
+  error_code ec;
+  proactor_->Await([&] {
+    sock.reset(proactor_->CreateSocket());
+    ec = sock->Connect(listen_ep_);
+    ASSERT_FALSE(ec);
+  });
+
+  proactor_->Await([&] {
+    while (!conn_socket_) {
+      ThisFiber::SleepFor(1ms);
+    }
+    static_cast<UringSocket*>(conn_socket_.get())->EnableRecvMultishot();
+  });
+
+  unsigned res;
+  FiberSocketBase::ProvidedBuffer pbuf[8];
+
+  auto recv_fb = proactor_->LaunchFiber([&] {
+    res = conn_socket_->RecvProvided(8, pbuf);
+  });
+
+  uint8_t buf[128];
+  memset(buf, 'x', sizeof(buf));
+
+  proactor_->Await([&] {
+    auto wrt_ec = sock->Write(io::Bytes(buf));
+    ASSERT_FALSE(wrt_ec);
+  });
+
+  recv_fb.Join();
+
+  ASSERT_TRUE(res > 0 && res < 8);
+  size_t total_size = 0;
+  for (unsigned i = 0; i < res; ++i) {
+    ASSERT_FALSE(pbuf[i].buffer.empty());
+    total_size += pbuf[i].buffer.size();
+    for (uint8_t b : pbuf[i].buffer) {
+      ASSERT_EQ('x', b);
+    }
+  }
+
+  ASSERT_LE(total_size, sizeof(buf));
+
+  proactor_->Await([&] { std::ignore = sock->Close(); });
+  proactor_->Await([&] {
+    for (unsigned i = 0; i < res; ++i) {
+      conn_socket_->ReturnProvided(pbuf[i]);
+    }
+    res = conn_socket_->RecvProvided(8, pbuf);
+  });
+  ASSERT_EQ(res, 1);
+}
+
+
 TEST_P(FiberSocketTest, NotEmpty) {
   bool use_uring = GetParam() == "uring";
   bool has_poll_first = false;

util/fibers/uring_proactor.cc

@@ -516,7 +516,7 @@ void UringProactor::EnqueueMultishotCompletion(uint16_t group_id, IoResult res,
                                                uint16_t* tail) {
   DCHECK_LT(group_id, bufring_groups_.size());
   DCHECK(flags & IORING_CQE_F_BUFFER) << res;
-  DCHECK_GT(res, 0) << flags;
+  CHECK_GT(res, 0) << flags;
 
   auto& buf_group = bufring_groups_[group_id];
   if (!buf_group.multishot_arr) {
@@ -574,16 +574,9 @@ auto UringProactor::PullMultiShotCompletion(uint16_t group_id, uint16_t* tail) -
   // link the entry to the free list.
   head_entry.next = buf_group.free_multi_shot_id;
   buf_group.free_multi_shot_id = head_id;
-
-  MultiShotResult res;
-  if (head_entry.res >= 0) {
-    uint8_t* buf = buf_group.buf + head_entry.bid * buf_group.entry_size;
-    res.emplace(io::MutableBytes{buf, size_t(head_entry.res)});
-  } else {
-    res = nonstd::make_unexpected(-head_entry.res);
-  }
-
-  return res;
+  CHECK_GT(head_entry.res, 0);
+  uint8_t* buf = buf_group.buf + head_entry.bid * buf_group.entry_size;
+  return MultiShotResult{buf, size_t(head_entry.res)};
 }
 
 void UringProactor::RegrowCentries() {

util/fibers/uring_proactor.h

@@ -7,9 +7,9 @@
 #include <liburing.h>
 #include <pthread.h>
 
+#include "base/segment_pool.h"
 #include "util/fibers/proactor_base.h"
 #include "util/fibers/submit_entry.h"
-#include "base/segment_pool.h"
 
 namespace util {
 namespace fb2 {
@@ -131,7 +131,9 @@ class UringProactor : public ProactorBase {
   // Returns bufring entry size for the given group_id.
   // -1 if group_id is invalid.
   int BufRingEntrySize(unsigned group_id) const {
-    return group_id < bufring_groups_.size() ? bufring_groups_[group_id].entry_size : -1;
+    return group_id < bufring_groups_.size() && bufring_groups_[group_id].ring != nullptr
+               ? bufring_groups_[group_id].entry_size
+               : -1;
   }
 
   // Returns number of available entries at the time of the call.
@@ -156,8 +158,7 @@ class UringProactor : public ProactorBase {
   // Returns a new head.
   void EnqueueMultishotCompletion(uint16_t group_id, IoResult res, uint32_t flags, uint16_t* tail);
 
-  // in case of error, returns errno.
-  using MultiShotResult = nonstd::expected<io::Bytes, unsigned>;
+  using MultiShotResult = io::Bytes;
 
   // Pulls a single range of a multishot completion. head must point to a valid id.
   // Once the queue of completions is exhausted, head is set to kMultiShotUndef.
@@ -219,8 +220,8 @@ class UringProactor : public ProactorBase {
     io_uring_buf_ring* ring = nullptr;
     uint8_t* buf = nullptr;
     MultiShot* multishot_arr = nullptr;  // Array of a cardinality of nentries.
-    uint8_t nentries_exp = 0;  // 2^nentries_exp is the number of entries.
-    uint8_t multishot_exp = 0;  // 2^multishot_exp is the number of multishot entries.
+    uint8_t nentries_exp = 0;            // 2^nentries_exp is the number of entries.
+    uint8_t multishot_exp = 0;           // 2^multishot_exp is the number of multishot entries.
     uint16_t free_multi_shot_id = 0;
     uint32_t entry_size = 0;
   };
@@ -278,7 +279,7 @@ class FiberCall {
   UringProactor::IoResult io_res_ = 0;
   uint32_t res_flags_ = 0;  // set by waker upon completion.
   bool was_run_ = false;
-  timespec ts_;             // in case of timeout.
+  timespec ts_;  // in case of timeout.
 };
 
 }  // namespace fb2