commit a696a695888e04eaf3a84328509968990e79bed8
| author | 斟酌 鵬兄 <tgckpg@gmail.com> |
|---|
| date | 2026-05-26T21:19:30Z |
|---|
| subject | Fixed OOM issue |
|---|
commit a696a695888e04eaf3a84328509968990e79bed8
Author: 斟酌 鵬兄 <tgckpg@gmail.com>
Date: 2026-05-26T21:19:30Z
Fixed OOM issue
---
tcp_route.c | 47 +++++++--
tests/run_tests.py | 2 +
tests/test_tcp_backpressure.py | 231 +++++++++++++++++++++++++++++++++++++++++
3 files changed, 272 insertions(+), 8 deletions(-)
diff --git a/tcp_route.c b/tcp_route.c
index 62d8d5d..a6f9088 100644
--- a/tcp_route.c
+++ b/tcp_route.c
@@ -16,6 +16,9 @@
#include "route.h"
#include "tcp_route.h"
+#define BEV_READ_HIGH_WATER (256 * 1024)
+#define BEV_WRITE_RESUME_WATER (128 * 1024)
+
struct tcp_route_ctx {
struct event_base *accept_base;
struct worker *worker;
@@ -57,9 +60,9 @@ static void free_conn(conn_t *conn) {
free(conn);
}
-static void pipe_read_cb(struct bufferevent *src, void *arg) {
+static void pipe_read_cb(struct bufferevent *src, void *arg)
+{
conn_t *conn = arg;
-
struct bufferevent *dst;
if (src == conn->client) {
@@ -74,6 +77,30 @@ static void pipe_read_cb(struct bufferevent *src, void *arg) {
struct evbuffer *output = bufferevent_get_output(dst);
evbuffer_add_buffer(output, input);
+
+ if (evbuffer_get_length(output) >= BEV_READ_HIGH_WATER) {
+ bufferevent_disable(src, EV_READ);
+ }
+}
+
+static void pipe_write_cb(struct bufferevent *dst, void *arg)
+{
+ conn_t *conn = arg;
+ struct bufferevent *src;
+
+ if (dst == conn->client) {
+ src = conn->upstream;
+ } else if (dst == conn->upstream) {
+ src = conn->client;
+ } else {
+ return;
+ }
+
+ struct evbuffer *output = bufferevent_get_output(dst);
+
+ if (evbuffer_get_length(output) < BEV_WRITE_RESUME_WATER) {
+ bufferevent_enable(src, EV_READ);
+ }
}
static void event_cb(struct bufferevent *bev, short events, void *arg) {
@@ -137,10 +164,13 @@ static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac)
return;
}
+ bufferevent_setwatermark(conn->client, EV_READ, 0, BEV_READ_HIGH_WATER);
+ bufferevent_setwatermark(conn->upstream, EV_READ, 0, BEV_READ_HIGH_WATER);
+
bufferevent_setcb(
conn->client,
pipe_read_cb,
- NULL,
+ pipe_write_cb,
event_cb,
conn
);
@@ -148,7 +178,7 @@ static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac)
bufferevent_setcb(
conn->upstream,
pipe_read_cb,
- NULL,
+ pipe_write_cb,
event_cb,
conn
);
@@ -259,14 +289,15 @@ static void accept_cb(
dispatch_client_fd(ctx->worker, &ac);
}
-static void accept_error_cb(struct evconnlistener *listener, void *arg) {
- struct event_base *base = arg;
+static void accept_error_cb(struct evconnlistener *listener, void *arg)
+{
+ struct tcp_route_ctx *ctx = arg;
int err = EVUTIL_SOCKET_ERROR();
LOG_ERROR("accept error", "err", _LOGV(evutil_socket_error_to_string(err)));
- evconnlistener_free(listener);
- event_base_loopexit(base, NULL);
+ evconnlistener_disable(listener);
+ event_base_loopexit(ctx->accept_base, NULL);
}
int start_tcp_route(
diff --git a/tests/run_tests.py b/tests/run_tests.py
index 4837fc0..11a738d 100644
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -13,6 +13,7 @@ from .support import (
from . import test_tcp_basic
from . import test_tcp_idle_timeout
from . import test_tcp_stress
+from . import test_tcp_backpressure
from . import test_udp_basic
from . import test_udp_idle_timeout
from . import test_haproxy_proxy_v2
@@ -31,6 +32,7 @@ UDP_PROXY_TEST_MODULES = [
STANDALONE_TEST_MODULES = [
test_haproxy_proxy_v2,
test_udp_proxy_v2,
+ test_tcp_backpressure,
test_tcp_idle_timeout,
test_udp_idle_timeout,
]
diff --git a/tests/test_tcp_backpressure.py b/tests/test_tcp_backpressure.py
new file mode 100644
index 0000000..5089172
--- /dev/null
+++ b/tests/test_tcp_backpressure.py
@@ -0,0 +1,231 @@
+import asyncio
+import os
+
+from .support import (
+ LISTEN_HOST,
+ BACKEND_PORT,
+ PROXY_PORT,
+ SkipTest,
+ run_tinyproxy_with_conf,
+)
+
+BACKPRESSURE_MAX_BYTES = 32 * 1024 * 1024
+BACKPRESSURE_CHUNK_SIZE = 64 * 1024
+BACKPRESSURE_DRAIN_TIMEOUT = 0.5
+
+
+async def close_writer(writer: asyncio.StreamWriter) -> None:
+ try:
+ writer.close()
+ await writer.wait_closed()
+ except ConnectionResetError:
+ pass
+
+async def abort_writer(writer: asyncio.StreamWriter) -> None:
+ transport = writer.transport
+ transport.abort()
+
+ # Let the event loop process connection_lost().
+ await asyncio.sleep(0)
+
+async def blackhole_handler(
+ reader: asyncio.StreamReader,
+ writer: asyncio.StreamWriter,
+) -> None:
+ try:
+ # Intentionally do not read.
+ #
+ # This simulates an upstream that accepted the TCP connection but
+ # stopped draining. Without proxy-side backpressure, tinyproxy may keep
+ # reading from the client and buffer data until memory grows without
+ # bound.
+ await asyncio.sleep(3600)
+ finally:
+ await close_writer(writer)
+
+async def push_until_drain_blocks(
+ writer: asyncio.StreamWriter,
+ max_bytes: int = BACKPRESSURE_MAX_BYTES,
+ chunk_size: int = BACKPRESSURE_CHUNK_SIZE,
+ drain_timeout: float = BACKPRESSURE_DRAIN_TIMEOUT,
+) -> int:
+ chunk = b"x" * chunk_size
+ written = 0
+
+ while written < max_bytes:
+ writer.write(chunk)
+ written += len(chunk)
+
+ try:
+ await asyncio.wait_for(writer.drain(), timeout=drain_timeout)
+ except asyncio.TimeoutError:
+ return written
+ except ConnectionError:
+ return written
+ except BrokenPipeError:
+ return written
+
+ await asyncio.sleep(0)
+
+ return written
+
+async def test_tcp_backpressure_when_upstream_does_not_read() -> None:
+ proxy_bin = os.environ.get("TINYPROXY_BIN")
+ if not proxy_bin:
+ raise SkipTest("TINYPROXY_BIN is not set")
+
+ conf_text = (
+ f"{LISTEN_HOST}:{PROXY_PORT} "
+ f"{LISTEN_HOST}:{BACKEND_PORT} "
+ f"tcp idle_timeout=30\n"
+ )
+
+ stop_backend = asyncio.Event()
+
+ async def blackhole_handler(
+ reader: asyncio.StreamReader,
+ writer: asyncio.StreamWriter,
+ ) -> None:
+ try:
+ # Intentionally do not read from reader.
+ await stop_backend.wait()
+ finally:
+ await abort_writer(writer)
+
+ backend_server = await asyncio.start_server(
+ blackhole_handler,
+ LISTEN_HOST,
+ BACKEND_PORT,
+ backlog=128,
+ )
+
+ writer: asyncio.StreamWriter | None = None
+
+ try:
+ async with run_tinyproxy_with_conf(
+ proxy_bin=proxy_bin,
+ conf_text=conf_text,
+ listen_host=LISTEN_HOST,
+ listen_port=PROXY_PORT,
+ proto="tcp",
+ ):
+ reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
+
+ written = await asyncio.wait_for(
+ push_until_drain_blocks(writer),
+ timeout=10.0,
+ )
+
+ assert written < BACKPRESSURE_MAX_BYTES, (
+ "proxy accepted too much client data while upstream was not reading; "
+ "this likely means TCP backpressure/watermarks are missing. "
+ f"wrote={written} max={BACKPRESSURE_MAX_BYTES}"
+ )
+
+ _ = reader
+ finally:
+ if writer is not None:
+ await abort_writer(writer)
+
+ stop_backend.set()
+
+ backend_server.close()
+ await asyncio.wait_for(backend_server.wait_closed(), timeout=3.0)
+
+async def test_tcp_backpressure_when_client_does_not_read() -> None:
+ proxy_bin = os.environ.get("TINYPROXY_BIN")
+ if not proxy_bin:
+ raise SkipTest("TINYPROXY_BIN is not set")
+
+ conf_text = (
+ f"{LISTEN_HOST}:{PROXY_PORT} "
+ f"{LISTEN_HOST}:{BACKEND_PORT} "
+ f"tcp idle_timeout=30\n"
+ )
+
+ backend_result: asyncio.Future[tuple[str, int]] = (
+ asyncio.get_running_loop().create_future()
+ )
+
+ async def pushing_backend_handler(
+ reader: asyncio.StreamReader,
+ writer: asyncio.StreamWriter,
+ ) -> None:
+ try:
+ # Intentionally do not wait for a client trigger.
+ #
+ # This test is about upstream -> proxy -> non-reading client.
+ # Waiting for a trigger byte makes the test depend on the
+ # client->upstream direction, which is unrelated and flaky here.
+ _ = reader
+
+ written = await push_until_drain_blocks(writer)
+
+ if not backend_result.done():
+ if written < BACKPRESSURE_MAX_BYTES:
+ backend_result.set_result(("blocked", written))
+ else:
+ backend_result.set_result(("unbounded", written))
+ except ConnectionError:
+ if not backend_result.done():
+ backend_result.set_result(("blocked", 0))
+ except BrokenPipeError:
+ if not backend_result.done():
+ backend_result.set_result(("blocked", 0))
+ except Exception as e:
+ if not backend_result.done():
+ backend_result.set_exception(e)
+ finally:
+ await abort_writer(writer)
+
+ backend_server = await asyncio.start_server(
+ pushing_backend_handler,
+ LISTEN_HOST,
+ BACKEND_PORT,
+ backlog=128,
+ )
+
+ writer: asyncio.StreamWriter | None = None
+
+ try:
+ async with run_tinyproxy_with_conf(
+ proxy_bin=proxy_bin,
+ conf_text=conf_text,
+ listen_host=LISTEN_HOST,
+ listen_port=PROXY_PORT,
+ proto="tcp",
+ ):
+ reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
+
+ # Intentionally do not read from reader.
+ _ = reader
+
+ status, written = await asyncio.wait_for(
+ backend_result,
+ timeout=10.0,
+ )
+
+ assert status == "blocked", (
+ "proxy accepted too much upstream data while client was not reading; "
+ "this likely means reverse-direction TCP backpressure/watermarks "
+ "are missing. "
+ f"status={status} wrote={written} max={BACKPRESSURE_MAX_BYTES}"
+ )
+ finally:
+ if writer is not None:
+ await abort_writer(writer)
+
+ backend_server.close()
+ await asyncio.wait_for(backend_server.wait_closed(), timeout=3.0)
+
+
+TESTS = [
+ (
+ "test_tcp_backpressure_when_upstream_does_not_read",
+ test_tcp_backpressure_when_upstream_does_not_read,
+ ),
+ (
+ "test_tcp_backpressure_when_client_does_not_read",
+ test_tcp_backpressure_when_client_does_not_read,
+ ),
+]