io_uring: Add file registration management

This makes it possible to use SQPOLL and other io_uring options.

Change-Id: I543879b193461ff022b3fd567ba35f2d95754bd7
Signed-off-by: William A. Kennington III <wak@google.com>
diff --git a/src/stdplus/io_uring.cpp b/src/stdplus/io_uring.cpp
index 3ec5da1..1402b8a 100644
--- a/src/stdplus/io_uring.cpp
+++ b/src/stdplus/io_uring.cpp
@@ -27,6 +27,37 @@
     }
 }
 
+IoUring::FileHandle::FileHandle(unsigned slot, IoUring& ring) :
+    slot(slot, &ring)
+{
+}
+
+void IoUring::FileHandle::drop(unsigned&& slot, IoUring*& ring)
+{
+    ring->updateFile(slot, -1);
+}
+
+[[nodiscard]] IoUring::FileHandle IoUring::registerFile(int fd)
+{
+    unsigned slot = 0;
+    for (; slot < files.size(); slot++)
+    {
+        if (files[slot] == -1)
+        {
+            updateFile(slot, fd);
+            return FileHandle(slot, *this);
+        }
+    }
+
+    io_uring_unregister_files(&ring);
+    files.reserve(files.size() + 1);
+    files.resize(files.capacity(), -1);
+    files[slot] = fd;
+    CHECK_RET(io_uring_register_files(&ring, files.data(), files.size()),
+              "io_uring_register_files");
+    return FileHandle(slot, *this);
+}
+
 io_uring_sqe& IoUring::getSQE()
 {
     return *CHECK_ERRNO(io_uring_get_sqe(&ring), "io_uring_get_sqe");
@@ -108,4 +139,11 @@
     h->handleCQE(cqe);
 }
 
+void IoUring::updateFile(unsigned slot, int fd)
+{
+    files[slot] = fd;
+    CHECK_RET(io_uring_register_files_update(&ring, slot, &files[slot], 1),
+              "io_uring_register_files_update");
+}
+
 } // namespace stdplus
diff --git a/src/stdplus/io_uring.hpp b/src/stdplus/io_uring.hpp
index d3090ed..7ea2114 100644
--- a/src/stdplus/io_uring.hpp
+++ b/src/stdplus/io_uring.hpp
@@ -2,6 +2,7 @@
 #include <liburing.h>
 
 #include <stdplus/fd/managed.hpp>
+#include <stdplus/handle/managed.hpp>
 
 #include <optional>
 #include <vector>
@@ -24,6 +25,24 @@
         virtual void handleCQE(io_uring_cqe&) noexcept = 0;
     };
 
+    class FileHandle
+    {
+      public:
+        inline operator unsigned() const
+        {
+            return *slot;
+        }
+
+      private:
+        explicit FileHandle(unsigned slot, IoUring& ring);
+
+        static void drop(unsigned&& slot, IoUring*& ring);
+
+        Managed<unsigned, IoUring*>::Handle<drop> slot;
+
+        friend class IoUring;
+    };
+
     explicit IoUring(size_t queue_size = 10);
     IoUring(IoUring&&) = delete;
     IoUring& operator=(IoUring&&) = delete;
@@ -31,6 +50,14 @@
     IoUring& operator=(const IoUring&) = delete;
     ~IoUring();
 
+    /** @brief Registers a file descriptor with a slot on the ring
+     *
+     *  @param[in] fd - The file descriptor to register
+     *  @throws std::system_error if the allocation fails
+     *  @return A handle to the registered file on the ring
+     */
+    [[nodiscard]] FileHandle registerFile(int fd);
+
     /** @brief Gets an unused SQE from the ring
      *
      *  @throws std::system_error if the allocation fails
@@ -77,8 +104,10 @@
     io_uring ring;
     std::optional<stdplus::ManagedFd> event_fd;
     std::vector<CQEHandler*> handlers;
+    std::vector<int> files;
 
     void dropHandler(CQEHandler* h, io_uring_cqe& cqe) noexcept;
+    void updateFile(unsigned slot, int fd);
 };
 
 } // namespace stdplus
diff --git a/test/io_uring.cpp b/test/io_uring.cpp
index 5be1b91..6e9981d 100644
--- a/test/io_uring.cpp
+++ b/test/io_uring.cpp
@@ -131,4 +131,17 @@
     EXPECT_CALL(h[0], handleCQE(_));
 }
 
+TEST_F(IoUringTest, RegisterFiles)
+{
+    // Slots are always allocated linearly and re-used if invalidated
+    std::optional<IoUring::FileHandle> h;
+    h = ring.registerFile(0);
+    EXPECT_EQ(*h, 0);
+    h = ring.registerFile(1);
+    EXPECT_EQ(*h, 1);
+    // The first handle should have dropped and can be replaced
+    h = ring.registerFile(2);
+    EXPECT_EQ(*h, 0);
+}
+
 } // namespace stdplus