sark_client/connector/
session.rs

1use std::cell::Cell;
2use std::io::{self, Read};
3use std::pin::Pin;
4use std::rc::Rc;
5use std::time::{Duration, Instant};
6
7use cartel_core::{Arena, Limits};
8use dope::driver::token::Token;
9use dope::manifold::connector;
10use dope::manifold::timer::Timer;
11use dope::runtime::Idle;
12use dope_fiber::WaitQueue;
13use o3::buffer::{Lease, Pool as BufferPool, PoolLayout};
14use o3::cell::RawCell;
15use o3::collections::SlotQueue;
16use o3::mem::ByteBudget;
17use sark_core::http::Response;
18use sark_core::http::codec::{DecodeMode, HeaderLookup, Parse};
19
20use crate::connector::codec::{self, Head};
21use crate::connector::error::Error;
22use crate::connector::retry::RetryPolicy;
23
24pub(super) type Outcome = Result<Response, Error>;
25
26pub(super) const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
27
28pub(super) const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
29
30pub const DEFAULT_MAX_INFLIGHT_PER_CONNECTION: usize = 256;
31
32pub const DEFAULT_RESPONSE_BUFFER_CAPACITY: usize = 64 * 1024 * 1024;
33
34pub const DEFAULT_REQUEST_SLOTS: u32 = 256;
35
36pub const DEFAULT_REQUEST_CAPACITY: u32 = 64 * 1024;
37
38const KEEPALIVE_MARGIN: Duration = Duration::from_secs(1);
39
40#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
41pub enum DecompressionPolicy {
42    #[default]
43    Strict,
44    Lenient,
45}
46#[derive(Default)]
47pub struct ConnState {
48    pending_close: bool,
49}
50
51impl connector::Lifecycle for ConnState {
52    fn wants_close(&self) -> connector::Close {
53        if self.pending_close {
54            connector::Close::Reconnect
55        } else {
56            connector::Close::Keep
57        }
58    }
59
60    fn defer_close(&self) -> bool {
61        false
62    }
63
64    fn is_drained(&self) -> bool {
65        true
66    }
67}
68
69struct ConnEntry<'d> {
70    conn_id: Cell<Option<Token>>,
71    arena: Arena<'d, Outcome>,
72    last_activity: Cell<Option<Instant>>,
73    keepalive: Cell<Option<Duration>>,
74    queued: Cell<bool>,
75}
76
77struct Pool<'d> {
78    entries: Box<[ConnEntry<'d>]>,
79    ready: RawCell<SlotQueue<Token>>,
80    live: Cell<usize>,
81    _budget: Pin<Rc<ByteBudget>>,
82}
83
84impl<'d> Pool<'d> {
85    fn new(capacity: usize, max_inflight: usize, limit: usize) -> Self {
86        let budget = Rc::pin(ByteBudget::new(limit));
87        let limits = Limits::new(1, limit, 1);
88        Self {
89            entries: (0..capacity)
90                .map(|_| ConnEntry {
91                    conn_id: Cell::new(None),
92                    arena: Arena::with_shared_budget(max_inflight, limits, budget.clone()),
93                    last_activity: Cell::new(None),
94                    keepalive: Cell::new(None),
95                    queued: Cell::new(false),
96                })
97                .collect(),
98            ready: RawCell::new(SlotQueue::with_capacity(capacity)),
99            live: Cell::new(0),
100            _budget: budget,
101        }
102    }
103
104    fn entry<'a>(&'a self, conn_id: Token) -> Option<&'a ConnEntry<'d>> {
105        let entry = self.entries.get(conn_id.slot().raw() as usize)?;
106        (entry.conn_id.get() == Some(conn_id)).then_some(entry)
107    }
108
109    fn push_ready(&self, entry: &ConnEntry<'d>, conn_id: Token) {
110        if entry.queued.replace(true) {
111            return;
112        }
113        unsafe {
114            self.ready.with_mut(|ready| {
115                ready
116                    .vacant_entry(conn_id.slot().raw() as usize)
117                    .expect("ready queue entry must be vacant")
118                    .push_back(conn_id);
119            })
120        };
121    }
122
123    fn pop_ready(&self) -> Option<Token> {
124        unsafe { self.ready.with_mut(SlotQueue::pop_front) }
125    }
126
127    fn remove_ready(&self, conn_id: Token) {
128        unsafe {
129            self.ready
130                .with_mut(|ready| ready.remove(conn_id.slot().raw() as usize))
131        };
132    }
133}
134
135pub struct Shared<'d> {
136    pool: Pool<'d>,
137    active_waiters: Pin<Box<WaitQueue>>,
138    pub(super) host: String,
139    pub(super) origin: http::Uri,
140    pub(super) decompression: DecompressionPolicy,
141    pub(super) max_redirects: u32,
142    pub(super) retry: RetryPolicy,
143    pub(super) idle_timeout: Duration,
144    pub(super) request_timeout: Duration,
145}
146
147impl<'d> Shared<'d> {
148    pub fn has_connection(&self) -> bool {
149        self.pool.live.get() != 0
150    }
151
152    pub(super) fn connection_count(&self) -> usize {
153        self.pool.live.get()
154    }
155
156    pub(super) fn wake(&self) {
157        self.active_waiters.as_ref().wake();
158    }
159
160    pub(super) fn try_register_active(
161        &self,
162        waiter: Pin<&dope_fiber::Waiter<'d>>,
163        context: Pin<&dope_fiber::Context<'_, 'd>>,
164    ) -> bool {
165        self.active_waiters.as_ref().try_register(waiter, context)
166    }
167
168    fn note_connect(&self, conn_id: Token, now: Instant) {
169        let entry = &self.pool.entries[conn_id.slot().raw() as usize];
170        if entry.conn_id.replace(Some(conn_id)).is_none() {
171            self.pool.live.set(self.pool.live.get() + 1);
172        }
173        entry.last_activity.set(Some(now));
174        entry.keepalive.set(None);
175        self.pool.push_ready(entry, conn_id);
176        self.wake();
177    }
178
179    fn push_response(
180        &self,
181        conn_id: Token,
182        outcome: Outcome,
183        bytes: usize,
184        keepalive: Option<Duration>,
185        now: Instant,
186    ) {
187        let Some(entry) = self.pool.entry(conn_id) else {
188            return;
189        };
190        entry.last_activity.set(Some(now));
191        if keepalive.is_some() {
192            entry.keepalive.set(keepalive);
193        }
194        entry.arena.try_push(outcome, bytes, 1);
195        entry.arena.complete();
196        if entry.arena.can_register() {
197            self.pool.push_ready(entry, conn_id);
198        }
199        self.wake();
200    }
201
202    pub(super) fn close_connection(&self, conn_id: Token) {
203        let Some(entry) = self.pool.entry(conn_id) else {
204            return;
205        };
206        entry.arena.fail_all(|| Err(Error::Closed));
207        entry.conn_id.set(None);
208        self.pool.live.set(self.pool.live.get() - 1);
209        entry.last_activity.set(None);
210        entry.keepalive.set(None);
211        entry.queued.set(false);
212        self.pool.remove_ready(conn_id);
213        self.wake();
214    }
215
216    pub fn acquire(
217        &self,
218        now: Instant,
219        idle_timeout: Duration,
220        mut recycle: impl FnMut(Token),
221    ) -> Option<Token> {
222        while let Some(conn_id) = self.pool.pop_ready() {
223            let Some(entry) = self.pool.entry(conn_id) else {
224                continue;
225            };
226            entry.queued.set(false);
227            let limit = entry
228                .keepalive
229                .get()
230                .map(|k| k.saturating_sub(KEEPALIVE_MARGIN))
231                .unwrap_or(idle_timeout);
232            let stale = entry
233                .last_activity
234                .get()
235                .is_some_and(|last| now.saturating_duration_since(last) >= limit);
236            if stale {
237                self.close_connection(conn_id);
238                recycle(conn_id);
239                continue;
240            }
241            if !entry.arena.can_register() {
242                continue;
243            }
244            return Some(conn_id);
245        }
246        None
247    }
248
249    pub fn arena(&'d self, conn_id: Token) -> Option<&'d Arena<'d, Outcome>> {
250        Some(&self.pool.entry(conn_id)?.arena)
251    }
252
253    pub fn submitted(&self, conn_id: Token, now: Instant) {
254        if let Some(entry) = self.pool.entry(conn_id) {
255            entry.last_activity.set(Some(now));
256            if entry.arena.can_register() {
257                self.pool.push_ready(entry, conn_id);
258            }
259        }
260    }
261
262    pub fn make_available(&self, conn_id: Token) {
263        if let Some(entry) = self.pool.entry(conn_id)
264            && entry.arena.can_register()
265        {
266            self.pool.push_ready(entry, conn_id);
267        }
268    }
269}
270
271pub struct Config {
272    codec: codec::Codec,
273    host: String,
274    origin: http::Uri,
275    decompression: DecompressionPolicy,
276    max_redirects: u32,
277    retry: RetryPolicy,
278    idle_timeout: Duration,
279    request_timeout: Duration,
280    max_inflight_per_connection: usize,
281    response_buffer_capacity: usize,
282    request_slots: u32,
283    request_capacity: u32,
284}
285
286impl Config {
287    pub fn new(host: impl Into<String>) -> Self {
288        let host = host.into();
289        let origin = format!("http://{host}/")
290            .parse()
291            .expect("invalid HTTP host");
292        Self {
293            codec: codec::Codec::default(),
294            host,
295            origin,
296            decompression: DecompressionPolicy::Strict,
297            max_redirects: 10,
298            retry: RetryPolicy::default(),
299            idle_timeout: DEFAULT_IDLE_TIMEOUT,
300            request_timeout: DEFAULT_REQUEST_TIMEOUT,
301            max_inflight_per_connection: DEFAULT_MAX_INFLIGHT_PER_CONNECTION,
302            response_buffer_capacity: DEFAULT_RESPONSE_BUFFER_CAPACITY,
303            request_slots: DEFAULT_REQUEST_SLOTS,
304            request_capacity: DEFAULT_REQUEST_CAPACITY,
305        }
306    }
307
308    pub fn with_decompression(host: impl Into<String>, policy: DecompressionPolicy) -> Self {
309        let mut config = Self::new(host);
310        config.decompression = policy;
311        config
312    }
313
314    pub fn max_response_body(mut self, cap: usize) -> Self {
315        self.codec.max_response_body = cap;
316        self
317    }
318
319    pub fn max_redirects(mut self, max: u32) -> Self {
320        self.max_redirects = max;
321        self
322    }
323
324    pub fn retry(mut self, retry: RetryPolicy) -> Self {
325        self.retry = retry;
326        self
327    }
328
329    pub fn idle_timeout(mut self, idle: Duration) -> Self {
330        self.idle_timeout = idle;
331        self
332    }
333
334    pub fn request_timeout(mut self, dur: Duration) -> Self {
335        self.request_timeout = dur;
336        self
337    }
338
339    pub fn max_inflight_per_connection(mut self, max: usize) -> Self {
340        self.max_inflight_per_connection = max;
341        self
342    }
343
344    pub fn response_buffer_capacity(mut self, capacity: usize) -> Self {
345        self.response_buffer_capacity = capacity;
346        self
347    }
348
349    pub fn request_pool(mut self, slots: u32, capacity: u32) -> Self {
350        self.request_slots = slots;
351        self.request_capacity = capacity;
352        self
353    }
354
355    fn request_pool_layout(&self) -> io::Result<PoolLayout> {
356        if self.request_slots == 0 {
357            return Err(io::Error::new(
358                io::ErrorKind::InvalidInput,
359                "request pool must have slots",
360            ));
361        }
362        PoolLayout::new(self.request_slots, self.request_capacity)
363            .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))
364    }
365}
366
367pub struct Port<'d> {
368    pub(super) io: connector::Port<'d, Lease<'d>>,
369    pub(super) shared: Shared<'d>,
370    codec: codec::Codec,
371    timer: Timer<'d>,
372    pub(super) requests: Pin<Box<BufferPool>>,
373}
374
375pub struct PortFactory {
376    config: Config,
377    capacity: usize,
378    timer_capacity: usize,
379    request_pool: PoolLayout,
380}
381
382impl<'d> Port<'d> {
383    pub fn new(
384        config: Config,
385        capacity: usize,
386        timer_capacity: usize,
387        driver: dope::DriverRef<'d>,
388    ) -> io::Result<Self> {
389        let request_pool = config.request_pool_layout()?;
390        Ok(Self::new_with_pool(
391            config,
392            capacity,
393            timer_capacity,
394            request_pool,
395            driver,
396        ))
397    }
398
399    fn new_with_pool(
400        config: Config,
401        capacity: usize,
402        timer_capacity: usize,
403        request_pool: PoolLayout,
404        driver: dope::DriverRef<'d>,
405    ) -> Self {
406        let Config {
407            codec,
408            host,
409            origin,
410            decompression,
411            max_redirects,
412            retry,
413            idle_timeout,
414            request_timeout,
415            max_inflight_per_connection,
416            response_buffer_capacity,
417            request_slots: _,
418            request_capacity: _,
419        } = config;
420        let shared = Shared {
421            pool: Pool::new(
422                capacity,
423                max_inflight_per_connection,
424                response_buffer_capacity,
425            ),
426            active_waiters: Box::pin(WaitQueue::with_capacity(timer_capacity)),
427            host,
428            origin,
429            decompression,
430            max_redirects,
431            retry,
432            idle_timeout,
433            request_timeout,
434        };
435        Self {
436            io: connector::Port::with_capacity(capacity, driver),
437            shared,
438            codec,
439            timer: Timer::with_capacity(timer_capacity, driver),
440            requests: Box::pin(BufferPool::new(request_pool)),
441        }
442    }
443
444    pub fn capacity(&self) -> usize {
445        self.io.capacity()
446    }
447
448    pub fn factory(
449        config: Config,
450        capacity: usize,
451        timer_capacity: usize,
452    ) -> io::Result<PortFactory> {
453        let request_pool = config.request_pool_layout()?;
454        Ok(PortFactory {
455            config,
456            capacity,
457            timer_capacity,
458            request_pool,
459        })
460    }
461
462    pub(super) fn timer(&'d self) -> &'d Timer<'d> {
463        &self.timer
464    }
465}
466
467impl dope::runtime::StorageFactory for PortFactory {
468    type Output<'d> = Port<'d>;
469
470    fn build<'d>(self, driver: &mut dope::DriverContext<'_, 'd>) -> Self::Output<'d> {
471        Port::new_with_pool(
472            self.config,
473            self.capacity,
474            self.timer_capacity,
475            self.request_pool,
476            driver.driver_ref(),
477        )
478    }
479}
480
481pub struct Session<'d> {
482    pub(super) port: &'d Port<'d>,
483}
484
485impl<'d> Session<'d> {
486    pub fn new(port: &'d Port<'d>) -> Self {
487        Self { port }
488    }
489}
490
491#[dope_gen::connector_session(codec = port.codec, io = port.io)]
492impl<'d> connector::Session<'d> for Session<'d> {
493    type Codec = codec::Codec;
494    type ConnState = ConnState;
495    type Send = Lease<'d>;
496
497    fn connect(&mut self, ctx: &mut connector::Ctx<'_, 'd, Self>) {
498        self.port.shared.note_connect(ctx.conn_id, Instant::now());
499    }
500
501    fn response(&mut self, head: Head, ctx: &mut connector::Ctx<'_, 'd, Self>) {
502        if let Some(reason) = head.error {
503            let bytes = head.full.len();
504            self.port.shared.push_response(
505                ctx.conn_id,
506                Err(Error::Parse(reason.into())),
507                bytes,
508                None,
509                Instant::now(),
510            );
511            ctx.state.pending_close = true;
512            return;
513        }
514        let buffered = head.full.len();
515        let bytes = head.full.as_ref();
516        let (outcome, keep_alive, keepalive_timeout) =
517            match Parse::response(bytes, DecodeMode::Response) {
518                Ok(Some(mut resp)) => {
519                    let keep = Self::should_keep_alive(&resp);
520                    let timeout = Self::keepalive_timeout(&resp);
521                    let outcome = match Self::decompress(
522                        &mut resp,
523                        self.port.shared.decompression,
524                        self.port.codec.max_response_body,
525                    ) {
526                        Ok(()) => Ok(resp),
527                        Err(e) => Err(e),
528                    };
529                    (outcome, keep, timeout)
530                }
531                Ok(None) => (
532                    Err(Error::Parse("incomplete response frame".into())),
533                    true,
534                    None,
535                ),
536                Err(e) => (Err(Error::Parse(e.to_string())), true, None),
537            };
538        if !keep_alive {
539            ctx.state.pending_close = true;
540        }
541        let buffered = outcome
542            .as_ref()
543            .map_or(buffered, |response| buffered.max(response.body().len()));
544        self.port.shared.push_response(
545            ctx.conn_id,
546            outcome,
547            buffered,
548            keepalive_timeout,
549            Instant::now(),
550        );
551    }
552
553    fn disconnect(&mut self, ctx: &mut connector::Ctx<'_, 'd, Self>) {
554        self.port.io.deactivate(ctx.conn_id);
555        self.port.shared.close_connection(ctx.conn_id);
556        ctx.state.pending_close = false;
557    }
558
559    fn pre_park(&mut self) {
560        self.port.timer.expire(Instant::now());
561    }
562
563    fn idle(&self) -> Idle {
564        Idle::Park(self.port.timer.earliest())
565    }
566}
567
568impl Session<'_> {
569    fn should_keep_alive(resp: &Response) -> bool {
570        let headers = resp.headers();
571        if headers.has_token(http::header::CONNECTION, "close")
572            || headers.has_token(http::header::CONNECTION, "upgrade")
573        {
574            return false;
575        }
576        if headers.has_token(http::header::CONNECTION, "keep-alive") {
577            return true;
578        }
579
580        let status = resp.status().as_u16();
581        if status < 200 || status == 204 || status == 304 {
582            return true;
583        }
584
585        headers.contains_key(http::header::CONTENT_LENGTH)
586            || headers.value_eq_ascii_case(http::header::TRANSFER_ENCODING, "chunked")
587    }
588
589    fn keepalive_timeout(resp: &Response) -> Option<Duration> {
590        let name = http::header::HeaderName::from_static("keep-alive");
591        let raw = resp.headers().get(name)?;
592        let value = raw.to_str().ok()?;
593        for part in value.split(',') {
594            let part = part.trim();
595            if let Some(rest) = part.strip_prefix("timeout=")
596                && let Ok(secs) = rest.trim().parse::<u64>()
597            {
598                return Some(Duration::from_secs(secs));
599            }
600        }
601        None
602    }
603
604    fn decompress(
605        resp: &mut Response,
606        policy: DecompressionPolicy,
607        max_body: usize,
608    ) -> Result<(), Error> {
609        let is_gzip = resp
610            .headers()
611            .value_eq_ascii_case(http::header::CONTENT_ENCODING, "gzip");
612        if !is_gzip || resp.body().is_empty() {
613            return Ok(());
614        }
615
616        let limit = max_body as u64;
617        let mut decoder = flate2::read::GzDecoder::new(resp.body()).take(limit + 1);
618        let mut decompressed = Vec::new();
619        match decoder.read_to_end(&mut decompressed) {
620            Ok(_) if decompressed.len() as u64 > limit => Err(Error::Parse(
621                "decompressed response body exceeds size limit".into(),
622            )),
623            Ok(_) => {
624                resp.set_body(decompressed);
625                resp.headers_mut().remove("content-encoding");
626                resp.headers_mut().remove("content-length");
627                Ok(())
628            }
629            Err(e) if policy == DecompressionPolicy::Strict => {
630                Err(Error::Parse(format!("gzip decompression failed: {e}")))
631            }
632            Err(_) => Ok(()),
633        }
634    }
635}
sark_client/connector/session.rs

sark_client/connector/
session.rs
