Merge pull request 'Bandwidth throttling: --bwlimit <KB/s> with token bucket' (#15) from bwlimit into main

Reviewed-on: #15

src/client/client_cli.c

@@ -1,7 +1,10 @@
 #include "client_send.h"
 #include "config.h"
 #include "log.h"
+#include "protocol.h"
 #include "utils.h"
+#include <errno.h>
+#include <limits.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -42,6 +45,7 @@ static void print_usage(void) {
   printf("  --save-to-disk      Write received files to disk\n");
   printf("  --server-host <ip>  Server IP address (default: 127.0.0.1)\n");
   printf("  --server-port <n>   Server port (default: 8080)\n");
+  printf("  --bwlimit <KB/s>    Bandwidth limit in kilobytes per second\n");
   printf("  --help              Show this help\n");
 }
 
@@ -127,6 +131,20 @@ int main(int argc, char *argv[]) {
       server_host = str_dup(argv[++i]);
     } else if (strcmp(argv[i], "--server-port") == 0 && i + 1 < argc) {
       server_port = atoi(argv[++i]);
+    } else if (strcmp(argv[i], "--bwlimit") == 0 && i + 1 < argc) {
+      char *end;
+      errno = 0;
+      unsigned long long kbps = strtoull(argv[++i], &end, 10);
+      if (errno != 0 || *end != '\0' || kbps == 0) {
+        fprintf(stderr, "Error: --bwlimit must be a positive integer\n");
+        return 1;
+      }
+      if (kbps > ULLONG_MAX / 1024) {
+        fprintf(stderr, "Error: --bwlimit value too large\n");
+        return 1;
+      }
+      io_set_bwlimit(kbps * 1024);
+      log_message(LOG_LEVEL_INFO, "Set bandwidth limit to %llu KB/s", kbps);
     } else if (strcmp(argv[i], "--progress") == 0) {
       config->show_progress = true;
     } else if (strcmp(argv[i], "--chunk-size") == 0 && i + 1 < argc) {

src/shared/protocol.c

@@ -1,18 +1,59 @@
 #include "protocol.h"
 #include "log.h"
+#include <errno.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <time.h>
 #include <unistd.h>
 
 static __thread int io_read_fd = -1;
 static __thread int io_write_fd = -1;
 
+static unsigned long long io_bwlimit = 0;
+static long long bw_tokens = 0;
+static struct timespec bw_last_refill = {0, 0};
+
 void io_set_fds(int read_fd, int write_fd) {
   io_read_fd = read_fd;
   io_write_fd = write_fd;
 }
 
+void io_set_bwlimit(unsigned long long bytes_per_sec) {
+  io_bwlimit = bytes_per_sec;
+  bw_tokens = (long long)io_bwlimit;
+  clock_gettime(CLOCK_MONOTONIC, &bw_last_refill);
+}
+
+static void bw_throttle(size_t bytes_written) {
+  if (io_bwlimit == 0) return;
+
+  struct timespec now;
+  clock_gettime(CLOCK_MONOTONIC, &now);
+
+  long long elapsed_ns = (now.tv_sec - bw_last_refill.tv_sec) * 1000000000LL +
+                         (now.tv_nsec - bw_last_refill.tv_nsec);
+  bw_last_refill = now;
+
+  long long tokens_to_add = (long long)((double)io_bwlimit * elapsed_ns / 1000000000.0);
+  bw_tokens += tokens_to_add;
+  if (bw_tokens > (long long)io_bwlimit)
+    bw_tokens = (long long)io_bwlimit;
+
+  bw_tokens -= (long long)bytes_written;
+
+  if (bw_tokens < 0) {
+    long long deficit_ns = (long long)((double)(-bw_tokens) / io_bwlimit * 1000000000.0);
+    struct timespec sleep_time, remaining;
+    sleep_time.tv_sec = deficit_ns / 1000000000LL;
+    sleep_time.tv_nsec = deficit_ns % 1000000000LL;
+    while (nanosleep(&sleep_time, &remaining) < 0 && errno == EINTR)
+      sleep_time = remaining;
+    bw_tokens = 0;
+    clock_gettime(CLOCK_MONOTONIC, &bw_last_refill);
+  }
+}
+
 static int io_fd(int dir_fd, int file_descriptor) {
   return (dir_fd != -1) ? dir_fd : file_descriptor;
 }
@@ -22,12 +63,16 @@ bool send_n_data(int file_descriptor, void *data, size_t data_size) {
   int fd = io_fd(io_write_fd, file_descriptor);
   ssize_t total_bytes_send = 0;
   while (total_bytes_send < data_size) {
+    size_t chunk = data_size - total_bytes_send;
+    if (io_bwlimit > 0 && chunk > 65536)
+      chunk = 65536;
     ssize_t bytes_send =
-        write(fd, (char *)data + total_bytes_send, data_size - total_bytes_send);
+        write(fd, (char *)data + total_bytes_send, chunk);
     if (bytes_send <= 0) {
       log_message(LOG_LEVEL_ERROR, "Could not send data");
       return false;
     }
+    bw_throttle((size_t)bytes_send);
     total_bytes_send += bytes_send;
   }
   log_message(LOG_LEVEL_DEBUG, "    Send n Data: %zu", total_bytes_send);

src/shared/protocol.h

@@ -9,6 +9,7 @@ typedef int Status;
 enum NET_STATUS { STATUS_OK, STATUS_ERROR, STATUS_FINISHED, STATUS_NEXT, STATUS_CHUNK, STATUS_MANIFEST };
 
 void io_set_fds(int read_fd, int write_fd);
+void io_set_bwlimit(unsigned long long bytes_per_sec);
 bool send_n_data(int file_descriptor, void *data, size_t data_size);
 bool receive_n_data(int file_descriptor, void *data, size_t data_size);
 

test.py

@@ -383,6 +383,17 @@ def run_profile(profile_name, source_dir, dest_dir):
         except Exception as e:
             results.append({"name": "Progress (--progress)", "suite": profile_name, "status": "Error", "time": "N/A", "error": str(e)})
 
+        # Bandwidth limit (--bwlimit 10240 = 10 MB/s)
+        feature_flags = BASE_CLIENT_FLAGS + ["--bwlimit", "10240"]
+        cmd = client_prefix + BASE_CLIENT_CMD + ["--source-dir", source_dir, "--dest-dir", dest_dir] + feature_flags
+        print(f"\n  --- Bandwidth limit (--bwlimit 10240 KB/s) ---\n    Running: {' '.join(cmd)}")
+        try:
+            r = run_single_test(cmd, "Bandwidth limit (--bwlimit 10240)", source_dir, dest_dir)
+            r["suite"] = profile_name
+            results.append(r)
+        except Exception as e:
+            results.append({"name": "Bandwidth limit (--bwlimit 10240)", "suite": profile_name, "status": "Error", "time": "N/A", "error": str(e)})
+
         # Chunk size (--chunk-size 5242880)
         feature_flags = BASE_CLIENT_FLAGS + ["--chunk-size", "5242880"]
         cmd = client_prefix + BASE_CLIENT_CMD + ["--source-dir", source_dir, "--dest-dir", dest_dir] + feature_flags