kube_portforward/
session.rs

1use std::sync::atomic::{
2    AtomicBool,
3    AtomicU32,
4    Ordering,
5};
6
7use crossbeam_queue::ArrayQueue;
8use tokio_util::sync::CancellationToken;
9
10use crate::error::Error;
11use crate::stream::Stream;
12use crate::subprotocol::Subprotocol;
13
14const SPARE_STREAM_CAP: usize = 16;
15
16/// When spare count drops to or below this threshold, the background
17/// replenisher refills
18const SPARE_STREAM_LOW_WATERMARK: usize = 8;
19
20/// One header  carried on a SPDY SYN_STREAM frame.
21type SpdyHeader = (String, String);
22
23/// The error stream and data stream header lists for one paired
24/// `portforward.k8s.io` connection.
25type PortforwardHeaderPair = (Vec<SpdyHeader>, Vec<SpdyHeader>);
26
27/// One port forward session that multiplexes many concurrent
28/// local connections over a pool of upgraded connections to the apiserver.
29pub struct Session {
30    inner: spdy_mux::Session,
31    protocol: Subprotocol,
32    /// Target pod port. The kubelet expects this in the SYN_STREAM
33    /// `port` header for every paired stream we open.
34    port: u16,
35    /// request id counter, kubelet uses this header to
36    /// pair the data and error streams of one logical TCP connection.
37    next_request_id: AtomicU32,
38    /// Pre-opened spare streams for instant connect(). Background task
39    /// replenishes when count drops to or below `SPARE_STREAM_LOW_WATERMARK`.
40    spare_streams: ArrayQueue<Stream>,
41    /// Guard against concurrent replenishment. Set by `replenish_spare_streams`
42    /// on entry, cleared on exit.
43    replenishing: AtomicBool,
44}
45
46impl Session {
47    pub(crate) fn from_spdy(session: spdy_mux::Session, protocol: Subprotocol, port: u16) -> Self {
48        Self {
49            spare_streams: ArrayQueue::new(SPARE_STREAM_CAP),
50            replenishing: AtomicBool::new(false),
51            inner: session,
52            protocol,
53            port,
54            next_request_id: AtomicU32::new(0),
55        }
56    }
57
58    /// Build the K8s `portforward.k8s.io v1`  headers for one
59    /// stream-pair connection, header names are lowercase
60    fn portforward_headers(&self) -> PortforwardHeaderPair {
61        let request_id = self
62            .next_request_id
63            .fetch_add(1, Ordering::Relaxed)
64            .to_string();
65        let port = self.port.to_string();
66        let error_headers = vec![
67            ("streamtype".to_string(), "error".to_string()),
68            ("port".to_string(), port.clone()),
69            ("requestid".to_string(), request_id.clone()),
70        ];
71        let data_headers = vec![
72            ("streamtype".to_string(), "data".to_string()),
73            ("port".to_string(), port),
74            ("requestid".to_string(), request_id),
75        ];
76        (error_headers, data_headers)
77    }
78
79    /// Grab the next stream and return a bidirectional [`Stream`].
80    pub async fn connect(&self) -> Result<Stream, Error> {
81        while let Some(stream) = self.spare_streams.pop() {
82            if !stream.is_read_closed() {
83                return Ok(stream);
84            }
85            tracing::debug!("spare stream stale (remote closed while idle), discarding");
86        }
87        self.open_new_stream().await
88    }
89
90    async fn open_new_stream(&self) -> Result<Stream, Error> {
91        let (error_headers, data_headers) = self.portforward_headers();
92        self.inner
93            .open_stream_pair(error_headers, data_headers)
94            .await
95            .map(Stream::from_spdy)
96            .map_err(Error::from)
97    }
98
99    /// Pre open spare streams up to `SPARE_STREAM_CAP`.
100    pub async fn replenish_spare_streams(&self) {
101        if self
102            .replenishing
103            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
104            .is_err()
105        {
106            return;
107        }
108        let _guard = ReplenishGuard(&self.replenishing);
109
110        while self.spare_streams.len() < SPARE_STREAM_CAP {
111            if self.is_full() || self.cancellation_token().is_cancelled() {
112                break;
113            }
114            match self.open_new_stream().await {
115                Ok(stream) => {
116                    if self.spare_streams.push(stream).is_err() {
117                        break;
118                    }
119                }
120                Err(_) => break,
121            }
122        }
123    }
124
125    pub fn spare_count(&self) -> usize {
126        self.spare_streams.len()
127    }
128
129    pub fn needs_replenish(&self) -> bool {
130        self.spare_count() <= SPARE_STREAM_LOW_WATERMARK
131    }
132
133    pub const fn protocol(&self) -> Subprotocol {
134        self.protocol
135    }
136
137    /// Max concurrent streams this session can hold.
138    pub fn capacity(&self) -> usize {
139        self.inner.capacity()
140    }
141
142    pub fn operating_capacity(&self) -> usize {
143        self.inner.operating_capacity()
144    }
145
146    pub fn in_use(&self) -> usize {
147        self.inner.in_use()
148    }
149
150    pub fn available(&self) -> usize {
151        self.inner.available()
152    }
153
154    pub fn is_full(&self) -> bool {
155        self.inner.is_full()
156    }
157
158    pub fn is_drained(&self) -> bool {
159        self.inner.is_drained()
160    }
161
162    pub fn cancellation_token(&self) -> CancellationToken {
163        self.inner.cancellation_token()
164    }
165
166    /// Gracefully close the session.
167    pub async fn close(self) -> Result<(), Error> {
168        self.inner.close().await.map_err(Error::from)
169    }
170}
171
172/// guard that clears the `replenishing` flag on drop.
173struct ReplenishGuard<'a>(&'a AtomicBool);
174
175impl Drop for ReplenishGuard<'_> {
176    fn drop(&mut self) {
177        self.0.store(false, Ordering::Release);
178    }
179}
kube_portforward/session.rs

kube_portforward/
session.rs
