transport: add data_frame_budget to client Endpoint

tests/integration_tests/tests/data_frame_budget.rs

@@ -0,0 +1,79 @@
+/*
+ *
+ * Copyright 2025 gRPC authors.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *
+ */
+
+use std::time::Duration;
+
+use integration_tests::pb::{test_client, test_server, Input, Output};
+use tokio::sync::oneshot;
+use tonic::{
+    transport::{Endpoint, Server},
+    Request, Response, Status,
+};
+
+#[tokio::test]
+async fn data_frame_budget_on_client_endpoint() {
+    struct Svc;
+
+    #[tonic::async_trait]
+    impl test_server::Test for Svc {
+        async fn unary_call(&self, _: Request<Input>) -> Result<Response<Output>, Status> {
+            Ok(Response::new(Output {}))
+        }
+    }
+
+    let svc = test_server::TestServer::new(Svc);
+
+    let (tx, rx) = oneshot::channel::<()>();
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = format!("http://{}", listener.local_addr().unwrap());
+
+    let jh = tokio::spawn(async move {
+        let listener =
+            tonic::transport::server::TcpIncoming::from(listener).with_nodelay(Some(true));
+        Server::builder()
+            .add_service(svc)
+            .serve_with_incoming_shutdown(listener, async { drop(rx.await) })
+            .await
+            .unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    // 10x h2's own default (25,600 bytes as of h2 v0.4.18), matching the
+    // budget a caller would raise this to when receiving streams made of
+    // many small DATA frames.
+    let channel = Endpoint::from_shared(addr)
+        .unwrap()
+        .data_frame_budget(256_000usize)
+        .connect()
+        .await
+        .unwrap();
+
+    let mut client = test_client::TestClient::new(channel);
+
+    client.unary_call(Request::new(Input {})).await.unwrap();
+
+    tx.send(()).unwrap();
+    jh.await.unwrap();
+}

tonic/src/transport/channel/endpoint.rs

@@ -67,6 +67,7 @@ pub struct Endpoint {
     pub(crate) init_stream_window_size: Option<u32>,
     pub(crate) init_connection_window_size: Option<u32>,
     pub(crate) max_frame_size: Option<u32>,
+    pub(crate) data_frame_budget: Option<usize>,
     pub(crate) tcp_keepalive: Option<Duration>,
     pub(crate) tcp_keepalive_interval: Option<Duration>,
     pub(crate) tcp_keepalive_retries: Option<u32>,
@@ -117,6 +118,7 @@ impl Endpoint {
             init_stream_window_size: None,
             init_connection_window_size: None,
             max_frame_size: None,
+            data_frame_budget: None,
             tcp_keepalive: None,
             tcp_keepalive_interval: None,
             tcp_keepalive_retries: None,
@@ -148,6 +150,7 @@ impl Endpoint {
             init_stream_window_size: None,
             init_connection_window_size: None,
             max_frame_size: None,
+            data_frame_budget: None,
             tcp_keepalive: None,
             tcp_keepalive_interval: None,
             tcp_keepalive_retries: None,
@@ -533,6 +536,36 @@ impl Endpoint {
         }
     }
 
+    /// Sets a connection-level budget for limiting memory overhead from
+    /// received small DATA frames.
+    ///
+    /// This guards against `GOAWAY(ENHANCE_YOUR_CALM)` disconnects on
+    /// streams that receive many small DATA frames back-to-back, such as a
+    /// server-streaming RPC that emits one small message per frame: HTTP/2
+    /// flow control accounts for DATA payload bytes but not the per-frame
+    /// buffering overhead, so an excessive number of small frames can
+    /// consume disproportionate memory relative to the bytes they carry.
+    ///
+    /// Passing `None` will do nothing.
+    ///
+    /// If not set, will default from underlying transport. As of `h2`
+    /// v0.4.18, that default is 25,600 bytes.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use tonic::transport::Endpoint;
+    /// # let builder = Endpoint::from_static("https://example.com");
+    /// let endpoint = builder.data_frame_budget(256_000usize);
+    /// ```
+    #[must_use]
+    pub fn data_frame_budget(self, budget: impl Into<Option<usize>>) -> Self {
+        Endpoint {
+            data_frame_budget: budget.into(),
+            ..self
+        }
+    }
+
     /// Sets the executor used to spawn async tasks.
     ///
     /// Uses `tokio::spawn` by default.

tonic/src/transport/channel/service/connection.rs

@@ -67,6 +67,10 @@
             settings.max_frame_size(val);
         }
 
+        if let Some(val) = endpoint.data_frame_budget {
+            settings.data_frame_budget(val);
+        }
+
         if let Some(val) = endpoint.http2_keep_alive_timeout {
             settings.keep_alive_timeout(val);
         }