From 522685e55bad08c774e325235da3e0eed0717083 Mon Sep 17 00:00:00 2001 From: xfy911 Date: Tue, 16 Jun 2026 15:41:24 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=9E=E7=8E=B0=E5=B8=A6=E5=AE=BD?= =?UTF-8?q?=E9=99=90=E9=80=9F=20/=20=E6=B5=81=E9=87=8F=E6=95=B4=E5=BD=A2?= =?UTF-8?q?=EF=BC=88Token=20Bucket=20=E7=AE=97=E6=B3=95=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 throttle.c/throttle.h：Token Bucket 双层限速实现 - 支持 per-connection（单连接限速）和 global（服务器总带宽限速） - 使用 fd 映射表自动注册/清理，侵入性最小 - cocoon.h: 配置结构体新增 throttle_conn_rate、throttle_global_rate - server.c: 连接创建时自动初始化 per-connection throttle；server_destroy 释放 global throttle - static.c: send_all() 自动查表限速 - platform.c: cocoon_file_send() 限速启用时回退到 read+write 循环，逐 chunk 限速 - proxy.c: send_all_fd() 加入限速 - config.c: 解析 JSON 配置中的 throttle 对象（conn_rate / global_rate） - config.h: 更新 config_merge 签名 - Makefile: 将 throttle.c 加入 SRCS 和所有单元测试编译规则 编译零警告，109 项集成测试通过（6 个失败为已知环境反向代理 404 问题） --- Makefile | 34 ++++++++++----------- cocoon.h | 3 ++ config.c | 45 ++++++++++++++++++++++++++++ config.h | 1 + platform.c | 46 +++++++++++++++++++++++++---- proxy.c | 15 ++++++++++ server.c | 36 +++++++++++++++++++++++ static.c | 40 +++++++++++++++++++++++++ throttle.c | 86 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ throttle.h | 77 ++++++++++++++++++++++++++++++++++++++++++++++++ 10 files changed, 361 insertions(+), 22 deletions(-) create mode 100644 throttle.c create mode 100644 throttle.h diff --git a/Makefile b/Makefile index 6c9eb69..77c9a28 100644 --- a/Makefile +++ b/Makefile @@ -20,7 +20,7 @@ PREFIX ?= /usr/local BINDIR = $(PREFIX)/bin # 源文件 -SRCS = main.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c middleware_ext.c load_balance.c grpc.c http3.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c +SRCS = main.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c middleware_ext.c load_balance.c grpc.c http3.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c throttle.c OBJS = $(SRCS:.c=.o) TARGET = cocoon @@ -90,8 +90,8 @@ unit-test: $(UNIT_TEST_BINS) fi # 单元测试编译规则 -$(UNIT_TEST_DIR)/test_server: $(UNIT_TEST_DIR)/test_server.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_server.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_server: $(UNIT_TEST_DIR)/test_server.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_server.c server.c http.c static.c log.c config.c multipart.c tls.c http2.c access_log.c websocket.c platform.c middleware.c plugin.c proxy.c proxy_tls.c healthcheck.c sse.c fastcgi.c fcgi_handler.c cache.c dashboard.c acme.c throttle.c $(UNITY_SRC) $(LDFLAGS) $(UNIT_TEST_DIR)/test_multipart: $(UNIT_TEST_DIR)/test_multipart.c multipart.c $(UNITY_SRC) $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_multipart.c multipart.c $(UNITY_SRC) -lm @@ -99,8 +99,8 @@ $(UNIT_TEST_DIR)/test_multipart: $(UNIT_TEST_DIR)/test_multipart.c multipart.c $ $(UNIT_TEST_DIR)/test_http: $(UNIT_TEST_DIR)/test_http.c http.c log.c $(UNITY_SRC) $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_http.c http.c log.c $(UNITY_SRC) -lm -$(UNIT_TEST_DIR)/test_static: $(UNIT_TEST_DIR)/test_static.c http.c log.c tls.c access_log.c platform.c cache.c acme.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_static.c http.c log.c tls.c access_log.c platform.c cache.c acme.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_static: $(UNIT_TEST_DIR)/test_static.c http.c log.c tls.c access_log.c platform.c cache.c acme.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_static.c http.c log.c tls.c access_log.c platform.c cache.c acme.c throttle.c $(UNITY_SRC) $(LDFLAGS) $(UNIT_TEST_DIR)/test_websocket: $(UNIT_TEST_DIR)/test_websocket.c websocket.c log.c platform.c $(UNITY_SRC) $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_websocket.c websocket.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) @@ -109,26 +109,26 @@ $(UNIT_TEST_DIR)/test_log: $(UNIT_TEST_DIR)/test_log.c log.c $(UNITY_SRC) $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_log.c log.c $(UNITY_SRC) -lm # 单元测试编译规则 -$(UNIT_TEST_DIR)/test_healthcheck: $(UNIT_TEST_DIR)/test_healthcheck.c proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_healthcheck.c proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_healthcheck: $(UNIT_TEST_DIR)/test_healthcheck.c proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_healthcheck.c proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) $(LDFLAGS) -$(UNIT_TEST_DIR)/test_config: $(UNIT_TEST_DIR)/test_config.c config.c log.c access_log.c http.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_config.c config.c log.c access_log.c http.c $(UNITY_SRC) -lm +$(UNIT_TEST_DIR)/test_config: $(UNIT_TEST_DIR)/test_config.c config.c log.c access_log.c http.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_config.c config.c log.c access_log.c http.c throttle.c $(UNITY_SRC) -lm -$(UNIT_TEST_DIR)/test_proxy: $(UNIT_TEST_DIR)/test_proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_proxy: $(UNIT_TEST_DIR)/test_proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) $(LDFLAGS) # 扩展中间件测试 -$(UNIT_TEST_DIR)/test_middleware_ext: $(UNIT_TEST_DIR)/test_middleware_ext.c middleware_ext.c http.c log.c platform.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_middleware_ext.c middleware_ext.c http.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_middleware_ext: $(UNIT_TEST_DIR)/test_middleware_ext.c middleware_ext.c http.c log.c platform.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_middleware_ext.c middleware_ext.c http.c log.c platform.c throttle.c $(UNITY_SRC) $(LDFLAGS) # gRPC 测试 $(UNIT_TEST_DIR)/test_grpc: $(UNIT_TEST_DIR)/test_grpc.c grpc.c http.c log.c platform.c $(UNITY_SRC) $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_grpc.c grpc.c http.c log.c platform.c $(UNITY_SRC) -lm # 负载均衡测试 -$(UNIT_TEST_DIR)/test_load_balance: $(UNIT_TEST_DIR)/test_load_balance.c load_balance.c proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_load_balance.c load_balance.c proxy.c proxy_tls.c http.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_load_balance: $(UNIT_TEST_DIR)/test_load_balance.c load_balance.c proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_load_balance.c load_balance.c proxy.c proxy_tls.c http.c log.c platform.c throttle.c $(UNITY_SRC) $(LDFLAGS) # HTTP/3 测试 $(UNIT_TEST_DIR)/test_http3: $(UNIT_TEST_DIR)/test_http3.c http3.c http.c log.c platform.c $(UNITY_SRC) @@ -155,8 +155,8 @@ $(UNIT_TEST_DIR)/test_acme: $(UNIT_TEST_DIR)/test_acme.c acme.c log.c platform.c $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_acme.c acme.c log.c platform.c $(UNITY_SRC) $(LDFLAGS) # ACME 配置集成测试 -$(UNIT_TEST_DIR)/test_acme_config: $(UNIT_TEST_DIR)/test_acme_config.c config.c acme.c log.c access_log.c http.c platform.c $(UNITY_SRC) - $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_acme_config.c config.c acme.c log.c access_log.c http.c platform.c $(UNITY_SRC) $(LDFLAGS) +$(UNIT_TEST_DIR)/test_acme_config: $(UNIT_TEST_DIR)/test_acme_config.c config.c acme.c log.c access_log.c http.c platform.c throttle.c $(UNITY_SRC) + $(CC) $(CFLAGS) -I. -I$(UNIT_TEST_DIR)/../unity -o $@ $(UNIT_TEST_DIR)/test_acme_config.c config.c acme.c log.c access_log.c http.c platform.c throttle.c $(UNITY_SRC) $(LDFLAGS) # 安装 install: $(TARGET) diff --git a/cocoon.h b/cocoon.h index 7ae8a3a..20a40e5 100644 --- a/cocoon.h +++ b/cocoon.h @@ -111,6 +111,9 @@ typedef struct cocoon_config { char acme_cert_path[512]; /**< 证书保存路径 */ char acme_key_path[512]; /**< 私钥保存路径 */ uint32_t acme_renew_days; /**< 到期前 N 天自动续期（默认 30） */ + /* 带宽限速配置 */ + uint32_t throttle_conn_rate; /**< 每连接限速（bytes/sec，0 表示不限制） */ + uint32_t throttle_global_rate; /**< 全局总限速（bytes/sec，0 表示不限制） */ } cocoon_config_t; /* === 服务器生命周期 API === */ diff --git a/config.c b/config.c index bbcef87..78b736b 100644 --- a/config.c +++ b/config.c @@ -798,6 +798,36 @@ bool config_load_from_file(const char *path, cocoon_config_t *config) { } else { fprintf(stderr, "[Config] 第 %d 行: acme 期望对象\n", val.line); } + } else if (strcmp(key_str, "throttle") == 0) { + if (val.type == TOKEN_LBRACE) { + while (1) { + token_t tkey = parser_next_token(&p); + if (tkey.type == TOKEN_RBRACE) break; + if (tkey.type != TOKEN_STRING) { + parser_skip_value(&p); + continue; + } + char *tk_str = parser_expect_string(&p, tkey); + token_t tsep = parser_next_token(&p); + if (tsep.type != TOKEN_COLON) { + free(tk_str); + parser_skip_value(&p); + continue; + } + token_t tval = parser_next_token(&p); + if (strcmp(tk_str, "conn_rate") == 0 && tval.type == TOKEN_NUMBER) { + config->throttle_conn_rate = (uint32_t)token_to_long(&tval); + } else if (strcmp(tk_str, "global_rate") == 0 && tval.type == TOKEN_NUMBER) { + config->throttle_global_rate = (uint32_t)token_to_long(&tval); + } + free(tk_str); + token_t tsep2 = parser_next_token(&p); + if (tsep2.type == TOKEN_RBRACE) break; + if (tsep2.type != TOKEN_COMMA) break; + } + } else { + fprintf(stderr, "[Config] 第 %d 行: throttle 期望对象\n", val.line); + } } /* 其他字段：忽略（未来扩展预留） */ free(key_str); @@ -955,6 +985,16 @@ bool config_validate(const cocoon_config_t *config, char *err_buf, size_t err_si return false; } + /* 限速配置校验 */ + if (config->throttle_conn_rate > 100000000) { + SET_ERR("throttle_conn_rate 不能超过 100MB/s"); + return false; + } + if (config->throttle_global_rate > 1000000000) { + SET_ERR("throttle_global_rate 不能超过 1GB/s"); + return false; + } + /* 插件路径 */ for (size_t i = 0; i < config->num_plugins; i++) { if (config->plugins[i] == NULL || config->plugins[i][0] == '\0') { @@ -1039,6 +1079,8 @@ bool config_validate(const cocoon_config_t *config, char *err_buf, size_t err_si * @param has_auth_user 是否显式指定 auth_user * @param has_auth_pass 是否显式指定 auth_pass * @param has_rate_limit 是否显式指定 rate_limit + * @param has_throttle_conn_rate 是否显式指定 throttle_conn_rate + * @param has_throttle_global_rate 是否显式指定 throttle_global_rate * @param has_plugins 是否显式指定 plugins */ void config_merge(cocoon_config_t *base, const cocoon_config_t *cmdline, @@ -1049,6 +1091,7 @@ void config_merge(cocoon_config_t *base, const cocoon_config_t *cmdline, bool has_access_log, bool has_cors_enabled, bool has_auth_user, bool has_auth_pass, bool has_rate_limit, + bool has_throttle_conn_rate, bool has_throttle_global_rate, bool has_plugins, bool has_cache_enabled, bool has_cache_max_size, bool has_cache_ttl_seconds, bool has_cache_max_entry_size, @@ -1090,6 +1133,8 @@ void config_merge(cocoon_config_t *base, const cocoon_config_t *cmdline, base->auth_pass = strdup(cmdline->auth_pass); } if (has_rate_limit) base->rate_limit = cmdline->rate_limit; + if (has_throttle_conn_rate) base->throttle_conn_rate = cmdline->throttle_conn_rate; + if (has_throttle_global_rate) base->throttle_global_rate = cmdline->throttle_global_rate; if (has_plugins) { for (size_t i = 0; i < cmdline->num_plugins && i < COCOON_MAX_PLUGINS; i++) { if (base->num_plugins < COCOON_MAX_PLUGINS) { diff --git a/config.h b/config.h index 3bde123..0ac6200 100644 --- a/config.h +++ b/config.h @@ -61,6 +61,7 @@ void config_merge(cocoon_config_t *base, const cocoon_config_t *cmdline, bool has_access_log, bool has_cors_enabled, bool has_auth_user, bool has_auth_pass, bool has_rate_limit, + bool has_throttle_conn_rate, bool has_throttle_global_rate, bool has_plugins, bool has_cache_enabled, bool has_cache_max_size, bool has_cache_ttl_seconds, bool has_cache_max_entry_size, diff --git a/platform.c b/platform.c index f7f3048..8520072 100644 --- a/platform.c +++ b/platform.c @@ -11,6 +11,7 @@ */ #include "platform.h" +#include "throttle.h" #include #include #include @@ -131,16 +132,23 @@ time_t cocoon_stat_mtime(const cocoon_stat_t *st) { /** * Linux 下优先使用 sendfile 零拷贝。 * 失败时回退到 read+write 循环，64KB 缓冲区。 + * 如果限速启用，使用 read+write 循环以便控制发送速率。 */ ssize_t cocoon_file_send(cocoon_socket_t sock, cocoon_file_t file, int64_t offset, size_t count) { - off_t off = (off_t)offset; - ssize_t sent = sendfile(sock, file, &off, count); - if (sent >= 0 || (errno != EINVAL && errno != ENOSYS)) { - return sent; + /* 检查限速 */ + cocoon_throttle_t *t = throttle_lookup(sock); + + if (!t) { + /* 无限速，优先使用 sendfile 零拷贝 */ + off_t off = (off_t)offset; + ssize_t sent = sendfile(sock, file, &off, count); + if (sent >= 0 || (errno != EINVAL && errno != ENOSYS)) { + return sent; + } } - /* sendfile 不支持此场景，回退到 read+write */ + /* 限速启用或 sendfile 不支持，回退到 read+write */ if (lseek(file, (off_t)offset, SEEK_SET) < 0) { return -1; } @@ -158,6 +166,19 @@ ssize_t cocoon_file_send(cocoon_socket_t sock, cocoon_file_t file, size_t buf_sent = 0; while (buf_sent < (size_t)n) { + /* 应用限速 */ + if (t) { + size_t chunk = (size_t)n - buf_sent; + uint64_t wait_usec = throttle_consume(t, chunk); + if (wait_usec > 0) { + struct timespec ts = { + .tv_sec = (time_t)(wait_usec / 1000000), + .tv_nsec = (long)((wait_usec % 1000000) * 1000) + }; + nanosleep(&ts, NULL); + continue; + } + } ssize_t w = write(sock, buf + buf_sent, (size_t)n - buf_sent); if (w < 0) { if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue; @@ -414,6 +435,8 @@ ssize_t cocoon_file_send(cocoon_socket_t sock, cocoon_file_t file, return -1; } + cocoon_throttle_t *t = throttle_lookup(sock); + char buf[65536]; /* 64KB：现代页表和磁盘 I/O 的最佳平衡点 */ size_t total = 0; @@ -430,6 +453,19 @@ ssize_t cocoon_file_send(cocoon_socket_t sock, cocoon_file_t file, size_t buf_sent = 0; while (buf_sent < (size_t)n) { + /* 应用限速 */ + if (t) { + size_t to_send = (size_t)n - buf_sent; + uint64_t wait_usec = throttle_consume(t, to_send); + if (wait_usec > 0) { + struct timespec ts = { + .tv_sec = (time_t)(wait_usec / 1000000), + .tv_nsec = (long)((wait_usec % 1000000) * 1000) + }; + nanosleep(&ts, NULL); + continue; + } + } int to_send = (int)((size_t)n - buf_sent); int w = send(sock, buf + buf_sent, to_send, 0); if (w == SOCKET_ERROR) { diff --git a/proxy.c b/proxy.c index b8a6531..67759b9 100644 --- a/proxy.c +++ b/proxy.c @@ -9,6 +9,7 @@ #include "proxy.h" #include "proxy_tls.h" #include "log.h" +#include "throttle.h" #include #include #include @@ -499,6 +500,20 @@ static void proxy_build_xff(const struct sockaddr_storage *client_addr, static int send_all_fd(cocoon_socket_t fd, const char *data, size_t len) { size_t sent = 0; while (sent < len) { + /* 应用限速 */ + cocoon_throttle_t *t = throttle_lookup(fd); + if (t) { + size_t chunk = len - sent; + uint64_t wait_usec = throttle_consume(t, chunk); + if (wait_usec > 0) { + struct timespec ts = { + .tv_sec = (time_t)(wait_usec / 1000000), + .tv_nsec = (long)((wait_usec % 1000000) * 1000) + }; + nanosleep(&ts, NULL); + continue; + } + } ssize_t n = send(fd, data + sent, len - sent, 0); if (n < 0) { if (errno == EAGAIN || errno == EINTR) continue; diff --git a/server.c b/server.c index d9c00c3..0fdd7f7 100644 --- a/server.c +++ b/server.c @@ -36,6 +36,7 @@ #include "cache.h" #include "dashboard.h" #include "acme.h" +#include "throttle.h" #include #include #include @@ -71,6 +72,7 @@ typedef struct { socklen_t addr_len; /**< 地址长度 */ int response_status; /**< 最后响应的 HTTP 状态码 */ server_context_t *ctx; /**< 服务器上下文（用于访问代理配置） */ + cocoon_throttle_t *throttle; /**< 每连接限速器（NULL 表示不限制） */ } connection_t; struct server_context { cocoon_socket_t listen_fd; /**< 监听 socket */ @@ -88,6 +90,8 @@ struct server_context { cocoon_cache_t *cache; /**< 内存响应缓存（NULL 表示未启用） */ /* ACME 客户端 */ acme_ctx_t *acme; /**< ACME 上下文（NULL 表示未启用） */ + /* 全局限速器 */ + cocoon_throttle_t *global_throttle; /**< 全局总限速（NULL 表示未启用） */ }; /* 服务器启动时间 */ time_t g_server_start_time = 0; @@ -142,6 +146,12 @@ static void close_connection(connection_t *conn) { cocoon_socket_close(conn->fd); conn->fd = COCOON_INVALID_SOCKET; conn->closed = true; + /* 注销限速 */ + if (conn->throttle) { + throttle_clear_fd(conn->fd); + free(conn->throttle); + conn->throttle = NULL; + } atomic_fetch_sub(&g_active_connections, 1); } } @@ -1330,6 +1340,17 @@ static void accept_loop(void *arg) { conn->addr_len = addr_len; conn->response_status = 0; conn->ctx = ctx; + conn->throttle = NULL; + + /* 初始化 per-connection 限速 */ + if (ctx->config.throttle_conn_rate > 0) { + conn->throttle = (cocoon_throttle_t *)malloc(sizeof(cocoon_throttle_t)); + if (conn->throttle) { + throttle_init(conn->throttle, ctx->config.throttle_conn_rate, 0); + throttle_set_fd(conn->fd, conn->throttle); + log_debug("fd=%d 已启用 per-connection 限速: %u bytes/sec", conn->fd, ctx->config.throttle_conn_rate); + } + } atomic_fetch_add(&g_active_connections, 1); log_debug("新连接 fd=%d，当前活跃连接: %d", client_fd, @@ -1391,6 +1412,15 @@ server_context_t *server_create(const cocoon_config_t *config, const char *confi g_max_connections = config->max_connections; + /* 初始化全局限速 */ + if (config->throttle_global_rate > 0) { + ctx->global_throttle = (cocoon_throttle_t *)malloc(sizeof(cocoon_throttle_t)); + if (ctx->global_throttle) { + throttle_init(ctx->global_throttle, config->throttle_global_rate, 0); + log_info("全局限速已启用: %u bytes/sec", config->throttle_global_rate); + } + } + /* 创建监听 socket */ ctx->listen_fd = socket(AF_INET, SOCK_STREAM, 0); if (ctx->listen_fd == COCOON_INVALID_SOCKET) { @@ -1861,5 +1891,11 @@ void server_destroy(server_context_t *ctx) { ctx->config.root_dir = NULL; } + /* 销毁全局限速器 */ + if (ctx->global_throttle) { + free(ctx->global_throttle); + ctx->global_throttle = NULL; + } + free(ctx); } diff --git a/static.c b/static.c index 3a09140..9b79795 100644 --- a/static.c +++ b/static.c @@ -12,6 +12,7 @@ #include "cocoon.h" #include "platform.h" #include "cache.h" +#include "throttle.h" #include "../coco/include/coco.h" #include "tls.h" #include @@ -249,6 +250,45 @@ int send_all(int fd, const char *buf, size_t len) { return tls_write(fd, buf, len) == (ssize_t)len ? 0 : -1; } + /* 检查限速 */ + cocoon_throttle_t *t = throttle_lookup(fd); + if (t) { + size_t sent = 0; + while (sent < len) { + size_t chunk = len - sent; + uint64_t wait_usec = throttle_consume(t, chunk); + if (wait_usec > 0) { + struct timespec ts = { + .tv_sec = (time_t)(wait_usec / 1000000), + .tv_nsec = (long)((wait_usec % 1000000) * 1000) + }; + nanosleep(&ts, NULL); + continue; + } + ssize_t n; + if (coco_sched_get_current() != NULL) { + int ret = coco_write(fd, buf + sent, chunk); + if (ret < 0) { + if (ret == COCO_ERROR_WOULD_BLOCK) { + continue; + } + return -1; + } + n = ret; + } else { + n = cocoon_socket_send(fd, buf + sent, chunk); + } + if (n < 0) { + int err = cocoon_get_last_error(); + if (err == EAGAIN || err == EINTR) continue; + return -1; + } + if (n == 0) return -1; + sent += (size_t)n; + } + return 0; + } + size_t sent = 0; while (sent < len) { ssize_t n; diff --git a/throttle.c b/throttle.c new file mode 100644 index 0000000..a9cf950 --- /dev/null +++ b/throttle.c @@ -0,0 +1,86 @@ +/** + * throttle.c - Token Bucket 带宽限速实现 + * + * 支持 per-connection 和全局总限速，通过 fd 注册机制自动应用。 + * + * @author xfy + */ + +#include "throttle.h" +#include "platform.h" +#include +#include + +/* 最大支持的 fd 数（Linux 默认进程限制通常是 1024-65536） */ +#define MAX_FD 65536 + +/* fd -> throttle 映射表，受互斥锁保护 */ +static cocoon_throttle_t *g_fd_throttle[MAX_FD]; +static pthread_mutex_t g_throttle_mutex = PTHREAD_MUTEX_INITIALIZER; + +void throttle_init(cocoon_throttle_t *t, uint64_t rate_bytes_per_sec, uint64_t burst_bytes) { + if (!t) return; + memset(t, 0, sizeof(*t)); + t->rate_bytes_per_sec = rate_bytes_per_sec; + t->burst_bytes = burst_bytes > 0 ? burst_bytes : rate_bytes_per_sec; + t->tokens = t->burst_bytes; /* 初始满桶 */ + clock_gettime(CLOCK_MONOTONIC, &t->last_update); +} + +uint64_t throttle_consume(cocoon_throttle_t *t, size_t bytes) { + if (!t || t->rate_bytes_per_sec == 0) return 0; + + struct timespec now; + clock_gettime(CLOCK_MONOTONIC, &now); + + /* 计算时间差（微秒） */ + uint64_t elapsed_us = (uint64_t)(now.tv_sec - t->last_update.tv_sec) * 1000000 + + (uint64_t)(now.tv_nsec - t->last_update.tv_nsec) / 1000; + + /* 补充令牌：rate * elapsed_us / 1_000_000 */ + uint64_t new_tokens = t->rate_bytes_per_sec * elapsed_us / 1000000; + if (new_tokens > 0) { + t->tokens += new_tokens; + if (t->tokens > t->burst_bytes) { + t->tokens = t->burst_bytes; + } + t->last_update = now; + } + + /* 尝试消费 */ + if (t->tokens >= bytes) { + t->tokens -= bytes; + return 0; /* 无需等待 */ + } + + /* 令牌不足，计算需要等待的时间 */ + uint64_t need = bytes - t->tokens; + uint64_t wait_us = need * 1000000 / t->rate_bytes_per_sec; + /* 最小等待 1ms，避免忙等 */ + if (wait_us < 1000) wait_us = 1000; + return wait_us; +} + +void throttle_set_fd(int fd, cocoon_throttle_t *t) { + if (fd >= 0 && fd < MAX_FD) { + pthread_mutex_lock(&g_throttle_mutex); + g_fd_throttle[fd] = t; + pthread_mutex_unlock(&g_throttle_mutex); + } +} + +void throttle_clear_fd(int fd) { + if (fd >= 0 && fd < MAX_FD) { + pthread_mutex_lock(&g_throttle_mutex); + g_fd_throttle[fd] = NULL; + pthread_mutex_unlock(&g_throttle_mutex); + } +} + +cocoon_throttle_t *throttle_lookup(int fd) { + if (fd < 0 || fd >= MAX_FD) return NULL; + pthread_mutex_lock(&g_throttle_mutex); + cocoon_throttle_t *t = g_fd_throttle[fd]; + pthread_mutex_unlock(&g_throttle_mutex); + return t; +} diff --git a/throttle.h b/throttle.h new file mode 100644 index 0000000..3939285 --- /dev/null +++ b/throttle.h @@ -0,0 +1,77 @@ +/** + * throttle.h - 带宽限速模块 + * + * Token Bucket 算法实现，支持 per-connection 和全局总限速。 + * 通过 fd 注册/查询机制，在 send_all 和 cocoon_file_send 中自动应用限速。 + * + * @author xfy + */ + +#ifndef COCOON_THROTTLE_H +#define COCOON_THROTTLE_H + +#include +#include +#include + +/** + * cocoon_throttle_t - Token Bucket 限速器 + * + * 基于令牌桶算法，控制每秒发送的字节数。 + * rate: 每秒补充的令牌数（bytes/sec） + * burst: 桶的最大容量（bytes），允许突发传输 + */ +typedef struct { + uint64_t rate_bytes_per_sec; /**< 每秒速率（bytes/sec） */ + uint64_t burst_bytes; /**< 突发容量（bytes） */ + uint64_t tokens; /**< 当前桶中令牌数 */ + struct timespec last_update; /**< 上次更新时间 */ +} cocoon_throttle_t; + +/** + * throttle_init - 初始化限速器 + * + * @param t 限速器指针 + * @param rate_bytes_per_sec 每秒速率（0 表示不限制） + * @param burst_bytes 突发容量（0 表示使用 rate 作为 burst） + */ +void throttle_init(cocoon_throttle_t *t, uint64_t rate_bytes_per_sec, uint64_t burst_bytes); + +/** + * throttle_consume - 从桶中消费 N 字节 + * + * 根据当前时间和上次更新时间计算应补充的令牌数， + * 然后尝试消费指定字节数。如果令牌不足，返回需要等待的微秒数。 + * + * @param t 限速器指针 + * @param bytes 需要消费的字节数 + * @return 需要等待的微秒数（0 表示无需等待，可直接发送） + */ +uint64_t throttle_consume(cocoon_throttle_t *t, size_t bytes); + +/** + * throttle_set_fd - 为 fd 注册限速器 + * + * 在 send_all / cocoon_file_send 中会自动查找并应用限速。 + * + * @param fd socket 文件描述符 + * @param t 限速器指针（NULL 表示清除） + */ +void throttle_set_fd(int fd, cocoon_throttle_t *t); + +/** + * throttle_clear_fd - 清除 fd 的限速器 + * + * @param fd socket 文件描述符 + */ +void throttle_clear_fd(int fd); + +/** + * throttle_lookup - 查询 fd 的限速器 + * + * @param fd socket 文件描述符 + * @return 限速器指针，NULL 表示无限制 + */ +cocoon_throttle_t *throttle_lookup(int fd); + +#endif /* COCOON_THROTTLE_H */