1 /* SPDX-License-Identifier: MIT */
2 /*
3 * gcc -Wall -O2 -D_GNU_SOURCE -o ucontext-cp ucontext-cp.c -luring
4 */
5 #define _POSIX_C_SOURCE 199309L
6 #include <stdio.h>
7 #include <fcntl.h>
8 #include <string.h>
9 #include <stdlib.h>
10 #include <unistd.h>
11 #include <assert.h>
12 #include <errno.h>
13 #include <ucontext.h>
14 #include <signal.h>
15 #include <inttypes.h>
16 #include <sys/types.h>
17 #include <sys/ioctl.h>
18 #include <sys/timerfd.h>
19 #include <poll.h>
20 #include "liburing.h"
21
22 #define QD 64
23 #define BS 1024
24
25 #ifndef SIGSTKSZ
26 #define SIGSTKSZ 8192
27 #endif
28
29 typedef struct {
30 struct io_uring *ring;
31 unsigned char *stack_buf;
32 ucontext_t ctx_main, ctx_fnew;
33 } async_context;
34
35 typedef struct {
36 async_context *pctx;
37 int *psuccess;
38 int *pfailure;
39 int infd;
40 int outfd;
41 } arguments_bundle;
42
43 #define DEFINE_AWAIT_OP(operation) \
44 static ssize_t await_##operation( \
45 async_context *pctx, \
46 int fd, \
47 const struct iovec *ioves, \
48 unsigned int nr_vecs, \
49 off_t offset) \
50 { \
51 struct io_uring_sqe *sqe = io_uring_get_sqe(pctx->ring); \
52 struct io_uring_cqe *cqe; \
53 \
54 if (!sqe) \
55 return -1; \
56 \
57 io_uring_prep_##operation(sqe, fd, ioves, nr_vecs, offset); \
58 io_uring_sqe_set_data(sqe, pctx); \
59 swapcontext(&pctx->ctx_fnew, &pctx->ctx_main); \
60 io_uring_peek_cqe(pctx->ring, &cqe); \
61 assert(cqe); \
62 io_uring_cqe_seen(pctx->ring, cqe); \
63 \
64 return cqe->res; \
65 }
66
67 DEFINE_AWAIT_OP(readv)
DEFINE_AWAIT_OP(writev)68 DEFINE_AWAIT_OP(writev)
69 #undef DEFINE_AWAIT_OP
70
71 int await_poll(async_context *pctx, int fd, short poll_mask) {
72 struct io_uring_sqe *sqe = io_uring_get_sqe(pctx->ring);
73 struct io_uring_cqe *cqe;
74 if (!sqe)
75 return -1;
76
77 io_uring_prep_poll_add(sqe, fd, poll_mask);
78 io_uring_sqe_set_data(sqe, pctx);
79 swapcontext(&pctx->ctx_fnew, &pctx->ctx_main);
80 io_uring_peek_cqe(pctx->ring, &cqe);
81 assert(cqe);
82 io_uring_cqe_seen(pctx->ring, cqe);
83
84 return cqe->res;
85 }
86
await_delay(async_context * pctx,time_t seconds)87 int await_delay(async_context *pctx, time_t seconds) {
88 struct io_uring_sqe *sqe = io_uring_get_sqe(pctx->ring);
89 struct io_uring_cqe *cqe;
90 struct __kernel_timespec ts = {
91 .tv_sec = seconds,
92 .tv_nsec = 0
93 };
94
95 if (!sqe)
96 return -1;
97
98 io_uring_prep_timeout(sqe, &ts, 0, 0);
99 io_uring_sqe_set_data(sqe, pctx);
100 swapcontext(&pctx->ctx_fnew, &pctx->ctx_main);
101 io_uring_peek_cqe(pctx->ring, &cqe);
102 assert(cqe);
103 io_uring_cqe_seen(pctx->ring, cqe);
104
105 return 0;
106 }
107
setup_context(async_context * pctx,struct io_uring * ring)108 static int setup_context(async_context *pctx, struct io_uring *ring)
109 {
110 int ret;
111
112 pctx->ring = ring;
113 ret = getcontext(&pctx->ctx_fnew);
114 if (ret < 0) {
115 perror("getcontext");
116 return -1;
117 }
118 pctx->stack_buf = malloc(SIGSTKSZ);
119 if (!pctx->stack_buf) {
120 perror("malloc");
121 return -1;
122 }
123 pctx->ctx_fnew.uc_stack.ss_sp = pctx->stack_buf;
124 pctx->ctx_fnew.uc_stack.ss_size = SIGSTKSZ;
125 pctx->ctx_fnew.uc_link = &pctx->ctx_main;
126
127 return 0;
128 }
129
copy_file(async_context * pctx,int infd,int outfd,struct iovec * piov)130 static int copy_file(async_context *pctx, int infd, int outfd, struct iovec* piov)
131 {
132 off_t offset = 0;
133
134 for (;;) {
135 ssize_t bytes_read;
136
137 printf("%d->%d: readv %ld bytes from %ld\n", infd, outfd, (long) piov->iov_len, (long) offset);
138 if ((bytes_read = await_readv(pctx, infd, piov, 1, offset)) < 0) {
139 perror("await_readv");
140 return 1;
141 }
142 if (bytes_read == 0)
143 return 0;
144
145 piov->iov_len = bytes_read;
146
147 printf("%d->%d: writev %ld bytes from %ld\n", infd, outfd, (long) piov->iov_len, (long) offset);
148 if (await_writev(pctx, outfd, piov, 1, offset) != bytes_read) {
149 perror("await_writev");
150 return 1;
151 }
152 if (bytes_read < BS)
153 return 0;
154 offset += bytes_read;
155
156 printf("%d->%d: wait %ds\n", infd, outfd, 1);
157 await_delay(pctx, 1);
158 }
159 }
160
copy_file_wrapper(arguments_bundle * pbundle)161 static void copy_file_wrapper(arguments_bundle *pbundle)
162 {
163 struct iovec iov = {
164 .iov_base = malloc(BS),
165 .iov_len = BS,
166 };
167 async_context *pctx = pbundle->pctx;
168
169 int ret = copy_file(pctx, pbundle->infd, pbundle->outfd, &iov);
170
171 printf("%d->%d: done with ret code %d\n", pbundle->infd, pbundle->outfd, ret);
172
173 if (ret == 0) {
174 ++*pbundle->psuccess;
175 } else {
176 ++*pbundle->pfailure;
177 }
178
179 free(iov.iov_base);
180 close(pbundle->infd);
181 close(pbundle->outfd);
182 free(pbundle->pctx->stack_buf);
183 free(pbundle->pctx);
184 free(pbundle);
185
186 swapcontext(&pctx->ctx_fnew, &pctx->ctx_main);
187 }
188
main(int argc,char * argv[])189 int main(int argc, char *argv[])
190 {
191 struct io_uring ring;
192 int i, req_count, ret;
193 int success = 0, failure = 0;
194
195 if (argc < 3) {
196 fprintf(stderr, "%s: infile1 outfile1 [infile2 outfile2 [...]]\n", argv[0]);
197 return 1;
198 }
199
200 ret = io_uring_queue_init(QD, &ring, 0);
201 if (ret < 0) {
202 fprintf(stderr, "queue_init: %s\n", strerror(-ret));
203 return -1;
204 }
205
206 req_count = (argc - 1) / 2;
207 printf("copying %d files...\n", req_count);
208
209 for (i = 1; i < argc; i += 2) {
210 int infd, outfd;
211
212 async_context *pctx = malloc(sizeof(*pctx));
213
214 if (!pctx || setup_context(pctx, &ring))
215 return 1;
216
217 infd = open(argv[i], O_RDONLY);
218 if (infd < 0) {
219 perror("open infile");
220 return 1;
221 }
222 outfd = open(argv[i + 1], O_WRONLY | O_CREAT | O_TRUNC, 0644);
223 if (outfd < 0) {
224 perror("open outfile");
225 return 1;
226 }
227
228 arguments_bundle *pbundle = malloc(sizeof(*pbundle));
229 pbundle->pctx = pctx;
230 pbundle->psuccess = &success;
231 pbundle->pfailure = &failure;
232 pbundle->infd = infd;
233 pbundle->outfd = outfd;
234
235 makecontext(&pctx->ctx_fnew, (void (*)(void)) copy_file_wrapper, 1, pbundle);
236
237 if (swapcontext(&pctx->ctx_main, &pctx->ctx_fnew)) {
238 perror("swapcontext");
239 return 1;
240 }
241 }
242
243 /* event loop */
244 while (success + failure < req_count) {
245 struct io_uring_cqe *cqe;
246
247 /* usually be timed waiting */
248 ret = io_uring_submit_and_wait(&ring, 1);
249 if (ret < 0) {
250 fprintf(stderr, "submit_and_wait: %s\n", strerror(-ret));
251 return 1;
252 }
253
254 ret = io_uring_wait_cqe(&ring, &cqe);
255 if (ret < 0) {
256 fprintf(stderr, "wait_cqe: %s\n", strerror(-ret));
257 return 1;
258 }
259
260 async_context *pctx = io_uring_cqe_get_data(cqe);
261
262 if (swapcontext(&pctx->ctx_main, &pctx->ctx_fnew)) {
263 perror("swapcontext");
264 return 1;
265 }
266 }
267
268 io_uring_queue_exit(&ring);
269
270 printf("finished with %d success(es) and %d failure(s)\n", success, failure);
271
272 return failure > 0;
273 }
274