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

base_http_client/
driver.rs
