sark_client/connector/
session.rs

1use std::io::Read;
2use std::ptr::NonNull;
3use std::time::{Duration, Instant};
4
5use cartel_core::{FatalSlot, Slab};
6use dope::WakerSet;
7use dope::manifold::connector;
8use dope::runtime::token::Token;
9use o3::buffer::Owned;
10use sark_core::http::Response;
11use sark_core::http::codec::{DecodeMode, HeaderLookup, Parse};
12
13use crate::connector::codec::{self, Head};
14use crate::connector::error::Error;
15use crate::connector::retry::RetryPolicy;
16
17pub(super) type Outcome = Result<Response, Error>;
18
19pub(super) const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
20
21pub(super) const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
22
23const KEEPALIVE_MARGIN: Duration = Duration::from_secs(1);
24
25#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
26pub enum DecompressionPolicy {
27    #[default]
28    Strict,
29    Lenient,
30}
31
32#[derive(Default)]
33pub struct ConnState {
34    pending_close: bool,
35}
36
37impl connector::Lifecycle for ConnState {
38    fn wants_close(&self) -> connector::Close {
39        if self.pending_close {
40            connector::Close::Reconnect
41        } else {
42            connector::Close::Keep
43        }
44    }
45
46    fn defer_close(&self) -> bool {
47        false
48    }
49
50    fn is_drained(&self) -> bool {
51        true
52    }
53}
54
55struct ConnEntry {
56    conn_id: Token,
57    slab_ix: usize,
58    last_activity: Instant,
59    keepalive: Option<Duration>,
60}
61
62pub struct Shared {
63    conns: Vec<ConnEntry>,
64    #[allow(clippy::vec_box)]
65    slabs: Vec<Box<Slab<Outcome>>>,
66    pub active_wakers: WakerSet,
67    pub fatal: FatalSlot<Error>,
68    pub host: String,
69    pub decompression: DecompressionPolicy,
70    pub max_redirects: u32,
71    pub retry: RetryPolicy,
72    pub idle_timeout: Duration,
73    pub request_timeout: Duration,
74}
75
76impl Shared {
77    fn new(host: String) -> Self {
78        Self {
79            conns: Vec::new(),
80            slabs: Vec::new(),
81            active_wakers: WakerSet::new(),
82            fatal: FatalSlot::default(),
83            host,
84            decompression: DecompressionPolicy::Strict,
85            max_redirects: 10,
86            retry: RetryPolicy::default(),
87            idle_timeout: DEFAULT_IDLE_TIMEOUT,
88            request_timeout: DEFAULT_REQUEST_TIMEOUT,
89        }
90    }
91
92    pub fn any_ready(&self) -> bool {
93        !self.conns.is_empty()
94    }
95
96    pub fn live_conns(&self) -> usize {
97        self.conns.len()
98    }
99
100    fn alloc_slab(&mut self) -> usize {
101        for ix in 0..self.slabs.len() {
102            let bound = self.conns.iter().any(|c| c.slab_ix == ix);
103            if !bound && self.slabs[ix].is_drained() {
104                return ix;
105            }
106        }
107        self.slabs.push(Box::new(Slab::new()));
108        self.slabs.len() - 1
109    }
110
111    fn note_connect(&mut self, conn_id: Token, now: Instant) {
112        if self.conns.iter().any(|c| c.conn_id == conn_id) {
113            return;
114        }
115        let slab_ix = self.alloc_slab();
116        self.conns.push(ConnEntry {
117            conn_id,
118            slab_ix,
119            last_activity: now,
120            keepalive: None,
121        });
122        self.active_wakers.drain_wake();
123    }
124
125    fn push_response(
126        &mut self,
127        conn_id: Token,
128        outcome: Outcome,
129        keepalive: Option<Duration>,
130        now: Instant,
131    ) {
132        let Some(pos) = self.conns.iter().position(|c| c.conn_id == conn_id) else {
133            return;
134        };
135        let slab_ix = {
136            let c = &mut self.conns[pos];
137            c.last_activity = now;
138            if keepalive.is_some() {
139                c.keepalive = keepalive;
140            }
141            c.slab_ix
142        };
143        self.slabs[slab_ix].push(outcome);
144        self.slabs[slab_ix].complete();
145        self.active_wakers.drain_wake();
146    }
147
148    fn fail_connection(&mut self, conn_id: Token, fatal: Option<String>) {
149        if let Some(pos) = self.conns.iter().position(|c| c.conn_id == conn_id) {
150            let slab_ix = self.conns[pos].slab_ix;
151            self.slabs[slab_ix].fail_all(|| match &fatal {
152                Some(m) => Err(Error::Http(m.clone())),
153                None => Err(Error::Closed),
154            });
155            self.conns.remove(pos);
156        }
157        self.active_wakers.drain_wake();
158    }
159
160    pub fn drop_conn(&mut self, conn_id: Token) {
161        if let Some(pos) = self.conns.iter().position(|c| c.conn_id == conn_id) {
162            let slab_ix = self.conns[pos].slab_ix;
163            self.slabs[slab_ix].fail_all(|| Err(Error::Closed));
164            self.conns.remove(pos);
165        }
166    }
167
168    pub fn acquire(&mut self, now: Instant, idle_timeout: Duration) -> (Option<Token>, Vec<Token>) {
169        let mut chosen: Option<usize> = None;
170        let mut best_depth = usize::MAX;
171        let mut recycle_idx: Vec<usize> = Vec::new();
172        for (i, c) in self.conns.iter().enumerate() {
173            let limit = c
174                .keepalive
175                .map(|k| k.saturating_sub(KEEPALIVE_MARGIN))
176                .unwrap_or(idle_timeout);
177            let stale = now.saturating_duration_since(c.last_activity) >= limit;
178            let depth = self.slabs[c.slab_ix].depth();
179            if stale {
180                recycle_idx.push(i);
181                continue;
182            }
183            if depth < best_depth {
184                best_depth = depth;
185                chosen = Some(i);
186            }
187        }
188        let chosen_tok = chosen.map(|i| self.conns[i].conn_id);
189        let mut recycle = Vec::with_capacity(recycle_idx.len());
190        recycle_idx.sort_unstable();
191        for &i in recycle_idx.iter().rev() {
192            recycle.push(self.conns[i].conn_id);
193            self.conns.remove(i);
194        }
195        (chosen_tok, recycle)
196    }
197
198    pub fn slab_ptr_for(&mut self, conn_id: Token) -> Option<NonNull<Slab<Outcome>>> {
199        let pos = self.conns.iter().position(|c| c.conn_id == conn_id)?;
200        let slab_ix = self.conns[pos].slab_ix;
201        Some(NonNull::from(&mut *self.slabs[slab_ix]))
202    }
203
204    pub fn touch(&mut self, conn_id: Token, now: Instant) {
205        if let Some(c) = self.conns.iter_mut().find(|c| c.conn_id == conn_id) {
206            c.last_activity = now;
207        }
208    }
209}
210
211pub struct Session {
212    codec: codec::Codec,
213    pub shared: Shared,
214}
215
216impl Session {
217    pub fn new(host: impl Into<String>) -> Self {
218        Self {
219            codec: codec::Codec::default(),
220            shared: Shared::new(host.into()),
221        }
222    }
223
224    pub fn with_decompression(host: impl Into<String>, policy: DecompressionPolicy) -> Self {
225        let mut session = Self::new(host);
226        session.shared.decompression = policy;
227        session
228    }
229
230    pub fn max_response_body(&mut self, cap: usize) -> &mut Self {
231        self.codec.max_response_body = cap;
232        self
233    }
234
235    pub fn with_max_redirects(mut self, max: u32) -> Self {
236        self.shared.max_redirects = max;
237        self
238    }
239
240    pub fn with_retry(mut self, retry: RetryPolicy) -> Self {
241        self.shared.retry = retry;
242        self
243    }
244
245    pub fn with_idle_timeout(mut self, idle: Duration) -> Self {
246        self.shared.idle_timeout = idle;
247        self
248    }
249
250    pub fn with_request_timeout(mut self, dur: Duration) -> Self {
251        self.shared.request_timeout = dur;
252        self
253    }
254}
255
256impl connector::Session for Session {
257    type Codec = codec::Codec;
258    type ConnState = ConnState;
259
260    fn codec(&self) -> &codec::Codec {
261        &self.codec
262    }
263
264    fn connect(&mut self, ctx: &mut connector::Ctx<'_, Self>) {
265        self.shared.fatal.clear();
266        self.shared.note_connect(ctx.conn_id, Instant::now());
267    }
268
269    fn response(&mut self, head: Head, ctx: &mut connector::Ctx<'_, Self>) {
270        if let Some(reason) = head.error {
271            self.shared.push_response(
272                ctx.conn_id,
273                Err(Error::Parse(reason.into())),
274                None,
275                Instant::now(),
276            );
277            ctx.state.pending_close = true;
278            return;
279        }
280        let bytes = head.full.as_ref();
281        let (outcome, keep_alive, keepalive_timeout) =
282            match Parse::response(bytes, DecodeMode::Response) {
283                Ok(Some(mut resp)) => {
284                    let keep = Self::should_keep_alive(&resp);
285                    let timeout = Self::keepalive_timeout(&resp);
286                    let outcome = match Self::decompress(
287                        &mut resp,
288                        self.shared.decompression,
289                        self.codec.max_response_body,
290                    ) {
291                        Ok(()) => Ok(resp),
292                        Err(e) => Err(e),
293                    };
294                    (outcome, keep, timeout)
295                }
296                Ok(None) => (
297                    Err(Error::Parse("incomplete response frame".into())),
298                    true,
299                    None,
300                ),
301                Err(e) => (Err(Error::Parse(e.to_string())), true, None),
302            };
303        if !keep_alive {
304            ctx.state.pending_close = true;
305        }
306        self.shared
307            .push_response(ctx.conn_id, outcome, keepalive_timeout, Instant::now());
308    }
309
310    fn disconnect(&mut self, ctx: &mut connector::Ctx<'_, Self>) {
311        let fatal_msg = self.shared.fatal.as_ref().map(|e| e.to_string());
312        self.shared.fail_connection(ctx.conn_id, fatal_msg);
313        ctx.state.pending_close = false;
314    }
315}
316
317impl Session {
318    fn should_keep_alive(resp: &Response) -> bool {
319        let headers = resp.headers();
320        if headers.has_token(http::header::CONNECTION, "close")
321            || headers.has_token(http::header::CONNECTION, "upgrade")
322        {
323            return false;
324        }
325        if headers.has_token(http::header::CONNECTION, "keep-alive") {
326            return true;
327        }
328
329        let status = resp.status().as_u16();
330        if status < 200 || status == 204 || status == 304 {
331            return true;
332        }
333
334        headers.contains_key(http::header::CONTENT_LENGTH)
335            || headers.value_eq_ascii_case(http::header::TRANSFER_ENCODING, "chunked")
336    }
337
338    fn keepalive_timeout(resp: &Response) -> Option<Duration> {
339        let name = http::header::HeaderName::from_static("keep-alive");
340        let raw = resp.headers().get(name)?;
341        let value = raw.to_str().ok()?;
342        for part in value.split(',') {
343            let part = part.trim();
344            if let Some(rest) = part.strip_prefix("timeout=")
345                && let Ok(secs) = rest.trim().parse::<u64>()
346            {
347                return Some(Duration::from_secs(secs));
348            }
349        }
350        None
351    }
352
353    fn decompress(
354        resp: &mut Response,
355        policy: DecompressionPolicy,
356        max_body: usize,
357    ) -> Result<(), Error> {
358        let is_gzip = resp
359            .headers()
360            .value_eq_ascii_case(http::header::CONTENT_ENCODING, "gzip");
361        if !is_gzip || resp.body().is_empty() {
362            return Ok(());
363        }
364
365        let limit = max_body as u64;
366        let mut decoder = flate2::read::GzDecoder::new(resp.body()).take(limit + 1);
367        let mut decompressed = Vec::new();
368        match decoder.read_to_end(&mut decompressed) {
369            Ok(_) if decompressed.len() as u64 > limit => Err(Error::Parse(
370                "decompressed response body exceeds size limit".into(),
371            )),
372            Ok(_) => {
373                resp.set_body(Owned::from(&decompressed[..]));
374                resp.headers_mut().remove("content-encoding");
375                resp.headers_mut().remove("content-length");
376                Ok(())
377            }
378            Err(e) if policy == DecompressionPolicy::Strict => {
379                Err(Error::Parse(format!("gzip decompression failed: {e}")))
380            }
381            Err(_) => Ok(()),
382        }
383    }
384}
sark_client/connector/session.rs

sark_client/connector/
session.rs
