This repository has been archived on 2022-05-27. You can view files and clone it, but cannot push or open issues or pull requests.
ocaml-cgroups2/examples/cp/cp.c
2021-11-29 09:56:59 +00:00

163 lines
3.2 KiB
C

#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/uio.h>
#include <unistd.h>
#include <liburing.h>
#include <linux/fs.h>
#define QUEUE_DEPTH 16
#define IO_BLOCK_SIZE 4096
#define BUFFER_NUM_BLOCKS 8
/*
* Data struct for SQEs and CQEs.
*/
struct io_data {
bool read;
off_t offset;
int len;
};
/*
* Get size of regular file or error otherwise.
*/
off_t get_file_size(int fd) {
struct stat st;
if (fstat(fd, &st) < 0) {
perror("fstat");
return -1;
}
if (S_ISREG(st.st_mode)) {
return st.st_size;
} else {
return -1;
}
}
/*
* Submit a linked pair of read/write requests
*/
int prep_read_write_pair(struct io_uring *ring, int src_fd, int dst_fd,
off_t offset, void *buf, int len) {
struct io_uring_sqe *sqe;
struct io_data *data;
sqe = io_uring_get_sqe(ring);
if (!sqe)
return 1;
io_uring_prep_read(sqe, src_fd, buf, len, offset);
data = malloc(sizeof(*data));
data->read = 1;
data->len = len;
data->offset = offset;
sqe->flags |= IOSQE_IO_LINK; // add as a dependency of the write
io_uring_sqe_set_data(sqe, data);
sqe = io_uring_get_sqe(ring);
if (!sqe) {
io_uring_submit_and_wait(ring, 1);
sqe = io_uring_get_sqe(ring);
if (!sqe)
return 2;
}
io_uring_prep_write(sqe, dst_fd, buf, len, offset);
data = malloc(sizeof(*data));
data->read = 0;
data->len = len;
data->offset = offset;
io_uring_sqe_set_data(sqe, data);
return 0;
}
/*
* Await and handle completion entries.
*/
int await_cqes(struct io_uring *ring, unsigned int wait_nr) {
struct io_uring_cqe *cqe;
struct io_data *data;
io_uring_submit(ring);
for (int i = 0; i < wait_nr; i++) {
int ret;
if ((ret = io_uring_wait_cqe(ring, &cqe))) {
fprintf(stderr, "io_uring_wait_cqe: %d\n", ret);
return 1;
}
}
return 0;
}
int main(int argc, char **argv) {
// argument validation
if (argc < 3) {
fprintf(stderr, "usage: %s <src> <dst>\n", argv[0]);
return 1;
}
// load files
int src_fd, dst_fd;
off_t src_len;
if ((src_fd = open(argv[1], O_RDONLY)) == -1) {
perror("open");
return 1;
}
if ((dst_fd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC)) == -1) {
perror("open");
return 1;
}
if ((src_len = get_file_size(src_fd)) == -1) {
return 1;
}
// prepare uring
struct io_uring ring;
io_uring_queue_init(QUEUE_DEPTH, &ring, 0);
// perform copy
void *buf;
if ((buf = mmap(0, IO_BLOCK_SIZE * BUFFER_NUM_BLOCKS, PROT_READ | PROT_WRITE,
MAP_SHARED | MAP_ANONYMOUS, -1, 0)) == MAP_FAILED) {
perror("mmap");
return 1;
}
off_t blocks =
(src_len / IO_BLOCK_SIZE) + ((src_len % IO_BLOCK_SIZE) > 0 ? 1 : 0);
for (off_t i = 0; i < blocks; i++) {
int section = i % BUFFER_NUM_BLOCKS;
void *base = buf + (IO_BLOCK_SIZE * section);
off_t len;
if (i == blocks - 1) {
len = src_len % IO_BLOCK_SIZE;
} else {
len = IO_BLOCK_SIZE;
}
prep_read_write_pair(&ring, src_fd, dst_fd, i * IO_BLOCK_SIZE, base, len);
if (await_cqes(&ring, 2)) {
return 1;
}
}
return 0;
}