Add support for polling the io_uring submission/completion queues

Summary: Add support for polling the io_uring submission/completion queues

Reviewed By: kevin-vigor

Differential Revision: D23566444

fbshipit-source-id: 7f8331d1f34db78dc4bca4aea3e5f5a1c8f00023

folly/experimental/io/IoUringBackend.cpp

@@ -101,6 +101,13 @@ IoUringBackend::IoUringBackend(Options options)
   params_.flags |= IORING_SETUP_CQSIZE;
   params_.cq_entries = options.capacity;
 
+  // poll SQ options
+  if (options.flags & Options::Flags::POLL_SQ) {
+    params_.flags |= IORING_SETUP_SQPOLL;
+    params_.sq_thread_idle = options.sqIdle.count();
+    params_.sq_thread_cpu = options.sqCpu;
+  }
+
   // allocate entries both for poll add and cancel
   if (::io_uring_queue_init_params(
           2 * options_.maxSubmit, &ioRing_, &params_)) {
@@ -199,7 +206,7 @@ bool IoUringBackend::isAvailable() {
 }
 
 void* IoUringBackend::allocSubmissionEntry() {
-  return ::io_uring_get_sqe(&ioRing_);
+  return get_sqe();
 }
 
 int IoUringBackend::submitOne(IoCb* /*unused*/) {
@@ -212,7 +219,7 @@ int IoUringBackend::cancelOne(IoCb* ioCb) {
     return 0;
   }
 
-  auto* sqe = ::io_uring_get_sqe(&ioRing_);
+  auto* sqe = get_sqe();
   CHECK(sqe);
 
   rentry->prepCancel(sqe, ioCb); // prev entry
@@ -233,7 +240,21 @@ int IoUringBackend::getActiveEvents(WaitForEventsMode waitForEvents) {
   // we can be called from the submitList() method
   // or with non blocking flags
   if (FOLLY_LIKELY(waitForEvents == WaitForEventsMode::WAIT)) {
+    // if polling the CQ, busy wait for one entry
+    if (options_.flags & Options::Flags::POLL_CQ) {
+      do {
+        ::io_uring_peek_cqe(&ioRing_, &cqe);
+        asm_volatile_pause();
+        // call the loop callback if installed
+        // we call it every time we poll for a CQE
+        // regardless of the io_uring_peek_cqe result
+        if (cqPollLoopCallback_) {
+          cqPollLoopCallback_();
+        }
+      } while (!cqe);
+    } else {
       ::io_uring_wait_cqe(&ioRing_, &cqe);
+    }
   } else {
     ::io_uring_peek_cqe(&ioRing_, &cqe);
   }
@@ -254,7 +275,11 @@ int IoUringBackend::submitBusyCheck(int num, WaitForEventsMode waitForEvents) {
   int res;
   while (i < num) {
     if (waitForEvents == WaitForEventsMode::WAIT) {
+      if (options_.flags & Options::Flags::POLL_CQ) {
+        res = ::io_uring_submit(&ioRing_);
+      } else {
         res = ::io_uring_submit_and_wait(&ioRing_, 1);
+      }
     } else {
       res = ::io_uring_submit(&ioRing_);
     }
@@ -278,6 +303,15 @@ int IoUringBackend::submitBusyCheck(int num, WaitForEventsMode waitForEvents) {
     }
 
     i += res;
+
+    // if polling the CQ, busy wait for one entry
+    if (waitForEvents == WaitForEventsMode::WAIT &&
+        options_.flags & Options::Flags::POLL_CQ && i == num) {
+      struct io_uring_cqe* cqe = nullptr;
+      while (!cqe) {
+        ::io_uring_peek_cqe(&ioRing_, &cqe);
+      }
+    }
   }
 
   return num;
@@ -292,7 +326,7 @@ size_t IoUringBackend::submitList(
   while (!ioCbs.empty()) {
     auto* entry = &ioCbs.front();
     ioCbs.pop_front();
-    auto* sqe = ::io_uring_get_sqe(&ioRing_);
+    auto* sqe = get_sqe();
     CHECK(sqe); // this should not happen
 
     entry->processSubmit(sqe);

folly/experimental/io/IoUringBackend.h

@@ -21,6 +21,7 @@
 #include <folly/Function.h>
 #include <folly/Range.h>
 #include <folly/experimental/io/PollIoBackend.h>
+#include <folly/portability/Asm.h>
 #include <folly/small_vector.h>
 
 #include <glog/logging.h>
@@ -41,6 +42,13 @@ class IoUringBackend : public PollIoBackend {
   // supports the io_uring backend
   static bool isAvailable();
 
+  // CQ poll mode loop callback
+  using CQPollLoopCallback = folly::Function<void()>;
+
+  void setCQPollLoopCallback(CQPollLoopCallback&& cb) {
+    cqPollLoopCallback_ = std::move(cb);
+  }
+
   // from PollIoBackend
   FdRegistrationRecord* registerFd(int fd) override {
     return fdRegistry_.alloc(fd);
@@ -393,6 +401,19 @@ class IoUringBackend : public PollIoBackend {
 
   void cleanup();
 
+  FOLLY_ALWAYS_INLINE struct io_uring_sqe* get_sqe() {
+    struct io_uring_sqe* ret = ::io_uring_get_sqe(&ioRing_);
+    // if running with SQ poll enabled
+    // we might have to wait for an sq entry to available
+    // before we can submit another one
+    while ((options_.flags & Options::Flags::POLL_SQ) && !ret) {
+      asm_volatile_pause();
+      ret = ::io_uring_get_sqe(&ioRing_);
+    }
+
+    return ret;
+  }
+
   size_t submit_internal();
 
   // io_uring related
@@ -400,5 +421,9 @@ class IoUringBackend : public PollIoBackend {
   struct io_uring ioRing_;
 
   FdRegistry fdRegistry_;
+
+  // poll callback to be invoked if POLL_CQ flag is set
+  // every time we poll for a CQE
+  CQPollLoopCallback cqPollLoopCallback_;
 };
 } // namespace folly

folly/experimental/io/PollIoBackend.h

@@ -29,6 +29,7 @@
 
 #include <folly/CPortability.h>
 #include <folly/CppAttributes.h>
+#include <folly/Function.h>
 #include <folly/io/async/EventBaseBackendBase.h>
 
 namespace folly {
@@ -36,6 +37,11 @@ namespace folly {
 class PollIoBackend : public EventBaseBackendBase {
  public:
   struct Options {
+    enum Flags {
+      POLL_SQ = 0x1,
+      POLL_CQ = 0x2,
+    };
+
     Options() = default;
 
     Options& setCapacity(size_t v) {
@@ -62,10 +68,39 @@ class PollIoBackend : public EventBaseBackendBase {
       return *this;
     }
 
+    Options& setFlags(uint32_t v) {
+      flags = v;
+
+      return *this;
+    }
+
+    Options& setSQIdle(std::chrono::milliseconds v) {
+      sqIdle = v;
+
+      return *this;
+    }
+
+    Options& setCQIdle(std::chrono::milliseconds v) {
+      cqIdle = v;
+
+      return *this;
+    }
+
+    Options& setCQCpu(uint32_t v) {
+      sqCpu = v;
+
+      return *this;
+    }
+
     size_t capacity{0};
     size_t maxSubmit{128};
     size_t maxGet{static_cast<size_t>(-1)};
     bool useRegisteredFds{false};
+    uint32_t flags{0};
+
+    std::chrono::milliseconds sqIdle{0};
+    std::chrono::milliseconds cqIdle{0};
+    uint32_t sqCpu{0};
   };
 
   explicit PollIoBackend(Options options);

folly/experimental/io/test/IoUringBackendTest.cpp

@@ -934,6 +934,23 @@ struct IoUringRegFdBackendProvider {
   }
 };
 
+// CQ polling
+struct IoUringPollCQBackendProvider {
+  static std::unique_ptr<folly::EventBaseBackendBase> getBackend() {
+    try {
+      folly::PollIoBackend::Options options;
+      options.setCapacity(kCapacity)
+          .setMaxSubmit(kMaxSubmit)
+          .setMaxGet(kMaxGet)
+          .setUseRegisteredFds(false)
+          .setFlags(folly::PollIoBackend::Options::Flags::POLL_CQ);
+      return std::make_unique<folly::IoUringBackend>(options);
+    } catch (const IoUringBackend::NotAvailable&) {
+      return nullptr;
+    }
+  }
+};
+
 // Instantiate the non registered fd tests
 INSTANTIATE_TYPED_TEST_CASE_P(IoUring, EventBaseTest, IoUringBackendProvider);
 INSTANTIATE_TYPED_TEST_CASE_P(IoUring, EventBaseTest1, IoUringBackendProvider);
@@ -947,5 +964,15 @@ INSTANTIATE_TYPED_TEST_CASE_P(
     IoUringRegFd,
     EventBaseTest1,
     IoUringRegFdBackendProvider);
+
+// Instantiate the poll CQ tests
+INSTANTIATE_TYPED_TEST_CASE_P(
+    IoUringPollCQ,
+    EventBaseTest,
+    IoUringPollCQBackendProvider);
+INSTANTIATE_TYPED_TEST_CASE_P(
+    IoUringPollCQ,
+    EventBaseTest1,
+    IoUringPollCQBackendProvider);
 } // namespace test
 } // namespace folly