tools/io_uring/io_uring-cp.c

0001 // SPDX-License-Identifier: GPL-2.0
0002 /*
0003  * Simple test program that demonstrates a file copy through io_uring. This
0004  * uses the API exposed by liburing.
0005  *
0006  * Copyright (C) 2018-2019 Jens Axboe
0007  */
0008 #include <stdio.h>
0009 #include <fcntl.h>
0010 #include <string.h>
0011 #include <stdlib.h>
0012 #include <unistd.h>
0013 #include <assert.h>
0014 #include <errno.h>
0015 #include <inttypes.h>
0016 #include <sys/types.h>
0017 #include <sys/stat.h>
0018 #include <sys/ioctl.h>
0019
0020 #include "liburing.h"
0021
0022 #define QD  64
0023 #define BS  (32*1024)
0024
0025 static int infd, outfd;
0026
0027 struct io_data {
0028     int read;
0029     off_t first_offset, offset;
0030     size_t first_len;
0031     struct iovec iov;
0032 };
0033
0034 static int setup_context(unsigned entries, struct io_uring *ring)
0035 {
0036     int ret;
0037
0038     ret = io_uring_queue_init(entries, ring, 0);
0039     if (ret < 0) {
0040         fprintf(stderr, "queue_init: %s\n", strerror(-ret));
0041         return -1;
0042     }
0043
0044     return 0;
0045 }
0046
0047 static int get_file_size(int fd, off_t *size)
0048 {
0049     struct stat st;
0050
0051     if (fstat(fd, &st) < 0)
0052         return -1;
0053     if (S_ISREG(st.st_mode)) {
0054         *size = st.st_size;
0055         return 0;
0056     } else if (S_ISBLK(st.st_mode)) {
0057         unsigned long long bytes;
0058
0059         if (ioctl(fd, BLKGETSIZE64, &bytes) != 0)
0060             return -1;
0061
0062         *size = bytes;
0063         return 0;
0064     }
0065
0066     return -1;
0067 }
0068
0069 static void queue_prepped(struct io_uring *ring, struct io_data *data)
0070 {
0071     struct io_uring_sqe *sqe;
0072
0073     sqe = io_uring_get_sqe(ring);
0074     assert(sqe);
0075
0076     if (data->read)
0077         io_uring_prep_readv(sqe, infd, &data->iov, 1, data->offset);
0078     else
0079         io_uring_prep_writev(sqe, outfd, &data->iov, 1, data->offset);
0080
0081     io_uring_sqe_set_data(sqe, data);
0082 }
0083
0084 static int queue_read(struct io_uring *ring, off_t size, off_t offset)
0085 {
0086     struct io_uring_sqe *sqe;
0087     struct io_data *data;
0088
0089     data = malloc(size + sizeof(*data));
0090     if (!data)
0091         return 1;
0092
0093     sqe = io_uring_get_sqe(ring);
0094     if (!sqe) {
0095         free(data);
0096         return 1;
0097     }
0098
0099     data->read = 1;
0100     data->offset = data->first_offset = offset;
0101
0102     data->iov.iov_base = data + 1;
0103     data->iov.iov_len = size;
0104     data->first_len = size;
0105
0106     io_uring_prep_readv(sqe, infd, &data->iov, 1, offset);
0107     io_uring_sqe_set_data(sqe, data);
0108     return 0;
0109 }
0110
0111 static void queue_write(struct io_uring *ring, struct io_data *data)
0112 {
0113     data->read = 0;
0114     data->offset = data->first_offset;
0115
0116     data->iov.iov_base = data + 1;
0117     data->iov.iov_len = data->first_len;
0118
0119     queue_prepped(ring, data);
0120     io_uring_submit(ring);
0121 }
0122
0123 static int copy_file(struct io_uring *ring, off_t insize)
0124 {
0125     unsigned long reads, writes;
0126     struct io_uring_cqe *cqe;
0127     off_t write_left, offset;
0128     int ret;
0129
0130     write_left = insize;
0131     writes = reads = offset = 0;
0132
0133     while (insize || write_left) {
0134         int had_reads, got_comp;
0135
0136         /*
0137          * Queue up as many reads as we can
0138          */
0139         had_reads = reads;
0140         while (insize) {
0141             off_t this_size = insize;
0142
0143             if (reads + writes >= QD)
0144                 break;
0145             if (this_size > BS)
0146                 this_size = BS;
0147             else if (!this_size)
0148                 break;
0149
0150             if (queue_read(ring, this_size, offset))
0151                 break;
0152
0153             insize -= this_size;
0154             offset += this_size;
0155             reads++;
0156         }
0157
0158         if (had_reads != reads) {
0159             ret = io_uring_submit(ring);
0160             if (ret < 0) {
0161                 fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret));
0162                 break;
0163             }
0164         }
0165
0166         /*
0167          * Queue is full at this point. Find at least one completion.
0168          */
0169         got_comp = 0;
0170         while (write_left) {
0171             struct io_data *data;
0172
0173             if (!got_comp) {
0174                 ret = io_uring_wait_cqe(ring, &cqe);
0175                 got_comp = 1;
0176             } else {
0177                 ret = io_uring_peek_cqe(ring, &cqe);
0178                 if (ret == -EAGAIN) {
0179                     cqe = NULL;
0180                     ret = 0;
0181                 }
0182             }
0183             if (ret < 0) {
0184                 fprintf(stderr, "io_uring_peek_cqe: %s\n",
0185                             strerror(-ret));
0186                 return 1;
0187             }
0188             if (!cqe)
0189                 break;
0190
0191             data = io_uring_cqe_get_data(cqe);
0192             if (cqe->res < 0) {
0193                 if (cqe->res == -EAGAIN) {
0194                     queue_prepped(ring, data);
0195                     io_uring_cqe_seen(ring, cqe);
0196                     continue;
0197                 }
0198                 fprintf(stderr, "cqe failed: %s\n",
0199                         strerror(-cqe->res));
0200                 return 1;
0201             } else if (cqe->res != data->iov.iov_len) {
0202                 /* Short read/write, adjust and requeue */
0203                 data->iov.iov_base += cqe->res;
0204                 data->iov.iov_len -= cqe->res;
0205                 data->offset += cqe->res;
0206                 queue_prepped(ring, data);
0207                 io_uring_cqe_seen(ring, cqe);
0208                 continue;
0209             }
0210
0211             /*
0212              * All done. if write, nothing else to do. if read,
0213              * queue up corresponding write.
0214              */
0215             if (data->read) {
0216                 queue_write(ring, data);
0217                 write_left -= data->first_len;
0218                 reads--;
0219                 writes++;
0220             } else {
0221                 free(data);
0222                 writes--;
0223             }
0224             io_uring_cqe_seen(ring, cqe);
0225         }
0226     }
0227
0228     /* wait out pending writes */
0229     while (writes) {
0230         struct io_data *data;
0231
0232         ret = io_uring_wait_cqe(ring, &cqe);
0233         if (ret) {
0234             fprintf(stderr, "wait_cqe=%d\n", ret);
0235             return 1;
0236         }
0237         if (cqe->res < 0) {
0238             fprintf(stderr, "write res=%d\n", cqe->res);
0239             return 1;
0240         }
0241         data = io_uring_cqe_get_data(cqe);
0242         free(data);
0243         writes--;
0244         io_uring_cqe_seen(ring, cqe);
0245     }
0246
0247     return 0;
0248 }
0249
0250 int main(int argc, char *argv[])
0251 {
0252     struct io_uring ring;
0253     off_t insize;
0254     int ret;
0255
0256     if (argc < 3) {
0257         printf("%s: infile outfile\n", argv[0]);
0258         return 1;
0259     }
0260
0261     infd = open(argv[1], O_RDONLY);
0262     if (infd < 0) {
0263         perror("open infile");
0264         return 1;
0265     }
0266     outfd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0644);
0267     if (outfd < 0) {
0268         perror("open outfile");
0269         return 1;
0270     }
0271
0272     if (setup_context(QD, &ring))
0273         return 1;
0274     if (get_file_size(infd, &insize))
0275         return 1;
0276
0277     ret = copy_file(&ring, insize);
0278
0279     close(infd);
0280     close(outfd);
0281     io_uring_queue_exit(&ring);
0282     return ret;
0283 }