storage/io_backend/iouring/io_uring_backend.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "../io_backend.h" | ||
| 2 | #include <cerrno> | ||
| 3 | |||
| 4 | thread_local int pending = 0; | ||
| 5 | thread_local std::vector<std::unique_ptr<coroutine<void>::pull_type>> coros; | ||
| 6 | #include <cstdlib> | ||
| 7 | #include <fcntl.h> | ||
| 8 | #include <liburing.h> | ||
| 9 | #include <sys/stat.h> | ||
| 10 | #include <unistd.h> | ||
| 11 | #include "base/factory.h" | ||
| 12 | #include "storage/kv_engine/base_kv.h" | ||
| 13 | |||
| 14 | class IoUringBackend : public IOBackend { | ||
| 15 | private: | ||
| 16 | std::string file_path; | ||
| 17 | int fd; | ||
| 18 | |||
| 19 | 383832 | static char* AllocateAligned(size_t bytes) { | |
| 20 | 383832 | void* ptr = nullptr; | |
| 21 | 383832 | int rc = posix_memalign(&ptr, PAGE_SIZE, bytes); | |
| 22 | 383832 | if (rc != 0 || ptr == nullptr) { | |
| 23 | ✗ | throw std::runtime_error( | |
| 24 | ✗ | "Failed to allocate aligned buffer, rc=" + std::to_string(rc)); | |
| 25 | } | ||
| 26 | 383832 | return reinterpret_cast<char*>(ptr); | |
| 27 | } | ||
| 28 | |||
| 29 | struct ThreadQpair { | ||
| 30 | struct io_uring qpair; | ||
| 31 | bool initialized = false; | ||
| 32 | 322 | ThreadQpair() = default; | |
| 33 | 322 | ~ThreadQpair() { release(); } | |
| 34 | |||
| 35 | 209281670 | static ThreadQpair& instance() { | |
| 36 | 209281670 | static thread_local ThreadQpair thread_qpair; | |
| 37 | 209281670 | return thread_qpair; | |
| 38 | } | ||
| 39 | |||
| 40 | 209281670 | struct io_uring* get(int queue_size) { | |
| 41 | 209281670 | if (!initialized) { | |
| 42 | 322 | int ret = io_uring_queue_init(queue_size, &qpair, 0); | |
| 43 | 322 | if (ret < 0) | |
| 44 | ✗ | throw std::runtime_error( | |
| 45 | ✗ | "Failed to initialize thread-local io_uring: " + | |
| 46 | ✗ | std::string(strerror(-ret))); | |
| 47 | 322 | initialized = true; | |
| 48 | } | ||
| 49 | 209281670 | return &qpair; | |
| 50 | } | ||
| 51 | |||
| 52 | 322 | void release() { | |
| 53 | 322 | if (initialized) { | |
| 54 | 322 | io_uring_queue_exit(&qpair); | |
| 55 | 322 | initialized = false; | |
| 56 | } | ||
| 57 | 322 | } | |
| 58 | }; | ||
| 59 | 209281670 | struct io_uring* get_thread_ring() { | |
| 60 | 209281670 | ThreadQpair& thread_qpair = ThreadQpair::instance(); | |
| 61 | 209281670 | return thread_qpair.get(queue_cnt); | |
| 62 | } | ||
| 63 | |||
| 64 | ✗ | void ReadPageAsync(coroutine<void>::push_type& sink, | |
| 65 | uint64_t index, | ||
| 66 | PageID_t page_id, | ||
| 67 | char* buffer) override { | ||
| 68 | ✗ | struct io_uring* ring = get_thread_ring(); | |
| 69 | ✗ | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 70 | ✗ | CHECK_NE(sqe, nullptr) << "Failed to get SQE for read operation"; | |
| 71 | ✗ | pending++; | |
| 72 | ✗ | io_uring_prep_read(sqe, fd, buffer, PAGE_SIZE, page_id * PAGE_SIZE); | |
| 73 | ✗ | sqe->user_data = index; | |
| 74 | ✗ | submit(); | |
| 75 | ✗ | sink(); | |
| 76 | ✗ | } | |
| 77 | 210796 | void ReadPageSync(PageID_t page_id, char* buffer) override { | |
| 78 | 210796 | struct io_uring* ring = get_thread_ring(); | |
| 79 | 210796 | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 80 | 210796 | CHECK_NE(sqe, nullptr) << "Failed to get SQE for sync read operation"; | |
| 81 | 210796 | io_uring_prep_read(sqe, fd, buffer, PAGE_SIZE, page_id * PAGE_SIZE); | |
| 82 | 210796 | int ret = io_uring_submit_and_wait(ring, 1); | |
| 83 | 210796 | CHECK_GE(ret, 0) << "Failed to submit sync read operation: " + | |
| 84 | ✗ | std::string(strerror(-ret)); | |
| 85 | 210796 | struct io_uring_cqe* cqe = nullptr; | |
| 86 | 210796 | ret = io_uring_wait_cqe(ring, &cqe); | |
| 87 | 210796 | CHECK_GE(ret, 0) << "Failed to wait CQE for sync read: " + | |
| 88 | ✗ | std::string(strerror(-ret)); | |
| 89 | 210796 | CHECK_NE(cqe, nullptr); | |
| 90 | 210796 | CHECK_GE(cqe->res, 0) << "Sync read failed: " + | |
| 91 | ✗ | std::string(strerror(-cqe->res)); | |
| 92 | 210796 | io_uring_cqe_seen(ring, cqe); | |
| 93 | 210796 | } | |
| 94 | |||
| 95 | ✗ | void WritePageAsync(coroutine<void>::push_type& sink, | |
| 96 | uint64_t index, | ||
| 97 | PageID_t page_id, | ||
| 98 | char* buffer) override { | ||
| 99 | ✗ | struct io_uring* ring = get_thread_ring(); | |
| 100 | ✗ | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 101 | ✗ | CHECK_NE(sqe, nullptr) << "Failed to get SQE for write operation"; | |
| 102 | ✗ | pending++; | |
| 103 | ✗ | io_uring_prep_write(sqe, fd, buffer, PAGE_SIZE, page_id * PAGE_SIZE); | |
| 104 | ✗ | sqe->user_data = index; | |
| 105 | ✗ | submit(); | |
| 106 | ✗ | sink(); | |
| 107 | ✗ | } | |
| 108 | ✗ | void WritePageSync(PageID_t page_id, char* buffer) override { | |
| 109 | ✗ | struct io_uring* ring = get_thread_ring(); | |
| 110 | ✗ | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 111 | ✗ | CHECK_NE(sqe, nullptr) << "Failed to get SQE for sync write operation"; | |
| 112 | ✗ | io_uring_prep_write(sqe, fd, buffer, PAGE_SIZE, page_id * PAGE_SIZE); | |
| 113 | ✗ | int ret = io_uring_submit_and_wait(ring, 1); | |
| 114 | ✗ | CHECK_GE(ret, 0) << "Failed to submit sync write operation: " + | |
| 115 | ✗ | std::string(strerror(-ret)); | |
| 116 | ✗ | struct io_uring_cqe* cqe = nullptr; | |
| 117 | ✗ | ret = io_uring_wait_cqe(ring, &cqe); | |
| 118 | ✗ | CHECK_GE(ret, 0) << "Failed to wait CQE for sync write: " + | |
| 119 | ✗ | std::string(strerror(-ret)); | |
| 120 | ✗ | CHECK_NE(cqe, nullptr); | |
| 121 | ✗ | CHECK_GE(cqe->res, 0) << "Sync write failed: " + | |
| 122 | ✗ | std::string(strerror(-cqe->res)); | |
| 123 | ✗ | io_uring_cqe_seen(ring, cqe); | |
| 124 | ✗ | } | |
| 125 | |||
| 126 | public: | ||
| 127 | 670 | IoUringBackend(const BaseKVConfig& config) : IOBackend(config), fd(-1) { | |
| 128 | 670 | file_path = config.json_config_.at("file_path").get<std::string>(); | |
| 129 | 670 | } | |
| 130 | 1340 | ~IoUringBackend() { | |
| 131 | 670 | if (fd >= 0) { | |
| 132 | 670 | close(fd); | |
| 133 | 670 | fd = -1; | |
| 134 | } | ||
| 135 | 670 | if (empty_page) { | |
| 136 | 670 | std::free(empty_page); | |
| 137 | 670 | empty_page = nullptr; | |
| 138 | } | ||
| 139 | 1340 | } | |
| 140 | |||
| 141 | 670 | void init() override { | |
| 142 | 670 | bool exists = (access(file_path.c_str(), F_OK) != -1); | |
| 143 | 670 | fd = | |
| 144 | 670 | open(file_path.c_str(), O_RDWR | O_CREAT | O_DIRECT, S_IRUSR | S_IWUSR); | |
| 145 | 670 | if (fd < 0) { | |
| 146 | ✗ | const int err = errno; | |
| 147 | ✗ | throw std::runtime_error( | |
| 148 | ✗ | "Failed to open file: " + file_path + ", errno=" + | |
| 149 | ✗ | std::to_string(err) + ", detail=" + std::string(strerror(err))); | |
| 150 | } | ||
| 151 | 670 | if (exists) { | |
| 152 | struct stat file_stat; | ||
| 153 | ✗ | fstat(fd, &file_stat); | |
| 154 | ✗ | next_page_id = file_stat.st_size / PAGE_SIZE; | |
| 155 | } else { | ||
| 156 | 670 | next_page_id = 1; | |
| 157 | } | ||
| 158 | 670 | empty_page = AllocateAligned(PAGE_SIZE); | |
| 159 | 670 | std::memset(empty_page, 0, PAGE_SIZE); | |
| 160 | 670 | } | |
| 161 | |||
| 162 | 210796 | char* AllocateBuffer() override { return AllocateAligned(PAGE_SIZE); } | |
| 163 | 172366 | char* AllocateBuffer(uint64_t page_count) override { | |
| 164 | 172366 | return AllocateAligned(PAGE_SIZE * page_count); | |
| 165 | } | ||
| 166 | 383162 | void FreeBuffer(char* buf) override { std::free(buf); } | |
| 167 | |||
| 168 | ✗ | void* GetPage(coroutine<void>::push_type& sink, | |
| 169 | uint64_t index, | ||
| 170 | PageID_t page_id) override { | ||
| 171 | ✗ | char* buffer = AllocateBuffer(); | |
| 172 | ✗ | ReadPageAsync(sink, index, page_id, buffer); | |
| 173 | ✗ | return buffer; | |
| 174 | } | ||
| 175 | ✗ | void* GetPage(PageID_t page_id) override { | |
| 176 | ✗ | char* buffer = AllocateBuffer(); | |
| 177 | ✗ | ReadPageSync(page_id, buffer); | |
| 178 | ✗ | return buffer; | |
| 179 | } | ||
| 180 | |||
| 181 | ✗ | void Unpin(coroutine<void>::push_type& sink, | |
| 182 | uint64_t index, | ||
| 183 | PageID_t page_id, | ||
| 184 | void* page_data, | ||
| 185 | bool is_dirty) override { | ||
| 186 | ✗ | if (is_dirty) | |
| 187 | ✗ | WritePage(sink, index, page_id, reinterpret_cast<char*>(page_data)); | |
| 188 | ✗ | FreeBuffer(reinterpret_cast<char*>(page_data)); | |
| 189 | ✗ | } | |
| 190 | |||
| 191 | ✗ | void Unpin(PageID_t page_id, void* page_data, bool is_dirty) override { | |
| 192 | ✗ | if (is_dirty) | |
| 193 | ✗ | WritePageSync(page_id, reinterpret_cast<char*>(page_data)); | |
| 194 | ✗ | FreeBuffer(reinterpret_cast<char*>(page_data)); | |
| 195 | ✗ | } | |
| 196 | |||
| 197 | ✗ | void PollCompletion() override { | |
| 198 | ✗ | struct io_uring* ring = get_thread_ring(); | |
| 199 | struct io_uring_cqe* cqe; | ||
| 200 | ✗ | int ret = io_uring_peek_cqe(ring, &cqe); | |
| 201 | ✗ | if (ret == -EAGAIN) | |
| 202 | ✗ | return; // No completion events | |
| 203 | ✗ | else if (ret < 0) | |
| 204 | ✗ | throw std::runtime_error( | |
| 205 | ✗ | "Failed to peek CQE: " + std::string(strerror(-ret))); | |
| 206 | ✗ | int id = cqe->user_data; | |
| 207 | ✗ | pending--; | |
| 208 | ✗ | io_uring_cqe_seen(ring, cqe); | |
| 209 | ✗ | if (coros[id] && *coros[id]) | |
| 210 | ✗ | (*coros[id])(); | |
| 211 | ✗ | return; | |
| 212 | } | ||
| 213 | |||
| 214 | 208799742 | void submit() override { | |
| 215 | 208799742 | struct io_uring* ring = get_thread_ring(); | |
| 216 | 208799742 | int ret = io_uring_submit(ring); | |
| 217 | 208799742 | CHECK_GE(ret, 0) << "Failed to submit io_uring operation: " + | |
| 218 | ✗ | std::string(strerror(-ret)); | |
| 219 | 208799742 | } | |
| 220 | |||
| 221 | 98766 | void BatchWritePages(const std::vector<IOEntry>& entries) override { | |
| 222 | 98766 | if (entries.empty()) | |
| 223 | ✗ | return; | |
| 224 | 98766 | struct io_uring* ring = get_thread_ring(); | |
| 225 | 98766 | int max_inflight = std::max(queue_cnt / 2, 1); | |
| 226 | 98766 | size_t submitted = 0; | |
| 227 | 98766 | int inflight = 0; | |
| 228 | |||
| 229 | 70963090 | while (submitted < entries.size() || inflight > 0) { | |
| 230 | 70963090 | while (submitted < entries.size() && inflight < max_inflight) { | |
| 231 | 98766 | auto& e = entries[submitted]; | |
| 232 | 98766 | uint64_t total_bytes = e.page_count * PAGE_SIZE; | |
| 233 | 98766 | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 234 | 98766 | CHECK_NE(sqe, nullptr) << "Failed to get SQE for batch write"; | |
| 235 | 98766 | io_uring_prep_write( | |
| 236 | 98766 | sqe, fd, e.buffer, total_bytes, e.page_id * PAGE_SIZE); | |
| 237 | 98766 | sqe->user_data = 0; | |
| 238 | 98766 | inflight++; | |
| 239 | 98766 | submitted++; | |
| 240 | } | ||
| 241 | 70864324 | submit(); | |
| 242 | struct io_uring_cqe* cqe; | ||
| 243 | 70963090 | while (inflight > 0 && io_uring_peek_cqe(ring, &cqe) == 0) { | |
| 244 | 98766 | CHECK_GE(cqe->res, 0) | |
| 245 | ✗ | << "Batch write failed: " + std::string(strerror(-cqe->res)); | |
| 246 | 98766 | io_uring_cqe_seen(ring, cqe); | |
| 247 | 98766 | inflight--; | |
| 248 | } | ||
| 249 | } | ||
| 250 | } | ||
| 251 | |||
| 252 | 172366 | void BatchReadPages(const std::vector<IOEntry>& entries) override { | |
| 253 | 172366 | if (entries.empty()) | |
| 254 | ✗ | return; | |
| 255 | 172366 | RecordRead(entries); | |
| 256 | 172366 | struct io_uring* ring = get_thread_ring(); | |
| 257 | 172366 | int max_inflight = std::max(queue_cnt / 2, 1); | |
| 258 | 172366 | size_t submitted = 0; | |
| 259 | 172366 | int inflight = 0; | |
| 260 | |||
| 261 | 138107784 | while (submitted < entries.size() || inflight > 0) { | |
| 262 | 138107784 | while (submitted < entries.size() && inflight < max_inflight) { | |
| 263 | 172366 | auto& e = entries[submitted]; | |
| 264 | 172366 | uint64_t total_bytes = e.page_count * PAGE_SIZE; | |
| 265 | 172366 | struct io_uring_sqe* sqe = io_uring_get_sqe(ring); | |
| 266 | 172366 | CHECK_NE(sqe, nullptr) << "Failed to get SQE for batch read"; | |
| 267 | 172366 | io_uring_prep_read( | |
| 268 | 172366 | sqe, fd, e.buffer, total_bytes, e.page_id * PAGE_SIZE); | |
| 269 | 172366 | sqe->user_data = 0; | |
| 270 | 172366 | inflight++; | |
| 271 | 172366 | submitted++; | |
| 272 | } | ||
| 273 | 137935418 | submit(); | |
| 274 | struct io_uring_cqe* cqe; | ||
| 275 | 138107784 | while (inflight > 0 && io_uring_peek_cqe(ring, &cqe) == 0) { | |
| 276 | 172366 | CHECK_GE(cqe->res, 0) | |
| 277 | ✗ | << "Batch read failed: " + std::string(strerror(-cqe->res)); | |
| 278 | 172366 | io_uring_cqe_seen(ring, cqe); | |
| 279 | 172366 | inflight--; | |
| 280 | } | ||
| 281 | } | ||
| 282 | } | ||
| 283 | }; | ||
| 284 | |||
| 285 | 12 | extern "C" void RecStoreForceLinkIoUringBackend() {} | |
| 286 | |||
| 287 | FACTORY_REGISTER(IOBackend, IOURING, IoUringBackend, const BaseKVConfig&); | ||
| 288 |