kube_portforward/
client.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use spdy_mux::{
5    split_fastws,
6    split_raw_spdy,
7};
8use tokio_util::sync::CancellationToken;
9
10use crate::connect::upgrade_spdy_with_fallback;
11use crate::error::Error;
12use crate::recovery::RecoveryCallback;
13use crate::session::Session;
14use crate::subprotocol::Subprotocol;
15
16const DEFAULT_PING: Duration = Duration::from_secs(15);
17const DEFAULT_WATCHDOG: Duration = Duration::from_secs(30);
18const DEFAULT_DRAIN: Duration = Duration::from_secs(2);
19
20/// Pool size for the  multiplexer
21const DEFAULT_SPDY_POOL_SIZE: usize = 6;
22
23/// entry point that bundles a `kube::Client` with its cluster URL.
24#[derive(Clone)]
25pub struct Client {
26    kube: kube::Client,
27    cluster_url: http::Uri,
28}
29
30impl Client {
31    pub const fn new(kube_client: kube::Client, cluster_url: http::Uri) -> Self {
32        Self {
33            kube: kube_client,
34            cluster_url,
35        }
36    }
37
38    pub fn builder() -> ClientBuilder {
39        ClientBuilder::default()
40    }
41
42    pub fn session(
43        &self, namespace: impl Into<String>, pod: impl Into<String>, port: u16,
44    ) -> SessionBuilder<'_> {
45        SessionBuilder {
46            client: self,
47            namespace: namespace.into(),
48            pod: pod.into(),
49            port,
50            ping_interval: DEFAULT_PING,
51            watchdog_timeout: DEFAULT_WATCHDOG,
52            drain_timeout: DEFAULT_DRAIN,
53            cancel: None,
54            recovery_callback: None,
55            spdy_pool_size: DEFAULT_SPDY_POOL_SIZE,
56        }
57    }
58
59    pub const fn kube_client(&self) -> &kube::Client {
60        &self.kube
61    }
62
63    pub const fn cluster_url(&self) -> &http::Uri {
64        &self.cluster_url
65    }
66}
67
68#[derive(Default)]
69pub struct ClientBuilder {
70    kube: Option<kube::Client>,
71    cluster_url: Option<http::Uri>,
72}
73
74impl ClientBuilder {
75    pub fn kube_client(mut self, c: kube::Client) -> Self {
76        self.kube = Some(c);
77        self
78    }
79
80    pub fn cluster_url(mut self, u: http::Uri) -> Self {
81        self.cluster_url = Some(u);
82        self
83    }
84
85    pub fn build(self) -> Result<Client, Error> {
86        let kube = self
87            .kube
88            .ok_or_else(|| Error::Configuration("kube_client is required".into()))?;
89        let cluster_url = self
90            .cluster_url
91            .ok_or_else(|| Error::Configuration("cluster_url is required".into()))?;
92        Ok(Client::new(kube, cluster_url))
93    }
94}
95
96/// Builder for opening a session
97pub struct SessionBuilder<'c> {
98    client: &'c Client,
99    namespace: String,
100    pod: String,
101    port: u16,
102    #[allow(dead_code)] // spdy-mux owns its own keepalive schedule
103    ping_interval: Duration,
104    #[allow(dead_code)] // spdy-mux owns its own watchdog
105    watchdog_timeout: Duration,
106    #[allow(dead_code)] // spdy-mux drains on cancel
107    drain_timeout: Duration,
108    cancel: Option<CancellationToken>,
109    recovery_callback: Option<RecoveryCallback>,
110    spdy_pool_size: usize,
111}
112
113impl SessionBuilder<'_> {
114    /// SPDY multiplexing doesn't have a pre-allocated channel pair pool.
115    pub const fn capacity(self, _n: usize) -> Self {
116        self
117    }
118
119    /// Accepted for API stability. The  multiplexer handles its own
120    /// keepalive schedule based on idle time.
121    pub const fn keepalive(mut self, ping: Duration, watchdog: Duration) -> Self {
122        self.ping_interval = ping;
123        self.watchdog_timeout = watchdog;
124        self
125    }
126
127    pub const fn shutdown_grace(mut self, drain: Duration) -> Self {
128        self.drain_timeout = drain;
129        self
130    }
131
132    pub fn cancellation_token(mut self, t: CancellationToken) -> Self {
133        self.cancel = Some(t);
134        self
135    }
136
137    /// Number of parallel upgraded connections in the SPDY pool. Each one
138    /// gets its own reader/writer task pair
139    pub fn spdy_pool_size(mut self, n: usize) -> Self {
140        self.spdy_pool_size = n.max(1);
141        self
142    }
143
144    pub fn on_recovery<F>(mut self, cb: F) -> Self
145    where
146        F: Fn(crate::recovery::RecoverySignal) + Send + Sync + 'static,
147    {
148        self.recovery_callback = Some(Arc::new(cb));
149        self
150    }
151
152    pub async fn open(self) -> Result<Session, Error> {
153        let cancel = self.cancel.unwrap_or_default();
154        let recovery_callback: RecoveryCallback = self
155            .recovery_callback
156            .unwrap_or_else(|| Arc::new(|_signal| {}));
157
158        open_spdy_session(
159            self.client,
160            &self.namespace,
161            &self.pod,
162            self.port,
163            self.spdy_pool_size,
164            cancel,
165            recovery_callback,
166        )
167        .await
168    }
169}
170
171/// Open a SPDY session, probe with the first upgrade, then fill the rest
172/// of the pool in parallel using whatever transport the probe picked.
173async fn open_spdy_session(
174    client: &Client, namespace: &str, pod: &str, port: u16, pool_size: usize,
175    cancel: CancellationToken, recovery_callback: RecoveryCallback,
176) -> Result<Session, Error> {
177    let first = upgrade_spdy_with_fallback(
178        client.kube_client(),
179        client.cluster_url(),
180        namespace,
181        pod,
182        &recovery_callback,
183    )
184    .await?;
185    let chosen_protocol = first.protocol;
186    let first_upgraded = first.upgraded;
187
188    tracing::info!(
189        pod = %pod,
190        pool_size,
191        protocol = %chosen_protocol,
192        "SPDY tunnel: probe succeeded"
193    );
194
195    let extra_upgrades = if pool_size > 1 {
196        let t_parallel = std::time::Instant::now();
197        let mut join_set = tokio::task::JoinSet::new();
198        for i in 1..pool_size {
199            let kube = client.kube_client().clone();
200            let url = client.cluster_url().clone();
201            let ns = namespace.to_owned();
202            let pod_name = pod.to_owned();
203            join_set.spawn(async move {
204                let result = match chosen_protocol {
205                    Subprotocol::Spdy31Tunnel => {
206                        crate::connect::upgrade_spdy_tunnel(&kube, &url, &ns, &pod_name).await
207                    }
208                    Subprotocol::LegacySpdy => {
209                        crate::connect::upgrade_legacy_spdy(&kube, &url, &ns, &pod_name).await
210                    }
211                };
212                (i, result)
213            });
214        }
215        let mut succeeded = Vec::with_capacity(pool_size - 1);
216        while let Some(join_result) = join_set.join_next().await {
217            match join_result {
218                Ok((_, Ok(upgraded))) => succeeded.push(upgraded.upgraded),
219                Ok((i, Err(e))) => {
220                    tracing::debug!("SPDY pool: connection {i}/{pool_size} failed: {e}");
221                }
222                Err(e) => {
223                    tracing::debug!("SPDY pool: connection task panicked: {e}");
224                }
225            }
226        }
227        tracing::info!(
228            pool_opened = succeeded.len() + 1,
229            pool_target = pool_size,
230            elapsed_ms = u64::try_from(t_parallel.elapsed().as_millis()).unwrap_or(u64::MAX),
231            "SPDY pool: parallel connections opened"
232        );
233        succeeded
234    } else {
235        Vec::new()
236    };
237
238    let config = spdy_mux::MuxConfig {
239        pool_size: extra_upgrades.len() + 1,
240        ..Default::default()
241    };
242    let all_upgrades = std::iter::once(first_upgraded).chain(extra_upgrades);
243    let t_pool = std::time::Instant::now();
244    let spdy_session = match chosen_protocol {
245        Subprotocol::Spdy31Tunnel => {
246            let pairs: Vec<_> = all_upgrades.map(split_fastws).collect();
247            spdy_mux::Session::with_config(pairs, cancel.clone(), config).await
248        }
249        Subprotocol::LegacySpdy => {
250            let pairs: Vec<_> = all_upgrades.map(split_raw_spdy).collect();
251            spdy_mux::Session::with_config(pairs, cancel.clone(), config).await
252        }
253    }
254    .map_err(Error::from)?;
255
256    tracing::info!(
257        pod = %pod,
258        pool_healthy = spdy_session.capacity() > 0,
259        pool_init_ms = u64::try_from(t_pool.elapsed().as_millis()).unwrap_or(u64::MAX),
260        protocol = %chosen_protocol,
261        "SPDY session ready"
262    );
263    Ok(Session::from_spdy(spdy_session, chosen_protocol, port))
264}
kube_portforward/client.rs

kube_portforward/
client.rs
