zebra_network/peer_set/
service.rs

1//! External ownership of the network's buffered service and background tasks.
2
3use std::{
4    fmt,
5    future::Future,
6    pin::Pin,
7    sync::Arc,
8    task::{Context, Poll},
9};
10
11use tokio::task::AbortHandle;
12use tower::{buffer::Buffer, util::BoxService, Service};
13
14use crate::{BoxError, Request, Response};
15
16/// A cloneable, buffered service for outbound network requests.
17///
18/// The network stays alive while an owning handle or an unfinished [`NetworkResponseFuture`]
19/// exists. Releasing the last owner aborts the worker and background tasks, releasing the
20/// listener and peer clients while the runtime is still running. Internal crawler and polling
21/// tasks use the underlying buffer without sharing this shutdown ownership.
22#[derive(Clone)]
23pub struct NetworkService {
24    inner: Buffer<BoxService<Request, Response, BoxError>, Request>,
25    tasks: Arc<NetworkTasks>,
26}
27
28// The boxed inner service is not `Debug`.
29impl fmt::Debug for NetworkService {
30    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
31        formatter
32            .debug_struct("NetworkService")
33            .field("tasks", &self.tasks)
34            .finish_non_exhaustive()
35    }
36}
37
38/// An outbound response that keeps its [`NetworkService`] alive until completion or cancellation.
39///
40/// This also preserves ownership when [`tower::ServiceExt::oneshot`] consumes the last service
41/// handle. Background handles returned by [`NetworkService::background_service`] do not create
42/// owning response futures.
43#[pin_project]
44pub struct NetworkResponseFuture {
45    #[pin]
46    inner: <Buffer<BoxService<Request, Response, BoxError>, Request> as Service<Request>>::Future,
47    tasks: Option<Arc<NetworkTasks>>,
48}
49
50impl Future for NetworkResponseFuture {
51    type Output = Result<Response, BoxError>;
52
53    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54        let this = self.project();
55        let response = this.inner.poll(cx);
56        if response.is_ready() {
57            drop(this.tasks.take());
58        }
59        response
60    }
61}
62
63#[derive(Debug)]
64struct NetworkTasks(Vec<AbortHandle>);
65
66impl Drop for NetworkTasks {
67    fn drop(&mut self) {
68        for task in &self.0 {
69            task.abort();
70        }
71    }
72}
73
74impl NetworkService {
75    /// Takes ownership of the existing buffer and its tasks without adding another buffer.
76    pub(crate) fn new(
77        inner: Buffer<BoxService<Request, Response, BoxError>, Request>,
78        tasks: Vec<AbortHandle>,
79    ) -> Self {
80        Self {
81            inner,
82            tasks: Arc::new(NetworkTasks(tasks)),
83        }
84    }
85
86    /// Returns a buffered handle that does not keep the network running.
87    ///
88    /// Use this handle in services retained by the network, including inbound handlers and
89    /// their downloaders, to avoid a shutdown ownership cycle. Keep a [`NetworkService`]
90    /// handle in an external owner for as long as the network is needed.
91    ///
92    /// After the last owning handle and unfinished response are released, requests through
93    /// this handle fail even if background handles remain. Neither this handle nor its response
94    /// futures keep the network alive. This does not add another buffer or worker.
95    pub fn background_service(&self) -> Buffer<BoxService<Request, Response, BoxError>, Request> {
96        self.inner.clone()
97    }
98
99    /// Registers tasks during initialization, before sharing external handles.
100    pub(crate) fn extend_tasks(&mut self, tasks: impl IntoIterator<Item = AbortHandle>) {
101        Arc::get_mut(&mut self.tasks)
102            .expect("network tasks are registered before external handles are shared")
103            .0
104            .extend(tasks);
105    }
106}
107
108impl Service<Request> for NetworkService {
109    type Response = Response;
110    type Error = BoxError;
111    type Future = NetworkResponseFuture;
112
113    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
114        self.inner.poll_ready(cx)
115    }
116
117    fn call(&mut self, request: Request) -> Self::Future {
118        NetworkResponseFuture {
119            inner: self.inner.call(request),
120            tasks: Some(self.tasks.clone()),
121        }
122    }
123}
zebra_network/peer_set/service.rs

zebra_network/peer_set/
service.rs
