Skip to content

Commit ec4c463

Browse files
committed
feat(gateway): add per-app connection rate limiting
Add configurable maximum concurrent connections per app to prevent a single slow or overloaded app from exhausting system resources and affecting other apps. - Add `max_connections_per_app` config option (default: 2000, 0 = unlimited) - Check connection count before establishing new backend connections - Log warning when rate limit is exceeded
1 parent f49f228 commit ec4c463

4 files changed

Lines changed: 39 additions & 4 deletions

File tree

gateway/gateway.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ app_address_ns_prefix = "_dstack-app-address"
7070
app_address_ns_compat = true
7171
workers = 32
7272
external_port = 443
73+
# Maximum concurrent connections per app. 0 means unlimited.
74+
max_connections_per_app = 2000
7375

7476
[core.proxy.timeouts]
7577
# Timeout for establishing a connection to the target app.

gateway/src/config.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ pub struct ProxyConfig {
8585
pub workers: usize,
8686
pub app_address_ns_prefix: String,
8787
pub app_address_ns_compat: bool,
88+
/// Maximum concurrent connections per app. 0 means unlimited.
89+
pub max_connections_per_app: u64,
8890
}
8991

9092
#[derive(Debug, Clone, Deserialize)]

gateway/src/proxy/tls_passthough.rs

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22
//
33
// SPDX-License-Identifier: Apache-2.0
44

5-
use anyhow::{Context, Result};
65
use std::fmt::Debug;
6+
use std::sync::atomic::Ordering;
7+
8+
use anyhow::{bail, Context, Result};
79
use tokio::{io::AsyncWriteExt, net::TcpStream, task::JoinSet, time::timeout};
8-
use tracing::{debug, info};
10+
use tracing::{debug, info, warn};
911

1012
use crate::{
1113
main_service::Proxy,
@@ -88,11 +90,38 @@ pub(crate) async fn proxy_with_sni(
8890
proxy_to_app(state, inbound, buffer, &addr.app_id, addr.port).await
8991
}
9092

93+
/// Check if app has reached max connections limit
94+
fn check_connection_limit(
95+
addresses: &AddressGroup,
96+
max_connections: u64,
97+
app_id: &str,
98+
) -> Result<()> {
99+
if max_connections == 0 {
100+
return Ok(());
101+
}
102+
let total: u64 = addresses
103+
.iter()
104+
.map(|a| a.counter.load(Ordering::Relaxed))
105+
.sum();
106+
if total >= max_connections {
107+
warn!(
108+
app_id,
109+
total, max_connections, "app connection limit exceeded"
110+
);
111+
bail!("app connection limit exceeded: {total}/{max_connections}");
112+
}
113+
Ok(())
114+
}
115+
91116
/// connect to multiple hosts simultaneously and return the first successful connection
92117
pub(crate) async fn connect_multiple_hosts(
93118
addresses: AddressGroup,
94119
port: u16,
120+
max_connections: u64,
121+
app_id: &str,
95122
) -> Result<(TcpStream, EnteredCounter)> {
123+
check_connection_limit(&addresses, max_connections, app_id)?;
124+
96125
let mut join_set = JoinSet::new();
97126
for addr in addresses {
98127
let counter = addr.counter.enter();
@@ -127,9 +156,10 @@ pub(crate) async fn proxy_to_app(
127156
port: u16,
128157
) -> Result<()> {
129158
let addresses = state.lock().select_top_n_hosts(app_id)?;
159+
let max_connections = state.config.proxy.max_connections_per_app;
130160
let (mut outbound, _counter) = timeout(
131161
state.config.proxy.timeouts.connect,
132-
connect_multiple_hosts(addresses.clone(), port),
162+
connect_multiple_hosts(addresses.clone(), port, max_connections, app_id),
133163
)
134164
.await
135165
.with_context(|| format!("connecting timeout to app {app_id}: {addresses:?}:{port}"))?

gateway/src/proxy/tls_terminate.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -318,9 +318,10 @@ impl Proxy {
318318
.with_context(|| format!("app {app_id} not found"))?;
319319
debug!("selected top n hosts: {addresses:?}");
320320
let tls_stream = self.tls_accept(inbound, buffer, h2).await?;
321+
let max_connections = self.config.proxy.max_connections_per_app;
321322
let (outbound, _counter) = timeout(
322323
self.config.proxy.timeouts.connect,
323-
connect_multiple_hosts(addresses, port),
324+
connect_multiple_hosts(addresses, port, max_connections, app_id),
324325
)
325326
.await
326327
.map_err(|_| anyhow!("connecting timeout"))?

0 commit comments

Comments
 (0)