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
25static int infd, outfd;
26
27struct io_data {
28	int read;
29	off_t first_offset, offset;
30	size_t first_len;
31	struct iovec iov;
32};
33
34static 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
47static 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
69static 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
84static 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
111static 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
123static 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		unsigned long had_reads;
135		int got_comp;
136
137		/*
138		 * Queue up as many reads as we can
139		 */
140		had_reads = reads;
141		while (insize) {
142			off_t this_size = insize;
143
144			if (reads + writes >= QD)
145				break;
146			if (this_size > BS)
147				this_size = BS;
148			else if (!this_size)
149				break;
150
151			if (queue_read(ring, this_size, offset))
152				break;
153
154			insize -= this_size;
155			offset += this_size;
156			reads++;
157		}
158
159		if (had_reads != reads) {
160			ret = io_uring_submit(ring);
161			if (ret < 0) {
162				fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret));
163				break;
164			}
165		}
166
167		/*
168		 * Queue is full at this point. Find at least one completion.
169		 */
170		got_comp = 0;
171		while (write_left) {
172			struct io_data *data;
173
174			if (!got_comp) {
175				ret = io_uring_wait_cqe(ring, &cqe);
176				got_comp = 1;
177			} else
178				ret = io_uring_peek_cqe(ring, &cqe);
179			if (ret < 0) {
180				fprintf(stderr, "io_uring_peek_cqe: %s\n",
181							strerror(-ret));
182				return 1;
183			}
184			if (!cqe)
185				break;
186
187			data = io_uring_cqe_get_data(cqe);
188			if (cqe->res < 0) {
189				if (cqe->res == -EAGAIN) {
190					queue_prepped(ring, data);
191					io_uring_cqe_seen(ring, cqe);
192					continue;
193				}
194				fprintf(stderr, "cqe failed: %s\n",
195						strerror(-cqe->res));
196				return 1;
197			} else if ((size_t) cqe->res != data->iov.iov_len) {
198				/* Short read/write, adjust and requeue */
199				data->iov.iov_base += cqe->res;
200				data->iov.iov_len -= cqe->res;
201				data->offset += cqe->res;
202				queue_prepped(ring, data);
203				io_uring_cqe_seen(ring, cqe);
204				continue;
205			}
206
207			/*
208			 * All done. if write, nothing else to do. if read,
209			 * queue up corresponding write.
210			 */
211			if (data->read) {
212				queue_write(ring, data);
213				write_left -= data->first_len;
214				reads--;
215				writes++;
216			} else {
217				free(data);
218				writes--;
219			}
220			io_uring_cqe_seen(ring, cqe);
221		}
222	}
223
224	return 0;
225}
226
227int main(int argc, char *argv[])
228{
229	struct io_uring ring;
230	off_t insize;
231	int ret;
232
233	if (argc < 3) {
234		printf("%s: infile outfile\n", argv[0]);
235		return 1;
236	}
237
238	infd = open(argv[1], O_RDONLY);
239	if (infd < 0) {
240		perror("open infile");
241		return 1;
242	}
243	outfd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0644);
244	if (outfd < 0) {
245		perror("open outfile");
246		return 1;
247	}
248
249	if (setup_context(QD, &ring))
250		return 1;
251	if (get_file_size(infd, &insize))
252		return 1;
253
254	ret = copy_file(&ring, insize);
255
256	close(infd);
257	close(outfd);
258	io_uring_queue_exit(&ring);
259	return ret;
260}
261