sark_client/connector/
session.rs

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

sark_client/connector/
session.rs
