1 /* SPDX-License-Identifier: MIT */
2 /*
3 * gcc -Wall -O2 -D_GNU_SOURCE -o io_uring-cp io_uring-cp.c -luring
4 */
5 #include <stdio.h>
6 #include <fcntl.h>
7 #include <string.h>
8 #include <stdlib.h>
9 #include <unistd.h>
10 #include <assert.h>
11 #include <errno.h>
12 #include <inttypes.h>
13 #include <sys/types.h>
14 #include <sys/stat.h>
15 #include <sys/ioctl.h>
16 #include "liburing.h"
17
18 #define QD 64
19 #define BS (32*1024)
20
21 static int infd, outfd;
22
23 struct io_data {
24 int read;
25 off_t first_offset, offset;
26 size_t first_len;
27 struct iovec iov;
28 };
29
setup_context(unsigned entries,struct io_uring * ring)30 static int setup_context(unsigned entries, struct io_uring *ring)
31 {
32 int ret;
33
34 ret = io_uring_queue_init(entries, ring, 0);
35 if (ret < 0) {
36 fprintf(stderr, "queue_init: %s\n", strerror(-ret));
37 return -1;
38 }
39
40 return 0;
41 }
42
get_file_size(int fd,off_t * size)43 static int get_file_size(int fd, off_t *size)
44 {
45 struct stat st;
46
47 if (fstat(fd, &st) < 0)
48 return -1;
49 if (S_ISREG(st.st_mode)) {
50 *size = st.st_size;
51 return 0;
52 } else if (S_ISBLK(st.st_mode)) {
53 unsigned long long bytes;
54
55 if (ioctl(fd, BLKGETSIZE64, &bytes) != 0)
56 return -1;
57
58 *size = bytes;
59 return 0;
60 }
61
62 return -1;
63 }
64
queue_prepped(struct io_uring * ring,struct io_data * data)65 static void queue_prepped(struct io_uring *ring, struct io_data *data)
66 {
67 struct io_uring_sqe *sqe;
68
69 sqe = io_uring_get_sqe(ring);
70 assert(sqe);
71
72 if (data->read)
73 io_uring_prep_readv(sqe, infd, &data->iov, 1, data->offset);
74 else
75 io_uring_prep_writev(sqe, outfd, &data->iov, 1, data->offset);
76
77 io_uring_sqe_set_data(sqe, data);
78 }
79
queue_read(struct io_uring * ring,off_t size,off_t offset)80 static int queue_read(struct io_uring *ring, off_t size, off_t offset)
81 {
82 struct io_uring_sqe *sqe;
83 struct io_data *data;
84
85 data = malloc(size + sizeof(*data));
86 if (!data)
87 return 1;
88
89 sqe = io_uring_get_sqe(ring);
90 if (!sqe) {
91 free(data);
92 return 1;
93 }
94
95 data->read = 1;
96 data->offset = data->first_offset = offset;
97
98 data->iov.iov_base = data + 1;
99 data->iov.iov_len = size;
100 data->first_len = size;
101
102 io_uring_prep_readv(sqe, infd, &data->iov, 1, offset);
103 io_uring_sqe_set_data(sqe, data);
104 return 0;
105 }
106
queue_write(struct io_uring * ring,struct io_data * data)107 static void queue_write(struct io_uring *ring, struct io_data *data)
108 {
109 data->read = 0;
110 data->offset = data->first_offset;
111
112 data->iov.iov_base = data + 1;
113 data->iov.iov_len = data->first_len;
114
115 queue_prepped(ring, data);
116 io_uring_submit(ring);
117 }
118
copy_file(struct io_uring * ring,off_t insize)119 static int copy_file(struct io_uring *ring, off_t insize)
120 {
121 unsigned long reads, writes;
122 struct io_uring_cqe *cqe;
123 off_t write_left, offset;
124 int ret;
125
126 write_left = insize;
127 writes = reads = offset = 0;
128
129 while (insize || write_left) {
130 unsigned long had_reads;
131 int got_comp;
132
133 /*
134 * Queue up as many reads as we can
135 */
136 had_reads = reads;
137 while (insize) {
138 off_t this_size = insize;
139
140 if (reads + writes >= QD)
141 break;
142 if (this_size > BS)
143 this_size = BS;
144 else if (!this_size)
145 break;
146
147 if (queue_read(ring, this_size, offset))
148 break;
149
150 insize -= this_size;
151 offset += this_size;
152 reads++;
153 }
154
155 if (had_reads != reads) {
156 ret = io_uring_submit(ring);
157 if (ret < 0) {
158 fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret));
159 break;
160 }
161 }
162
163 /*
164 * Queue is full at this point. Find at least one completion.
165 */
166 got_comp = 0;
167 while (write_left) {
168 struct io_data *data;
169
170 if (!got_comp) {
171 ret = io_uring_wait_cqe(ring, &cqe);
172 got_comp = 1;
173 } else {
174 ret = io_uring_peek_cqe(ring, &cqe);
175 if (ret == -EAGAIN) {
176 cqe = NULL;
177 ret = 0;
178 }
179 }
180 if (ret < 0) {
181 fprintf(stderr, "io_uring_peek_cqe: %s\n",
182 strerror(-ret));
183 return 1;
184 }
185 if (!cqe)
186 break;
187
188 data = io_uring_cqe_get_data(cqe);
189 if (cqe->res < 0) {
190 if (cqe->res == -EAGAIN) {
191 queue_prepped(ring, data);
192 io_uring_submit(ring);
193 io_uring_cqe_seen(ring, cqe);
194 continue;
195 }
196 fprintf(stderr, "cqe failed: %s\n",
197 strerror(-cqe->res));
198 return 1;
199 } else if ((size_t)cqe->res != data->iov.iov_len) {
200 /* Short read/write, adjust and requeue */
201 data->iov.iov_base += cqe->res;
202 data->iov.iov_len -= cqe->res;
203 data->offset += cqe->res;
204 queue_prepped(ring, data);
205 io_uring_submit(ring);
206 io_uring_cqe_seen(ring, cqe);
207 continue;
208 }
209
210 /*
211 * All done. if write, nothing else to do. if read,
212 * queue up corresponding write.
213 */
214 if (data->read) {
215 queue_write(ring, data);
216 write_left -= data->first_len;
217 reads--;
218 writes++;
219 } else {
220 free(data);
221 writes--;
222 }
223 io_uring_cqe_seen(ring, cqe);
224 }
225 }
226
227 /* wait out pending writes */
228 while (writes) {
229 struct io_data *data;
230
231 ret = io_uring_wait_cqe(ring, &cqe);
232 if (ret) {
233 fprintf(stderr, "wait_cqe=%d\n", ret);
234 return 1;
235 }
236 if (cqe->res < 0) {
237 fprintf(stderr, "write res=%d\n", cqe->res);
238 return 1;
239 }
240 data = io_uring_cqe_get_data(cqe);
241 free(data);
242 writes--;
243 io_uring_cqe_seen(ring, cqe);
244 }
245
246 return 0;
247 }
248
main(int argc,char * argv[])249 int main(int argc, char *argv[])
250 {
251 struct io_uring ring;
252 off_t insize;
253 int ret;
254
255 if (argc < 3) {
256 printf("%s: infile outfile\n", argv[0]);
257 return 1;
258 }
259
260 infd = open(argv[1], O_RDONLY);
261 if (infd < 0) {
262 perror("open infile");
263 return 1;
264 }
265 outfd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0644);
266 if (outfd < 0) {
267 perror("open outfile");
268 return 1;
269 }
270
271 if (setup_context(QD, &ring))
272 return 1;
273 if (get_file_size(infd, &insize))
274 return 1;
275
276 ret = copy_file(&ring, insize);
277
278 close(infd);
279 close(outfd);
280 io_uring_queue_exit(&ring);
281 return ret;
282 }
283