base_http_client/
driver.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4use crate::{
5    Connection, Endpoint, Error, Options, ResponseBody, connection,
6    outgoing::{Observer, OutgoingBody},
7};
8use base_http::{Request, Response, body::ReadableBody};
9use base_pool::{Acquisition, MultiplexBuildPermit, MultiplexPool, MultiplexUse};
10use base_scheduler::{
11    Cancellation, SocketClose, SocketConnect, SocketOpen, SocketStreamReading, SocketStreamWriting,
12};
13use event_listener::Event;
14use futures_channel::oneshot;
15use futures_util::stream::{FuturesUnordered, StreamExt};
16use pin_project_lite::pin_project;
17use std::{
18    collections::VecDeque,
19    future::{Future, poll_fn},
20    pin::{Pin, pin},
21    sync::{Arc, Mutex, MutexGuard, PoisonError, atomic::AtomicUsize},
22    task::Poll,
23};
24
25pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
26    mutex.lock().unwrap_or_else(PoisonError::into_inner)
27}
28pub(crate) type Work<'scope, T> = Pin<Box<dyn Future<Output = T> + Send + 'scope>>;
29type Factory<'scope> = Box<
30    dyn Fn(
31            MultiplexBuildPermit<Connection<'scope>>,
32            Cancellation,
33        ) -> Work<'scope, Result<(), Error>>
34        + Send
35        + Sync
36        + 'scope,
37>;
38pub(crate) type Reply = oneshot::Sender<Result<Response<ResponseBody>, Error>>;
39
40pub(crate) struct Operation<'scope> {
41    pub request: Request<OutgoingBody<'scope>>,
42    pub interim: Box<dyn Observer + 'scope>,
43    pub reply: Option<Reply>,
44    pub cancellation: Cancellation,
45    pub lifetime: Cancellation,
46    pub _pending: Option<Pending<'scope>>,
47}
48pub(crate) struct Assigned<'scope> {
49    pub operation: Operation<'scope>,
50    pub active: MultiplexUse<Connection<'scope>>,
51}
52struct State<'scope> {
53    requests: VecDeque<Operation<'scope>>,
54    work: VecDeque<Work<'scope, ()>>,
55    pending: usize,
56    closed: bool,
57    abandoned: bool,
58}
59pub(crate) struct Queue<'scope> {
60    state: Mutex<State<'scope>>,
61    pub changed: Event,
62    pub handles: AtomicUsize,
63    limit: usize,
64}
65impl<'scope> Queue<'scope> {
66    pub fn push(&self, operation: Operation<'scope>) -> Result<(), Error> {
67        let mut state = lock(&self.state);
68        if state.closed {
69            return Err(Error::Closed);
70        }
71        state.requests.push_back(operation);
72        drop(state);
73        self.changed.notify(usize::MAX);
74        Ok(())
75    }
76    pub(crate) fn push_work(&self, work: Work<'scope, ()>) -> Result<(), Error> {
77        let mut state = lock(&self.state);
78        if state.abandoned {
79            drop(state);
80            drop(work);
81            return Err(Error::Closed);
82        }
83        state.work.push_back(work);
84        drop(state);
85        self.changed.notify(usize::MAX);
86        Ok(())
87    }
88    pub fn begin_close(&self) {
89        lock(&self.state).closed = true;
90        self.changed.notify(usize::MAX);
91    }
92    fn abandon(&self) {
93        let discarded = {
94            let mut state = lock(&self.state);
95            state.closed = true;
96            state.abandoned = true;
97            (
98                std::mem::take(&mut state.requests),
99                std::mem::take(&mut state.work),
100            )
101        };
102        self.changed.notify(usize::MAX);
103        drop(discarded);
104    }
105}
106pub(crate) struct Pending<'scope>(Arc<Queue<'scope>>);
107impl<'scope> Pending<'scope> {
108    pub fn new(queue: Arc<Queue<'scope>>) -> Result<Self, Error> {
109        {
110            let mut state = lock(&queue.state);
111            if state.closed {
112                return Err(Error::Closed);
113            }
114            if state.pending >= queue.limit {
115                return Err(Error::Busy);
116            }
117            state.pending += 1;
118        }
119        Ok(Self(queue))
120    }
121}
122impl Drop for Pending<'_> {
123    fn drop(&mut self) {
124        let drained = {
125            let mut state = lock(&self.0.state);
126            state.pending -= 1;
127            state.closed && state.pending == 0
128        };
129        // Completion only changes the supervisor's work when a closing client
130        // has drained. Submission and deferred cleanup notify independently.
131        if drained {
132            self.0.changed.notify(usize::MAX);
133        }
134    }
135}
136
137pub(crate) struct Context<'scope> {
138    pool: MultiplexPool<Connection<'scope>>,
139    factory: Factory<'scope>,
140    pub(crate) cancellation: Cancellation,
141}
142impl<'scope> Context<'scope> {
143    pub(crate) async fn acquire(
144        &self,
145        cancellation: &Cancellation,
146    ) -> Result<MultiplexUse<Connection<'scope>>, Error> {
147        loop {
148            // Admission can change between reservation and activation. Either outcome without a use acquires again, which reports a closed pool.
149            let active = match self.pool.acquire_or_reserve(cancellation).await? {
150                Acquisition::Existing(reservation) => reservation.activate().ok(),
151                Acquisition::Build(build) => {
152                    (self.factory)(build, cancellation.clone()).await?;
153                    None
154                }
155            };
156            if let Some(active) = active {
157                return Ok(active);
158            }
159        }
160    }
161    async fn cleanup(&self) -> Result<(), Error> {
162        let mut failure = None;
163        while let Some(retired) = self
164            .pool
165            .next_retirement(&Cancellation::never())
166            .await
167            .expect("an uncancellable retirement wait cannot fail")
168        {
169            retired.stop();
170            if let Err(error) = retired.wait_closed().await {
171                failure.get_or_insert(error);
172            }
173            retired.complete();
174        }
175        failure.map_or(Ok(()), Err)
176    }
177}
178
179/// Explicitly supervised HTTP progress. Poll it concurrently with client calls
180/// and body consumers. Cancellation forces cleanup; await it to join all lane
181/// retirements and socket closes. Established plaintext direct requests run in
182/// their calling tasks. Direct request/response handles must be polled or dropped
183/// during cancellation so their codecs can return. The current scoped
184/// implementation polls all supervised connection work inside this one future;
185/// it does not distribute connections across task workers. Drop provides no
186/// asynchronous join.
187pub struct Driver<'scope> {
188    work: Work<'scope, Result<(), Error>>,
189    cancellation: Cancellation,
190    abandon: Box<dyn Fn() + Send + Sync + 'scope>,
191}
192impl Driver<'_> {
193    pub async fn run(mut self, cancellation: &Cancellation) -> Result<(), Error> {
194        let mut cancelled = pin!(cancellation.cancelled());
195        poll_fn(|cx| {
196            if !self.cancellation.is_cancelled() && cancelled.as_mut().poll(cx).is_ready() {
197                self.cancellation.cancel();
198            }
199            self.work.as_mut().poll(cx)
200        })
201        .await
202    }
203}
204impl Drop for Driver<'_> {
205    fn drop(&mut self) {
206        self.cancellation.cancel();
207        (self.abandon)();
208    }
209}
210
211pub(crate) struct ClientParts<'scope> {
212    pub queue: Arc<Queue<'scope>>,
213    pub pool: MultiplexPool<Connection<'scope>>,
214    pub driver: Driver<'scope>,
215    pub context: Arc<Context<'scope>>,
216}
217
218pub(crate) fn make<'scope, S>(
219    scheduler: &'scope S,
220    endpoint: Endpoint,
221    options: Options,
222) -> ClientParts<'scope>
223where
224    S: SocketOpen
225        + SocketConnect
226        + SocketClose
227        + SocketStreamReading
228        + SocketStreamWriting
229        + 'scope,
230    S::Socket: 'static,
231{
232    let pool =
233        MultiplexPool::new(options.pool.clone()).expect("Client::new validated the pool options");
234    let queue = Arc::new(Queue {
235        state: Mutex::new(State {
236            requests: VecDeque::new(),
237            work: VecDeque::new(),
238            pending: 0,
239            closed: false,
240            abandoned: false,
241        }),
242        changed: Event::new(),
243        handles: AtomicUsize::new(1),
244        limit: options.max_requests,
245    });
246    let cancellation = Cancellation::new();
247    let weak_queue = Arc::downgrade(&queue);
248    let connection_lifetime = cancellation.clone();
249    let factory = Box::new(move |build, cancelled: Cancellation| {
250        let options = options.clone();
251        let endpoint = endpoint.clone();
252        let lifetime = connection_lifetime.clone();
253        // Acquirers hold the queue: the driver's assignment jobs and client handles both retain it.
254        let queue = weak_queue.upgrade().expect("an acquirer retains the queue");
255        Box::pin(async move {
256            let (ready, established) = oneshot::channel();
257            queue.push_work(Box::pin(connection::run(
258                scheduler,
259                endpoint,
260                options,
261                build,
262                cancelled,
263                lifetime,
264                ready,
265                Arc::downgrade(&queue),
266            )))?;
267            // Only an abandoned driver drops construction before it reports readiness.
268            established.await.unwrap_or(Err(Error::Closed))
269        }) as Work<'scope, _>
270    });
271    let context = Arc::new(Context {
272        pool: pool.clone(),
273        factory,
274        cancellation: cancellation.clone(),
275    });
276    let abandoned = queue.clone();
277    let abandoned_pool = pool.clone();
278    let driver = Driver {
279        work: Box::pin(run(queue.clone(), context.clone())),
280        cancellation,
281        abandon: Box::new(move || {
282            abandoned.abandon();
283            abandoned_pool.begin_close();
284        }),
285    };
286    ClientParts {
287        queue,
288        pool,
289        driver,
290        context,
291    }
292}
293
294async fn assign<'scope>(mut operation: Operation<'scope>, context: Arc<Context<'scope>>) {
295    let result = context.acquire(&operation.cancellation).await;
296    match result {
297        Ok(active) => {
298            let connection = active.clone_handle();
299            connection.submit(Assigned { operation, active });
300        }
301        Err(error) => {
302            operation.request.body.close(None);
303            let reply = operation
304                .reply
305                .take()
306                .expect("an unassigned operation retains its reply");
307            let _ = reply.send(Err(error));
308        }
309    }
310}
311
312async fn run<'scope>(
313    queue: Arc<Queue<'scope>>,
314    context: Arc<Context<'scope>>,
315) -> Result<(), Error> {
316    let mut jobs: FuturesUnordered<Work<'scope, ()>> = FuturesUnordered::new();
317    let cleanup_context = context.clone();
318    let mut cleanup = pin!(cleanup_context.cleanup());
319    let mut cancellation = pin!(context.cancellation.cancelled());
320    let mut stopping = false;
321    let mut cleaned = None;
322    loop {
323        event_listener::listener!(queue.changed => listener);
324        let (requests, work) = {
325            let mut state = lock(&queue.state);
326            (
327                std::mem::take(&mut state.requests),
328                std::mem::take(&mut state.work),
329            )
330        };
331        for operation in requests {
332            let context = context.clone();
333            jobs.push(Box::pin(async move {
334                let local = operation.cancellation.clone();
335                let lifetime = operation.lifetime.clone();
336                let driver = context.cancellation.clone();
337                linked(assign(operation, context), &local, &lifetime, &driver).await;
338            }));
339        }
340        jobs.extend(work);
341        let close = {
342            let state = lock(&queue.state);
343            state.closed && state.pending == 0
344        };
345        if stopping || close {
346            context.pool.begin_close();
347        }
348        let mut listener = pin!(listener);
349        let finished = poll_fn(|cx| {
350            if !stopping && cancellation.as_mut().poll(cx).is_ready() {
351                stopping = true;
352                queue.begin_close();
353                context.pool.begin_close();
354            }
355            if cleaned.is_none() {
356                if let Poll::Ready(result) = cleanup.as_mut().poll(cx) {
357                    cleaned = Some(result);
358                }
359            }
360            if cleaned.is_some() && jobs.is_empty() && lock(&queue.state).work.is_empty() {
361                return Poll::Ready(cleaned.take());
362            }
363            if let Poll::Ready(Some(())) = jobs.poll_next_unpin(cx) {
364                return Poll::Ready(None);
365            }
366            if listener.as_mut().poll(cx).is_ready() {
367                return Poll::Ready(None);
368            }
369            Poll::Pending
370        })
371        .await;
372        if let Some(result) = finished {
373            return result;
374        }
375    }
376}
377
378pin_project! {
379    struct Checked<'a, F, C> {
380        #[pin]
381        work: F,
382        #[pin]
383        cancelled: C,
384        cancellation: &'a Cancellation,
385    }
386}
387impl<T, F, C> Future for Checked<'_, F, C>
388where
389    F: Future<Output = Result<T, Error>>,
390    C: Future<Output = ()>,
391{
392    type Output = Result<T, Error>;
393    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
394        let this = self.project();
395        if this.cancellation.is_cancelled() {
396            return Poll::Ready(Err(Error::Cancelled));
397        }
398        match this.work.poll(cx) {
399            ready @ Poll::Ready(_) => ready,
400            Poll::Pending => {
401                // Subscribe only when work actually waits. Polling cancelled
402                // after Pending also closes the race with the initial check.
403                if this.cancelled.poll(cx).is_ready() {
404                    Poll::Ready(Err(Error::Cancelled))
405                } else {
406                    Poll::Pending
407                }
408            }
409        }
410    }
411}
412pub(crate) fn checked<T>(
413    work: impl Future<Output = Result<T, Error>>,
414    cancellation: &Cancellation,
415) -> impl Future<Output = Result<T, Error>> {
416    // Store the nested future where it will be polled. An async wrapper with
417    // pin!(work) needs another large state slot and moves work on first poll.
418    Checked {
419        work,
420        cancelled: cancellation.cancelled(),
421        cancellation,
422    }
423}
424// Propagate lifetime cancellation without dropping owned construction/cleanup.
425pub(crate) async fn linked<T>(
426    work: impl Future<Output = T>,
427    cancellation: &Cancellation,
428    lifetime: &Cancellation,
429    driver: &Cancellation,
430) -> T {
431    let mut work = pin!(work);
432    let mut a = pin!(lifetime.cancelled());
433    let mut b = pin!(driver.cancelled());
434    poll_fn(|cx| {
435        if !cancellation.is_cancelled()
436            && (a.as_mut().poll(cx).is_ready() || b.as_mut().poll(cx).is_ready())
437        {
438            cancellation.cancel();
439        }
440        work.as_mut().poll(cx)
441    })
442    .await
443}
444pub(crate) struct CancelOnDrop(pub Cancellation, pub bool);
445impl Drop for CancelOnDrop {
446    fn drop(&mut self) {
447        if self.1 {
448            self.0.cancel();
449        }
450    }
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use crate::{Endpoint, Options};
457    use base_http::{Method, body::EmptyBody};
458    use base_scheduler_direct::DirectScheduler;
459    use std::{
460        sync::atomic::Ordering,
461        task::{Context, Wake, Waker},
462    };
463
464    #[derive(Default)]
465    struct Wakes(AtomicUsize);
466
467    // Work which records whether it was polled.
468    struct Probe<'a>(&'a std::cell::Cell<bool>);
469    impl Future for Probe<'_> {
470        type Output = Result<(), Error>;
471        fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
472            self.0.set(true);
473            Poll::Ready(Ok(()))
474        }
475    }
476    impl Wake for Wakes {
477        fn wake(self: Arc<Self>) {
478            self.0.fetch_add(1, Ordering::Relaxed);
479        }
480    }
481
482    #[test]
483    fn checked_prioritizes_existing_cancellation_and_wakes_blocked_work() {
484        use futures_lite::future::{block_on, poll_once};
485        use std::cell::Cell;
486        let polled = Cell::new(false);
487        block_on(checked(Probe(&polled), &Cancellation::new())).unwrap();
488        assert!(polled.get());
489
490        let cancel = Cancellation::new();
491        cancel.cancel();
492        let polled = Cell::new(false);
493        let error = block_on(checked(Probe(&polled), &cancel)).unwrap_err();
494        assert_eq!(error.to_string(), Error::Cancelled.to_string());
495        assert!(!polled.get());
496
497        let cancel = Cancellation::new();
498        let wakes = Arc::new(Wakes::default());
499        let waker = Waker::from(wakes.clone());
500        let mut cx = Context::from_waker(&waker);
501        let mut work = pin!(checked(
502            std::future::pending::<Result<(), Error>>(),
503            &cancel
504        ));
505        assert!(work.as_mut().poll(&mut cx).is_pending());
506        cancel.cancel();
507        assert!(wakes.0.load(Ordering::Relaxed) > 0);
508        assert_eq!(
509            format!("{:?}", work.as_mut().poll(&mut cx)),
510            "Ready(Err(Cancelled))"
511        );
512
513        // Cancellation occurring during a Pending poll must be noticed in
514        // that poll, even before a listener has ever been registered.
515        let cancel = Cancellation::new();
516        let pending = std::future::poll_fn(|_| {
517            cancel.cancel();
518            Poll::Pending::<Result<(), Error>>
519        });
520        let mut work = pin!(checked(pending, &cancel));
521        assert_eq!(
522            format!("{:?}", block_on(poll_once(work.as_mut()))),
523            "Some(Err(Cancelled))"
524        );
525    }
526
527    #[test]
528    fn pending_completion_wakes_supervisor_only_when_closed_and_drained() {
529        let queue = Arc::new(Queue {
530            state: Mutex::new(State {
531                requests: VecDeque::new(),
532                work: VecDeque::new(),
533                pending: 0,
534                closed: false,
535                abandoned: false,
536            }),
537            changed: Event::new(),
538            handles: AtomicUsize::new(1),
539            limit: 2,
540        });
541        let wakes = Arc::new(Wakes::default());
542        let waker = Waker::from(wakes.clone());
543        let mut cx = Context::from_waker(&waker);
544        let mut changed = pin!(queue.changed.listen());
545        assert!(changed.as_mut().poll(&mut cx).is_pending());
546
547        let first = Pending::new(queue.clone()).unwrap();
548        let second = Pending::new(queue.clone()).unwrap();
549        let error = Pending::new(queue.clone()).err().unwrap();
550        assert_eq!(error.to_string(), Error::Busy.to_string());
551        drop(first);
552        assert_eq!(wakes.0.load(Ordering::Relaxed), 0);
553        // Admission is a synchronous check; completion restores it without
554        // requiring any supervisor work.
555        let third = Pending::new(queue.clone()).unwrap();
556        queue.begin_close();
557        assert!(changed.as_mut().poll(&mut cx).is_ready());
558        let mut changed = pin!(queue.changed.listen());
559        assert!(changed.as_mut().poll(&mut cx).is_pending());
560        wakes.0.store(0, Ordering::Relaxed);
561        let error = Pending::new(queue.clone()).err().unwrap();
562        assert_eq!(error.to_string(), Error::Closed.to_string());
563        drop(second);
564        assert_eq!(wakes.0.load(Ordering::Relaxed), 0);
565        assert!(changed.as_mut().poll(&mut cx).is_pending());
566        drop(third);
567        assert_eq!(wakes.0.load(Ordering::Relaxed), 1);
568        assert!(changed.as_mut().poll(&mut cx).is_ready());
569        assert_eq!(lock(&queue.state).pending, 0);
570    }
571
572    fn parts(scheduler: &DirectScheduler) -> ClientParts<'_> {
573        let endpoint = Endpoint::new("127.0.0.1:1".parse().unwrap());
574        make(scheduler, endpoint, Options::default())
575    }
576
577    #[test]
578    fn requests_pushed_after_close_are_rejected() {
579        let scheduler = DirectScheduler::new(Default::default()).unwrap();
580        let parts = parts(&scheduler);
581        parts.queue.begin_close();
582        let operation = Operation {
583            request: Request::new(
584                Method::GET,
585                "/",
586                OutgoingBody::new(EmptyBody, Cancellation::never()),
587            ),
588            interim: Box::new(crate::outgoing::tests::ignore),
589            reply: None,
590            cancellation: Cancellation::new(),
591            lifetime: Cancellation::never(),
592            _pending: None,
593        };
594        let error = parts.queue.push(operation).unwrap_err();
595        assert_eq!(error.to_string(), Error::Closed.to_string());
596        assert!(lock(&parts.queue.state).requests.is_empty());
597    }
598
599    #[test]
600    fn construction_after_the_queue_is_abandoned_releases_its_build() {
601        let scheduler = DirectScheduler::new(Default::default()).unwrap();
602        let parts = parts(&scheduler);
603        // Abandonment closes the queue before the pool; a build reserved between the two cannot start.
604        parts.queue.abandon();
605        let error = futures_lite::future::block_on(parts.context.acquire(&Cancellation::never()))
606            .err()
607            .unwrap();
608        assert_eq!(error.to_string(), Error::Closed.to_string());
609        assert_eq!(parts.pool.stats().total(), 0);
610    }
611}
base_http_client/driver.rs

base_http_client/
driver.rs
