Skip to content

Commit 2eafdc5

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 2eafdc5

4 files changed

Lines changed: 35 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: 29 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,34 @@ 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_rate_limit(addresses: &AddressGroup, max_connections: u64, app_id: &str) -> Result<()> {
95+
if max_connections == 0 {
96+
return Ok(());
97+
}
98+
let total: u64 = addresses
99+
.iter()
100+
.map(|a| a.counter.load(Ordering::Relaxed))
101+
.sum();
102+
if total >= max_connections {
103+
warn!(
104+
app_id,
105+
total, max_connections, "app connection limit exceeded"
106+
);
107+
bail!("app connection limit exceeded: {total}/{max_connections}");
108+
}
109+
Ok(())
110+
}
111+
91112
/// connect to multiple hosts simultaneously and return the first successful connection
92113
pub(crate) async fn connect_multiple_hosts(
93114
addresses: AddressGroup,
94115
port: u16,
116+
max_connections: u64,
117+
app_id: &str,
95118
) -> Result<(TcpStream, EnteredCounter)> {
119+
check_rate_limit(&addresses, max_connections, app_id)?;
120+
96121
let mut join_set = JoinSet::new();
97122
for addr in addresses {
98123
let counter = addr.counter.enter();
@@ -127,9 +152,10 @@ pub(crate) async fn proxy_to_app(
127152
port: u16,
128153
) -> Result<()> {
129154
let addresses = state.lock().select_top_n_hosts(app_id)?;
155+
let max_connections = state.config.proxy.max_connections_per_app;
130156
let (mut outbound, _counter) = timeout(
131157
state.config.proxy.timeouts.connect,
132-
connect_multiple_hosts(addresses.clone(), port),
158+
connect_multiple_hosts(addresses.clone(), port, max_connections, app_id),
133159
)
134160
.await
135161
.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)