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

base_http_client/
driver.rs
