Add io_uring support to fastboot

Test: th
Bug: 31712568

Change-Id: If18bd4eb53ffe851c31d7dc61906921da111114f
diff --git a/fastboot/device/usb.cpp b/fastboot/device/usb.cpp
index 4115a6d..b77d772 100644
--- a/fastboot/device/usb.cpp
+++ b/fastboot/device/usb.cpp
@@ -15,6 +15,7 @@
  */
 
 #include "usb.h"
+#include "usb_iouring.h"
 
 #include <dirent.h>
 #include <errno.h>
@@ -28,6 +29,7 @@
 
 #include <linux/usb/ch9.h>
 #include <linux/usb/functionfs.h>
+#include <sys/utsname.h>
 
 #include <algorithm>
 #include <atomic>
@@ -38,6 +40,7 @@
 
 #include <android-base/logging.h>
 #include <android-base/properties.h>
+#include <liburing.h>
 
 using namespace std::chrono_literals;
 
@@ -65,8 +68,8 @@
     }
 }
 
-static int getMaxPacketSize(int ffs_fd) {
-    usb_endpoint_descriptor desc;
+int getMaxPacketSize(int ffs_fd) {
+    usb_endpoint_descriptor desc{};
     if (ioctl(ffs_fd, FUNCTIONFS_ENDPOINT_DESC, reinterpret_cast<unsigned long>(&desc))) {
         D("[ could not get endpoint descriptor! (%d) ]", errno);
         return MAX_PACKET_SIZE_HS;
@@ -128,11 +131,9 @@
 
 static int usb_ffs_do_aio(usb_handle* h, const void* data, int len, bool read) {
     aio_block* aiob = read ? &h->read_aiob : &h->write_aiob;
-    bool zero_packet = false;
 
     int num_bufs = len / h->io_size + (len % h->io_size == 0 ? 0 : 1);
     const char* cur_data = reinterpret_cast<const char*>(data);
-    int packet_size = getMaxPacketSize(aiob->fd);
 
     if (posix_madvise(const_cast<void*>(data), len, POSIX_MADV_SEQUENTIAL | POSIX_MADV_WILLNEED) <
         0) {
@@ -145,17 +146,6 @@
 
         len -= buf_len;
         cur_data += buf_len;
-
-        if (len == 0 && buf_len % packet_size == 0 && read) {
-            // adb does not expect the device to send a zero packet after data transfer,
-            // but the host *does* send a zero packet for the device to read.
-            zero_packet = h->reads_zero_packets;
-        }
-    }
-    if (zero_packet) {
-        io_prep(&aiob->iocb[num_bufs], aiob->fd, reinterpret_cast<const void*>(cur_data),
-                packet_size, 0, read);
-        num_bufs += 1;
     }
 
     while (true) {
@@ -204,21 +194,46 @@
     h->open_new_connection = true;
     h->lock.unlock();
     h->notify.notify_one();
+    if (h->aio_type == AIOType::IO_URING) {
+        exit_io_uring_ffs(h);
+    }
 }
 
-usb_handle* create_usb_handle(unsigned num_bufs, unsigned io_size) {
-    usb_handle* h = new usb_handle();
+bool DoesKernelSupportIouring() {
+    struct utsname uts {};
+    unsigned int major = 0, minor = 0;
+    if ((uname(&uts) != 0) || (sscanf(uts.release, "%u.%u", &major, &minor) != 2)) {
+        return false;
+    }
+    if (major > 5) {
+        return true;
+    }
+    // We will only support kernels from 5.6 onwards as IOSQE_ASYNC flag and
+    // IO_URING_OP_READ/WRITE opcodes were introduced only on 5.6 kernel
+    return minor >= 6;
+}
 
-    if (android::base::GetBoolProperty("sys.usb.ffs.aio_compat", false)) {
+std::unique_ptr<usb_handle> create_usb_handle(unsigned num_bufs, unsigned io_size) {
+    auto h = std::make_unique<usb_handle>();
+    if (DoesKernelSupportIouring() &&
+        android::base::GetBoolProperty("sys.usb.ffs.io_uring_enabled", false)) {
+        init_io_uring_ffs(h.get(), num_bufs);
+        h->aio_type = AIOType::IO_URING;
+        LOG(INFO) << "Using io_uring for usb ffs";
+    } else if (android::base::GetBoolProperty("sys.usb.ffs.aio_compat", false)) {
         // Devices on older kernels (< 3.18) will not have aio support for ffs
         // unless backported. Fall back on the non-aio functions instead.
         h->write = usb_ffs_write;
         h->read = usb_ffs_read;
+        h->aio_type = AIOType::SYNC_IO;
+        LOG(INFO) << "Using sync io for usb ffs";
     } else {
         h->write = usb_ffs_aio_write;
         h->read = usb_ffs_aio_read;
         aio_block_init(&h->read_aiob, num_bufs);
         aio_block_init(&h->write_aiob, num_bufs);
+        h->aio_type = AIOType::AIO;
+        LOG(INFO) << "Using aio for usb ffs";
     }
     h->io_size = io_size;
     h->close = usb_ffs_close;
diff --git a/fastboot/device/usb.h b/fastboot/device/usb.h
index 6c3f542..8996c31 100644
--- a/fastboot/device/usb.h
+++ b/fastboot/device/usb.h
@@ -18,8 +18,10 @@
 
 #include <linux/usb/functionfs.h>
 
+#include <liburing.h>
 #include <atomic>
 #include <condition_variable>
+#include <memory>
 #include <mutex>
 #include <vector>
 
@@ -35,9 +37,11 @@
     int fd;
 };
 
-struct usb_handle {
-    usb_handle() {}
+int getMaxPacketSize(int ffs_fd);
 
+enum class AIOType { SYNC_IO, AIO, IO_URING };
+
+struct usb_handle {
     std::condition_variable notify;
     std::mutex lock;
     bool open_new_connection = true;
@@ -56,8 +60,9 @@
     struct aio_block read_aiob;
     struct aio_block write_aiob;
 
-    bool reads_zero_packets;
+    io_uring ring;
     size_t io_size;
+    AIOType aio_type;
 };
 
-usb_handle* create_usb_handle(unsigned num_bufs, unsigned io_size);
+std::unique_ptr<usb_handle> create_usb_handle(unsigned num_bufs, unsigned io_size);
diff --git a/fastboot/device/usb_client.cpp b/fastboot/device/usb_client.cpp
index 3f9b0f0..d1b38d4 100644
--- a/fastboot/device/usb_client.cpp
+++ b/fastboot/device/usb_client.cpp
@@ -232,7 +232,6 @@
 
     h->read_aiob.fd = h->bulk_out.get();
     h->write_aiob.fd = h->bulk_in.get();
-    h->reads_zero_packets = false;
     return true;
 
 err:
diff --git a/fastboot/device/usb_iouring.cpp b/fastboot/device/usb_iouring.cpp
new file mode 100644
index 0000000..d987712
--- /dev/null
+++ b/fastboot/device/usb_iouring.cpp
@@ -0,0 +1,140 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <android-base/logging.h>
+#include <liburing.h>
+#include "liburing/io_uring.h"
+#include "usb.h"
+
+static int prep_async_read(struct io_uring* ring, int fd, void* data, size_t len, int64_t offset) {
+    if (io_uring_sq_space_left(ring) <= 0) {
+        LOG(ERROR) << "Submission queue run out of space.";
+        return -1;
+    }
+    auto sqe = io_uring_get_sqe(ring);
+    if (sqe == nullptr) {
+        return -1;
+    }
+    io_uring_sqe_set_flags(sqe, IOSQE_IO_LINK | IOSQE_ASYNC);
+    io_uring_prep_read(sqe, fd, data, len, offset);
+    return 0;
+}
+
+static int prep_async_write(struct io_uring* ring, int fd, const void* data, size_t len,
+                            int64_t offset) {
+    if (io_uring_sq_space_left(ring) <= 0) {
+        LOG(ERROR) << "Submission queue run out of space.";
+        return -1;
+    }
+    auto sqe = io_uring_get_sqe(ring);
+    if (sqe == nullptr) {
+        return -1;
+    }
+    io_uring_sqe_set_flags(sqe, IOSQE_IO_LINK | IOSQE_ASYNC);
+    io_uring_prep_write(sqe, fd, data, len, offset);
+    return 0;
+}
+
+template <bool read, typename T>
+int prep_async_io(struct io_uring* ring, int fd, T* data, size_t len, int64_t offset) {
+    if constexpr (read) {
+        return prep_async_read(ring, fd, data, len, offset);
+    } else {
+        return prep_async_write(ring, fd, data, len, offset);
+    }
+}
+
+template <typename T>
+static constexpr T DivRoundup(T x, T y) {
+    return (x + y - 1) / y;
+}
+
+extern int getMaxPacketSize(int ffs_fd);
+
+template <bool read, typename T>
+static int usb_ffs_do_aio(usb_handle* h, T* const data, const int len) {
+    const aio_block* aiob = read ? &h->read_aiob : &h->write_aiob;
+    const int num_requests = DivRoundup<int>(len, h->io_size);
+    auto cur_data = data;
+    const auto packet_size = getMaxPacketSize(aiob->fd);
+
+    for (int bytes_remain = len; bytes_remain > 0;) {
+        const int buf_len = std::min(bytes_remain, static_cast<int>(h->io_size));
+        const auto ret = prep_async_io<read>(&h->ring, aiob->fd, cur_data, buf_len, 0);
+        if (ret < 0) {
+            PLOG(ERROR) << "Failed to queue io_uring request";
+            return -1;
+        }
+
+        bytes_remain -= buf_len;
+        cur_data = reinterpret_cast<T*>(reinterpret_cast<size_t>(cur_data) + buf_len);
+    }
+    const int ret = io_uring_submit(&h->ring);
+    if (ret <= 0 || ret != num_requests) {
+        PLOG(ERROR) << "io_uring: failed to submit SQE entries to kernel";
+        return -1;
+    }
+    int res = 0;
+    bool success = true;
+    for (int i = 0; i < num_requests; ++i) {
+        struct io_uring_cqe* cqe{};
+        const auto ret = TEMP_FAILURE_RETRY(io_uring_wait_cqe(&h->ring, &cqe));
+        if (ret < 0 || cqe == nullptr) {
+            PLOG(ERROR) << "Failed to get CQE from kernel";
+            success = false;
+            continue;
+        }
+        res += cqe->res;
+        if (cqe->res < 0) {
+            LOG(ERROR) << "io_uring request failed:, i = " << i
+                       << ", num_requests = " << num_requests << ", res = " << cqe->res << ": "
+                       << strerror(cqe->res) << (read ? " read" : " write")
+                       << " request size: " << len << ", io_size: " << h->io_size
+                       << " max packet size: " << packet_size << ", fd: " << aiob->fd;
+            success = false;
+            errno = -cqe->res;
+        }
+        io_uring_cqe_seen(&h->ring, cqe);
+    }
+    if (!success) {
+        return -1;
+    }
+    return res;
+}
+
+static int usb_ffs_io_uring_read(usb_handle* h, void* data, int len, bool /* allow_partial */) {
+    return usb_ffs_do_aio<true>(h, data, len);
+}
+
+static int usb_ffs_io_uring_write(usb_handle* h, const void* data, int len) {
+    return usb_ffs_do_aio<false>(h, data, len);
+}
+
+void exit_io_uring_ffs(usb_handle* h) {
+    io_uring_queue_exit(&h->ring);
+}
+
+bool init_io_uring_ffs(usb_handle* h, size_t queue_depth) {
+    const auto err = io_uring_queue_init(queue_depth, &h->ring, 0);
+    if (err) {
+        LOG(ERROR) << "Failed to initialize io_uring of depth " << queue_depth << ": "
+                   << strerror(err);
+        return false;
+    }
+    h->write = usb_ffs_io_uring_write;
+    h->read = usb_ffs_io_uring_read;
+    return true;
+}
diff --git a/fastboot/device/usb_iouring.h b/fastboot/device/usb_iouring.h
new file mode 100644
index 0000000..7c14b81
--- /dev/null
+++ b/fastboot/device/usb_iouring.h
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "usb.h"
+
+bool init_io_uring_ffs(usb_handle* h, size_t queue_depth);
+
+void exit_io_uring_ffs(usb_handle* h);