feat(gateway): Multiplexing with yamux

Cargo.toml

@@ -52,6 +52,7 @@ members = [
     "verifier",
     "no_std_check",
     "size-parser",
+    "uds2yamux",
 ]
 resolver = "2"
 
@@ -159,7 +160,10 @@ rocket = { git = "https://github.com/rwf2/Rocket", branch = "master", features =
 ] }
 rocket-apitoken = { git = "https://github.com/kvinwang/rocket-apitoken", branch = "dev" }
 tokio = { version = "1.46.1" }
+tokio-util = { version = "0.7", features = ["compat"] }
 tokio-vsock = "0.7.0"
+yamux = { git = "https://github.com/kvinwang/rust-yamux", branch = "feat/ping-timeout" }
+quinn = { path = "/home/kvin/src/quinn-plain/quinn", default-features = false, features = ["runtime-tokio", "null-crypto"] }
 sysinfo = "0.35.2"
 default-net = "0.22.0"
 url = "2.5"

gateway/Cargo.toml

@@ -20,6 +20,7 @@ fs-err.workspace = true
 clap = { workspace = true, features = ["derive", "string"] }
 shared_child.workspace = true
 tokio = { workspace = true, features = ["full"] }
+tokio-util = { workspace = true, features = ["compat"] }
 rustls.workspace = true
 tokio-rustls = { workspace = true, features = ["ring"] }
 rinja.workspace = true
@@ -37,6 +38,7 @@ bytes.workspace = true
 safe-write.workspace = true
 smallvec.workspace = true
 futures.workspace = true
+yamux.workspace = true
 cmd_lib.workspace = true
 load_config.workspace = true
 dstack-kms-rpc.workspace = true
@@ -46,6 +48,7 @@ http-client = { workspace = true, features = ["prpc"] }
 sha2.workspace = true
 dstack-types.workspace = true
 serde-duration.workspace = true
+size-parser = { workspace = true, features = ["serde"] }
 reqwest = { workspace = true, features = ["json"] }
 hyper = { workspace = true, features = ["server", "http1"] }
 hyper-util = { version = "0.1", features = ["tokio"] }

gateway/dstack-app/builder/entrypoint.sh

@@ -150,6 +150,16 @@ write = "5s"
 shutdown = "5s"
 total = "5h"
 
+[core.proxy.yamux]
+listen_addr = "${YAMUX_LISTEN_ADDR:-0.0.0.0}"
+listen_port = ${YAMUX_LISTEN_PORT:-0}
+## Set to 0 to disable max connection receive window limit.
+max_connection_receive_window = ${YAMUX_MAX_CONNECTION_RECEIVE_WINDOW:-1073741824}
+max_num_streams = ${YAMUX_MAX_NUM_STREAMS:-512}
+read_after_close = ${YAMUX_READ_AFTER_CLOSE:-true}
+split_send_size = ${YAMUX_SPLIT_SEND_SIZE:-16384}
+ping_timeout = "${YAMUX_PING_TIMEOUT:-15s}"
+
 [core.recycle]
 enabled = true
 interval = "5m"

gateway/gateway.toml

@@ -71,6 +71,16 @@ app_address_ns_compat = true
 workers = 32
 external_port = 443
 
+[core.proxy.yamux]
+listen_addr = "0.0.0.0"
+listen_port = 4433
+# max_connection_receive_window = 0 to disable limit
+max_connection_receive_window = "1G"
+max_num_streams = 4096
+read_after_close = true
+split_send_size = "16K"
+ping_timeout = "15s"
+
 [core.proxy.timeouts]
 # Timeout for establishing a connection to the target app.
 connect = "5s"

gateway/src/config.rs

@@ -85,6 +85,8 @@ pub struct ProxyConfig {
     pub workers: usize,
     pub app_address_ns_prefix: String,
     pub app_address_ns_compat: bool,
+    #[serde(default)]
+    pub yamux: YamuxConfig,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -112,6 +114,35 @@ pub struct Timeouts {
     pub shutdown: Duration,
 }
 
+#[derive(Debug, Clone, Deserialize)]
+pub struct YamuxConfig {
+    pub listen_addr: Ipv4Addr,
+    pub listen_port: Option<u16>,
+    #[serde(with = "size_parser::human_size")]
+    pub max_connection_receive_window: usize,
+    pub max_num_streams: usize,
+    pub read_after_close: bool,
+    #[serde(with = "size_parser::human_size")]
+    pub split_send_size: usize,
+    /// Ping timeout for yamux connections. Set to 0s to disable.
+    #[serde(with = "serde_duration")]
+    pub ping_timeout: Duration,
+}
+
+impl Default for YamuxConfig {
+    fn default() -> Self {
+        Self {
+            listen_addr: Ipv4Addr::new(0, 0, 0, 0),
+            listen_port: None,
+            max_connection_receive_window: 1024 * 1024 * 1024,
+            max_num_streams: 4096,
+            read_after_close: true,
+            split_send_size: 16 * 1024,
+            ping_timeout: Duration::from_secs(15),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct RecycleConfig {
     pub enabled: bool,