tangetools/iothrottle/iothrottle.c

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <dlfcn.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <time.h>
#include <sys/uio.h>
#include <sys/sendfile.h>
#include <linux/fs.h>

// Function pointers for original read and write and other relevant calls
ssize_t (*original_read)(int fd, void *buf, size_t count);
ssize_t (*original_write)(int fd, const void *buf, size_t count);
size_t (*original_fwrite)(const void *ptr, size_t size, size_t nmemb, FILE *stream);
int (*original_fflush)(FILE *stream);
ssize_t (*original_pread)(int fd, void *buf, size_t count, off_t offset);
ssize_t (*original_pwrite)(int fd, const void *buf, size_t count, off_t offset);
ssize_t (*original_readv)(int fd, const struct iovec *iov, int iovcnt);
ssize_t (*original_writev)(int fd, const struct iovec *iov, int iovcnt);
off_t (*original_sendfile)(int out_fd, int in_fd, off_t *offset, size_t count);
ssize_t (*original_splice)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);
ssize_t (*original_copy_file_range)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);

// Rate limits in bytes per second
static size_t read_limit = 0; // 0 = no limit
static size_t write_limit = 0;
static int debug_mode = 0;

// Tracking variables
static size_t read_bytes = 0;
static size_t write_bytes = 0;
static struct timespec read_start, write_start;

void throttle(size_t *bytes, size_t limit, struct timespec *start) {
    if (limit == 0) {
        return; // No limit applied
    }

    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);

    double elapsed = (now.tv_sec - start->tv_sec) + (now.tv_nsec - start->tv_nsec) / 1e9;

    double allowed_bytes = limit * elapsed;
    if (debug_mode) fprintf(stderr, "allowed: %ld written: %ld\n",
			    (long)allowed_bytes, *bytes);
    if (*bytes > allowed_bytes) {
      *bytes -= allowed_bytes;
      double sleep_time = *bytes / (double)limit;

      if (debug_mode) fprintf(stderr, "sleep: %f\n", sleep_time);
      usleep((useconds_t)(sleep_time * 1e6));
      *start = now;
    }
}

size_t fwrite(const void *ptr, size_t size, size_t nmemb, FILE *stream) {
    if (!original_fwrite) {
        original_fwrite = dlsym(RTLD_NEXT, "fwrite");
        if (!original_fwrite) {
            if (debug_mode) fprintf(stderr, "Error loading original fwrite: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    size_t total_bytes = size * nmemb;
    throttle(&write_bytes, write_limit, &write_start);

    size_t result = original_fwrite(ptr, size, nmemb, stream);
    if (result > 0) {
        write_bytes += total_bytes;
    }
    return result;
}

int fflush(FILE *stream) {
    if (!original_fflush) {
        original_fflush = dlsym(RTLD_NEXT, "fflush");
        if (!original_fflush) {
            if (debug_mode) fprintf(stderr, "Error loading original fflush: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    return original_fflush(stream);
}

ssize_t read(int fd, void *buf, size_t count) {
    if (!original_read) {
        original_read = dlsym(RTLD_NEXT, "read");
        if (!original_read) {
            if (debug_mode) fprintf(stderr, "Error loading original read: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_read(fd, buf, count);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t write(int fd, const void *buf, size_t count) {
    if (!original_write) {
        original_write = dlsym(RTLD_NEXT, "write");
        if (!original_write) {
            if (debug_mode) fprintf(stderr, "Error loading original write: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_write(fd, buf, count);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t pread(int fd, void *buf, size_t count, off_t offset) {
    if (!original_pread) {
        original_pread = dlsym(RTLD_NEXT, "pread");
        if (!original_pread) {
            if (debug_mode) fprintf(stderr, "Error loading original pread: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_pread(fd, buf, count, offset);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t pwrite(int fd, const void *buf, size_t count, off_t offset) {
    if (!original_pwrite) {
        original_pwrite = dlsym(RTLD_NEXT, "pwrite");
        if (!original_pwrite) {
            if (debug_mode) fprintf(stderr, "Error loading original pwrite: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_pwrite(fd, buf, count, offset);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

ssize_t readv(int fd, const struct iovec *iov, int iovcnt) {
    if (!original_readv) {
        original_readv = dlsym(RTLD_NEXT, "readv");
        if (!original_readv) {
            if (debug_mode) fprintf(stderr, "Error loading original readv: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);

    ssize_t result = original_readv(fd, iov, iovcnt);
    if (result > 0) {
        read_bytes += result;
    }
    return result;
}

ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
    if (!original_writev) {
        original_writev = dlsym(RTLD_NEXT, "writev");
        if (!original_writev) {
            if (debug_mode) fprintf(stderr, "Error loading original writev: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_writev(fd, iov, iovcnt);
    if (result > 0) {
        write_bytes += result;
    }
    return result;
}

off_t sendfile(int out_fd, int in_fd, off_t *offset, size_t count) {
    if (!original_sendfile) {
        original_sendfile = dlsym(RTLD_NEXT, "sendfile");
        if (!original_sendfile) {
            if (debug_mode) fprintf(stderr, "Error loading original sendfile: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    off_t result = original_sendfile(out_fd, in_fd, offset, count);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

ssize_t splice(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
    if (!original_splice) {
        original_splice = dlsym(RTLD_NEXT, "splice");
        if (!original_splice) {
            if (debug_mode) fprintf(stderr, "Error loading original splice: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_splice(fd_in, off_in, fd_out, off_out, len, flags);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

ssize_t copy_file_range(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
    if (!original_copy_file_range) {
        original_copy_file_range = dlsym(RTLD_NEXT, "copy_file_range");
        if (!original_copy_file_range) {
            if (debug_mode) fprintf(stderr, "Error loading original copy_file_range: %s\n", dlerror());
            exit(EXIT_FAILURE);
        }
    }

    throttle(&read_bytes, read_limit, &read_start);
    throttle(&write_bytes, write_limit, &write_start);

    ssize_t result = original_copy_file_range(fd_in, off_in, fd_out, off_out, len, flags);
    if (result > 0) {
        read_bytes += result;
        write_bytes += result;
    }
    return result;
}

__attribute__((constructor))
void init() {
    clock_gettime(CLOCK_MONOTONIC, &read_start);
    clock_gettime(CLOCK_MONOTONIC, &write_start);

    // Initialize limits from environment variables (bytes/s)
    char *read_env = getenv("IOTHROTTLE_READ");
    char *write_env = getenv("IOTHROTTLE_WRITE");
    char *debug_env = getenv("IOTHROTTLE_DEBUG");

    if (debug_env && atoi(debug_env) != 0) {
        debug_mode = 1;
    }

    if (debug_mode) fprintf(stderr, "init called: default read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);

    if (read_env) {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ=%s\n", read_env);
        read_limit = strtoull(read_env, NULL, 10);
    } else {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ not set\n");
    }

    if (write_env) {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE=%s\n", write_env);
        write_limit = strtoull(write_env, NULL, 10);
    } else {
        if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE not set\n");
    }

    if (debug_mode) fprintf(stderr, "Final read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);
}