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