xfy/cocoon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 实现带宽限速 / 流量整形（Token Bucket 算法）
Some checks failed
CI / build (push) Failing after 1m0s
Details

Browse Source 
- 新增 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 问题）
This commit is contained in:
xfy911 2026-06-16 15:41:24 +08:00
parent c970c7f73a
commit 522685e55b
10 changed files with 361 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
Makefile
View File

@ -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)

3
cocoon.h
View File

@ -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 === */

45
config.c
View File

@ -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) {

1
config.h
View File

@ -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,

38
platform.c
View File

@ -11,6 +11,7 @@
*/
#include "platform.h"
#include "throttle.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
@ -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) {
/* 检查限速 */
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) {

15
proxy.c
View File

@ -9,6 +9,7 @@
#include "proxy.h"
#include "proxy_tls.h"
#include "log.h"
#include "throttle.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
@ -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;

36
server.c
View File

@ -36,6 +36,7 @@
#include "cache.h"
#include "dashboard.h"
#include "acme.h"
#include "throttle.h"
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
@ -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);
}

40
static.c
View File

@ -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 <stdio.h>
@ -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;

86
throttle.c Normal file
View File

@ -0,0 +1,86 @@
/**
* throttle.c - Token Bucket 带宽限速实现
*
* 支持 per-connection 和全局总限速，通过 fd 注册机制自动应用。
*
* @author xfy
*/
#include "throttle.h"
#include "platform.h"
#include <pthread.h>
#include <string.h>
/* 最大支持的 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;
}

77
throttle.h Normal file
View File

@ -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 <stdint.h>
#include <stddef.h>
#include <time.h>
/**
* 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 */