-
-
Save mathslimin/87adcdb25f778d1c743e4e53c733ab89 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#include <stdio.h> | |
#include <fcntl.h> | |
#include <string.h> | |
#include <stdlib.h> | |
#include <unistd.h> | |
#include <assert.h> | |
#include <errno.h> | |
#include <sys/stat.h> | |
#include <sys/ioctl.h> | |
#include <liburing.h> | |
#define QD 2 | |
#define BS (16 * 1024) | |
static int infd, outfd; | |
static int flags; | |
struct io_data { | |
int read; | |
off_t first_offset, offset; | |
size_t first_len; | |
struct iovec iov; | |
}; | |
static int setup_context(unsigned entries, struct io_uring *ring) { | |
int ret; | |
//flags = IORING_SETUP_SQPOLL; | |
flags = 0; | |
ret = io_uring_queue_init(entries, ring, flags); | |
if( ret < 0) { | |
fprintf(stderr, "queue_init: %s\n", strerror(-ret)); | |
return -1; | |
} | |
return 0; | |
} | |
static int get_file_size(int fd, off_t *size) { | |
struct stat st; | |
if (fstat(fd, &st) < 0 ) | |
return -1; | |
if(S_ISREG(st.st_mode)) { | |
*size = st.st_size; | |
return 0; | |
} else if (S_ISBLK(st.st_mode)) { | |
unsigned long long bytes; | |
if (ioctl(fd, BLKGETSIZE64, &bytes) != 0) | |
return -1; | |
*size = bytes; | |
return 0; | |
} | |
return -1; | |
} | |
static void queue_prepped(struct io_uring *ring, struct io_data *data) { | |
struct io_uring_sqe *sqe; | |
sqe = io_uring_get_sqe(ring); | |
assert(sqe); | |
if (data->read) | |
io_uring_prep_readv(sqe, infd, &data->iov, 1, data->offset); | |
else | |
io_uring_prep_writev(sqe, outfd, &data->iov, 1, data->offset); | |
io_uring_sqe_set_data(sqe, data); | |
} | |
static int queue_read(struct io_uring *ring, off_t size, off_t offset) { | |
struct io_uring_sqe *sqe; | |
struct io_data *data; | |
data = malloc(size + sizeof(*data)); | |
if (!data) | |
return 1; | |
sqe = io_uring_get_sqe(ring); | |
if (!sqe) { | |
free(data); | |
return 1; | |
} | |
data->read = 1; | |
data->offset = data->first_offset = offset; | |
data->iov.iov_base = data + 1; | |
data->iov.iov_len = size; | |
data->first_len = size; | |
io_uring_prep_readv(sqe, infd, &data->iov, 1, offset); | |
io_uring_sqe_set_data(sqe, data); | |
return 0; | |
} | |
static void queue_write(struct io_uring *ring, struct io_data *data) { | |
data->read = 0; | |
data->offset = data->first_offset; | |
data->iov.iov_base = data + 1; | |
data->iov.iov_len = data->first_len; | |
queue_prepped(ring, data); | |
io_uring_submit(ring); | |
} | |
int copy_file(struct io_uring *ring, off_t insize) { | |
unsigned long reads, writes; | |
struct io_uring_cqe *cqe; | |
struct io_uring_sqe *sqe; | |
off_t write_left, offset; | |
int ret; | |
write_left = insize; | |
writes = reads = offset = 0; | |
while (insize || write_left) { | |
int had_reads, got_comp; | |
/* Queue up as many reads as we can */ | |
had_reads = reads; | |
while (insize) { | |
off_t this_size = insize; | |
if (reads + writes >= QD) | |
break; | |
if (this_size > BS) | |
this_size = BS; | |
else if (!this_size) | |
break; | |
if (queue_read(ring, this_size, offset)) | |
break; | |
insize -= this_size; | |
offset += this_size; | |
reads++; | |
} | |
if (had_reads != reads) { | |
ret = io_uring_submit(ring); | |
if (ret < 0) { | |
fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret)); | |
break; | |
} | |
} | |
/* Queue is full at this point. Let's find at least one completion */ | |
got_comp = 0; | |
while (write_left) { | |
struct io_data *data; | |
if (!got_comp) { | |
ret = io_uring_wait_cqe(ring, &cqe); | |
got_comp = 1; | |
} else { | |
ret = io_uring_peek_cqe(ring, &cqe); | |
if (ret == -EAGAIN) { | |
cqe = NULL; | |
ret = 0; | |
} | |
} | |
if (ret < 0) { | |
fprintf(stderr, "io_uring_peek_cqe: %s\n", | |
strerror(-ret)); | |
return 1; | |
} | |
if (!cqe) | |
break; | |
/* | |
* need to call io_uring_enter() to make the kernel notice the new IO | |
* if polled and the thread is now sleeping. | |
*/ | |
sqe = io_uring_get_sqe(ring); | |
if (sqe == NULL) { | |
perror("io_uring_get_sqe"); | |
exit(1); | |
} | |
if ((sqe->flags) & IORING_SQ_NEED_WAKEUP) { | |
printf("IORING_ENTER_SQ_WAKEUP\n"); | |
io_uring_enter(ring, 1, 1, IORING_ENTER_SQ_WAKEUP); | |
} | |
data = io_uring_cqe_get_data(cqe); | |
if (cqe->res < 0) { | |
if (cqe->res == -EAGAIN) { | |
queue_prepped(ring, data); | |
io_uring_cqe_seen(ring, cqe); | |
continue; | |
} | |
fprintf(stderr, "cqe failed: %s\n", | |
strerror(-cqe->res)); | |
return 1; | |
} else if (cqe->res != data->iov.iov_len) { | |
/* short read/write; adjust and requeue */ | |
data->iov.iov_base += cqe->res; | |
data->iov.iov_len -= cqe->res; | |
queue_prepped(ring, data); | |
io_uring_cqe_seen(ring, cqe); | |
continue; | |
} | |
/* | |
* All done. If write, nothing else to do. If read, | |
* queue up corresponding write. | |
* */ | |
if (data->read) { | |
queue_write(ring, data); | |
write_left -= data->first_len; | |
reads--; | |
writes++; | |
} else { | |
free(data); | |
writes--; | |
} | |
io_uring_cqe_seen(ring, cqe); | |
} | |
} | |
return 0; | |
} | |
int main(int argc, char *argv[]) { | |
struct io_uring ring; | |
off_t insize; | |
int ret; | |
if (argc < 3) { | |
printf("Usage: %s <infile> <outfile>\n", argv[0]); | |
return 1; | |
} | |
infd = open(argv[1], O_RDONLY); | |
if (infd < 0) { | |
perror("open infile"); | |
return 1; | |
} | |
outfd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0644); | |
if (outfd < 0) { | |
perror("open outfile"); | |
return 1; | |
} | |
if (setup_context(QD, &ring)) | |
return 1; | |
if (get_file_size(infd, &insize)) | |
return 1; | |
ret = copy_file(&ring, insize); | |
close(infd); | |
close(outfd); | |
io_uring_queue_exit(&ring); | |
return ret; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment