]> git.sesse.net Git - plocate/blob - io_uring_engine.cpp
Support --debug for plocate-build, and unbreak some debug printfs there.
[plocate] / io_uring_engine.cpp
1 #include <assert.h>
2 #include <errno.h>
3 #include <fcntl.h>
4 #include <stdio.h>
5 #include <stdlib.h>
6 #include <string.h>
7 #ifndef WITHOUT_URING
8 #include <liburing.h>
9 #endif
10 #include "dprintf.h"
11 #include "io_uring_engine.h"
12
13 #include <functional>
14 #include <iosfwd>
15 #include <string>
16 #include <sys/stat.h>
17 #include <unistd.h>
18 #include <utility>
19
20 using namespace std;
21
22 IOUringEngine::IOUringEngine(size_t slop_bytes)
23         : slop_bytes(slop_bytes)
24 {
25 #ifdef WITHOUT_URING
26         int ret = -1;
27         dprintf("Compiled without liburing support; not using io_uring.\n");
28 #else
29         int ret = io_uring_queue_init(queue_depth, &ring, 0);
30         if (ret < 0) {
31                 dprintf("io_uring_queue_init() failed; not using io_uring.\n");
32         }
33 #endif
34         using_uring = (ret >= 0);
35
36 #ifndef WITHOUT_URING
37         if (using_uring) {
38                 io_uring_probe *probe = io_uring_get_probe_ring(&ring);
39                 supports_stat = (probe != nullptr && io_uring_opcode_supported(probe, IORING_OP_STATX));
40                 if (!supports_stat) {
41                         dprintf("io_uring on this kernel does not support statx(); will do synchronous access checking.\n");
42                 }
43                 free(probe);
44         }
45 #endif
46 }
47
48 void IOUringEngine::submit_stat(const char *path, std::function<void()> cb)
49 {
50         assert(supports_stat);
51
52 #ifndef WITHOUT_URING
53         if (pending_reads < queue_depth) {
54                 io_uring_sqe *sqe = io_uring_get_sqe(&ring);
55                 if (sqe == nullptr) {
56                         fprintf(stderr, "io_uring_get_sqe: %s\n", strerror(errno));
57                         exit(1);
58                 }
59                 submit_stat_internal(sqe, strdup(path), move(cb));
60         } else {
61                 QueuedStat qs;
62                 qs.cb = move(cb);
63                 qs.pathname = strdup(path);
64                 queued_stats.push(move(qs));
65         }
66 #endif
67 }
68
69 void IOUringEngine::submit_read(int fd, size_t len, off_t offset, function<void(string_view)> cb)
70 {
71         if (!using_uring) {
72                 // Synchronous read.
73                 string s;
74                 s.resize(len + slop_bytes);
75                 complete_pread(fd, &s[0], len, offset);
76                 cb(string_view(s.data(), len));
77                 return;
78         }
79
80 #ifndef WITHOUT_URING
81         if (pending_reads < queue_depth) {
82                 io_uring_sqe *sqe = io_uring_get_sqe(&ring);
83                 if (sqe == nullptr) {
84                         fprintf(stderr, "io_uring_get_sqe: %s\n", strerror(errno));
85                         exit(1);
86                 }
87                 submit_read_internal(sqe, fd, len, offset, move(cb));
88         } else {
89                 queued_reads.push(QueuedRead{ fd, len, offset, move(cb) });
90         }
91 #endif
92 }
93
94 #ifndef WITHOUT_URING
95 void IOUringEngine::submit_read_internal(io_uring_sqe *sqe, int fd, size_t len, off_t offset, function<void(string_view)> cb)
96 {
97         void *buf;
98         if (posix_memalign(&buf, /*alignment=*/4096, len + slop_bytes)) {
99                 fprintf(stderr, "Couldn't allocate %zu bytes: %s\n", len, strerror(errno));
100                 exit(1);
101         }
102
103         PendingRead *pending = new PendingRead;
104         pending->op = OP_READ;
105         pending->read_cb = move(cb);
106         pending->read.buf = buf;
107         pending->read.len = len;
108         pending->read.fd = fd;
109         pending->read.offset = offset;
110         pending->read.iov = iovec{ buf, len };
111
112         io_uring_prep_readv(sqe, fd, &pending->read.iov, 1, offset);
113         io_uring_sqe_set_data(sqe, pending);
114         ++pending_reads;
115 }
116
117 void IOUringEngine::submit_stat_internal(io_uring_sqe *sqe, char *path, std::function<void()> cb)
118 {
119         PendingRead *pending = new PendingRead;
120         pending->op = OP_STAT;
121         pending->stat_cb = move(cb);
122         pending->stat.pathname = path;
123         pending->stat.buf = new struct statx;
124
125         io_uring_prep_statx(sqe, /*fd=*/-1, pending->stat.pathname, AT_STATX_SYNC_AS_STAT, STATX_MODE, pending->stat.buf);
126         io_uring_sqe_set_data(sqe, pending);
127         ++pending_reads;
128 }
129 #endif
130
131 void IOUringEngine::finish()
132 {
133         if (!using_uring) {
134                 return;
135         }
136
137 #ifndef WITHOUT_URING
138         bool anything_to_submit = true;
139         while (pending_reads > 0) {
140                 io_uring_cqe *cqe;
141                 if (io_uring_peek_cqe(&ring, &cqe) != 0) {
142                         if (anything_to_submit) {
143                                 // Nothing ready, so submit whatever is pending and then do a blocking wait.
144                                 int ret = io_uring_submit_and_wait(&ring, 1);
145                                 if (ret < 0) {
146                                         fprintf(stderr, "io_uring_submit(queued): %s\n", strerror(-ret));
147                                         exit(1);
148                                 }
149                                 anything_to_submit = false;
150                         } else {
151                                 int ret = io_uring_wait_cqe(&ring, &cqe);
152                                 if (ret < 0) {
153                                         fprintf(stderr, "io_uring_wait_cqe: %s\n", strerror(-ret));
154                                         exit(1);
155                                 }
156                         }
157                 }
158
159                 unsigned head;
160                 io_uring_for_each_cqe(&ring, head, cqe)
161                 {
162                         PendingRead *pending = reinterpret_cast<PendingRead *>(cqe->user_data);
163                         if (pending->op == OP_STAT) {
164                                 io_uring_cqe_seen(&ring, cqe);
165                                 --pending_reads;
166
167                                 size_t old_pending_reads = pending_reads;
168                                 pending->stat_cb();
169                                 free(pending->stat.pathname);
170                                 delete pending->stat.buf;
171                                 delete pending;
172
173                                 if (pending_reads != old_pending_reads) {
174                                         // A new read was made in the callback (and not queued),
175                                         // so we need to re-submit.
176                                         anything_to_submit = true;
177                                 }
178                         } else {
179                                 if (cqe->res <= 0) {
180                                         fprintf(stderr, "async read failed: %s\n", strerror(-cqe->res));
181                                         exit(1);
182                                 }
183
184                                 if (size_t(cqe->res) < pending->read.iov.iov_len) {
185                                         // Incomplete read, so resubmit it.
186                                         pending->read.iov.iov_base = (char *)pending->read.iov.iov_base + cqe->res;
187                                         pending->read.iov.iov_len -= cqe->res;
188                                         pending->read.offset += cqe->res;
189                                         io_uring_cqe_seen(&ring, cqe);
190
191                                         io_uring_sqe *sqe = io_uring_get_sqe(&ring);
192                                         if (sqe == nullptr) {
193                                                 fprintf(stderr, "No free SQE for resubmit; this shouldn't happen.\n");
194                                                 exit(1);
195                                         }
196                                         io_uring_prep_readv(sqe, pending->read.fd, &pending->read.iov, 1, pending->read.offset);
197                                         io_uring_sqe_set_data(sqe, pending);
198                                         anything_to_submit = true;
199                                 } else {
200                                         io_uring_cqe_seen(&ring, cqe);
201                                         --pending_reads;
202
203                                         size_t old_pending_reads = pending_reads;
204                                         pending->read_cb(string_view(reinterpret_cast<char *>(pending->read.buf), pending->read.len));
205                                         free(pending->read.buf);
206                                         delete pending;
207
208                                         if (pending_reads != old_pending_reads) {
209                                                 // A new read was made in the callback (and not queued),
210                                                 // so we need to re-submit.
211                                                 anything_to_submit = true;
212                                         }
213                                 }
214                         }
215                 }
216
217                 // See if there are any queued stats we can submit now.
218                 // Running a stat means we're very close to printing out a match,
219                 // which is more important than reading more blocks from disk.
220                 // (Even if those blocks returned early, they would only generate
221                 // more matches that would be blocked by this one in Serializer.)
222                 // Thus, prioritize stats.
223                 while (!queued_stats.empty() && pending_reads < queue_depth) {
224                         io_uring_sqe *sqe = io_uring_get_sqe(&ring);
225                         if (sqe == nullptr) {
226                                 fprintf(stderr, "io_uring_get_sqe: %s\n", strerror(errno));
227                                 exit(1);
228                         }
229                         QueuedStat &qs = queued_stats.front();
230                         submit_stat_internal(sqe, qs.pathname, move(qs.cb));
231                         queued_stats.pop();
232                         anything_to_submit = true;
233                 }
234
235                 // See if there are any queued reads we can submit now.
236                 while (!queued_reads.empty() && pending_reads < queue_depth) {
237                         io_uring_sqe *sqe = io_uring_get_sqe(&ring);
238                         if (sqe == nullptr) {
239                                 fprintf(stderr, "io_uring_get_sqe: %s\n", strerror(errno));
240                                 exit(1);
241                         }
242                         QueuedRead &qr = queued_reads.front();
243                         submit_read_internal(sqe, qr.fd, qr.len, qr.offset, move(qr.cb));
244                         queued_reads.pop();
245                         anything_to_submit = true;
246                 }
247         }
248 #endif
249 }
250
251 void complete_pread(int fd, void *ptr, size_t len, off_t offset)
252 {
253         while (len > 0) {
254                 ssize_t ret = pread(fd, ptr, len, offset);
255                 if (ret == -1 && errno == EINTR) {
256                         continue;
257                 }
258                 if (ret <= 0) {
259                         perror("pread");
260                         exit(1);
261                 }
262                 ptr = reinterpret_cast<char *>(ptr) + ret;
263                 len -= ret;
264                 offset -= ret;
265         }
266 }