rama_net/client/
connector.rs

1use std::{future::Future, net::SocketAddr};
2
3use crate::address::HostWithPort;
4
5use rama_core::{
6    error::{BoxError, BoxErrorExt as _},
7    extensions::Extensions,
8    futures::{
9        Stream, StreamExt as _,
10        stream::{BoxStream, FuturesUnordered},
11    },
12};
13use rama_macros::Extension;
14
15#[derive(Debug, Clone, PartialEq, Eq, Hash, Extension)]
16#[extension(tags(net))]
17/// Target [`HostWithPort`] which if found in extensions
18/// is to be used by a connector such as a TCPConnector instead
19/// of the requested address, unless a proxy is requested in
20/// which case a proxy is to be used instead.
21pub struct ConnectorTarget(pub HostWithPort);
22
23/// A lazily-resolved source of connection target [`SocketAddr`]esses.
24///
25/// This is the abstraction boundary between resolving a target and
26/// connecting to it. An upstream connector (e.g. the `rama-dns` address
27/// resolver) stamps a [`ConnectorTargetStream`] into the [`Extensions`], and a
28/// transport connector consumes it: dialing (and optionally racing) the
29/// yielded addresses. The transport stays resolver-agnostic: it only ever sees
30/// a stream of [`SocketAddr`]s, never a DNS resolver.
31pub trait AddressCandidates: Send + Sync + 'static {
32    /// Stream the candidate [`SocketAddr`]esses, in the order they should be
33    /// attempted. The given [`Extensions`] carry per-request resolve config.
34    fn stream<'a>(
35        &'a self,
36        extensions: &'a Extensions,
37    ) -> BoxStream<'a, Result<SocketAddr, BoxError>>;
38}
39
40#[derive(Extension)]
41#[extension(tags(net))]
42/// [`Extensions`] carrier for an [`AddressCandidates`] source.
43pub struct ConnectorTargetStream(pub Box<dyn AddressCandidates>);
44
45impl ConnectorTargetStream {
46    /// Wrap an [`AddressCandidates`] implementor.
47    #[must_use]
48    pub fn new(candidates: impl AddressCandidates) -> Self {
49        Self(Box::new(candidates))
50    }
51
52    /// Stream the candidate addresses (see [`AddressCandidates::stream`]).
53    pub fn stream<'a>(
54        &'a self,
55        extensions: &'a Extensions,
56    ) -> BoxStream<'a, Result<SocketAddr, BoxError>> {
57        self.0.stream(extensions)
58    }
59}
60
61impl core::fmt::Debug for ConnectorTargetStream {
62    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
63        f.debug_struct("ConnectorTargetStream")
64            .finish_non_exhaustive()
65    }
66}
67
68/// Race connection attempts over a stream of candidate [`SocketAddr`]esses.
69///
70/// Pulls candidates from `candidates` in order (e.g. happy-eyeballs order from a
71/// [`ConnectorTargetStream`]), keeps up to `max_in_flight` dials racing
72/// concurrently via `dial`, and returns the first successful connection together
73/// with the address it connected to. If every candidate fails (to resolve or to
74/// connect), the last error is returned.
75pub async fn race_connect<S, C, F, Fut>(
76    candidates: S,
77    max_in_flight: usize,
78    dial: F,
79) -> Result<(SocketAddr, C), BoxError>
80where
81    S: Stream<Item = Result<SocketAddr, BoxError>>,
82    F: Fn(SocketAddr) -> Fut + Sync,
83    Fut: Future<Output = Result<C, BoxError>> + Send,
84    C: Send,
85{
86    let max_in_flight = max_in_flight.max(1);
87
88    let dial = &dial;
89    let mut candidates = std::pin::pin!(candidates);
90    let mut in_flight = FuturesUnordered::new();
91    let mut candidates_done = false;
92    let mut last_err: Option<BoxError> = None;
93
94    enum Event<C> {
95        Candidate(Option<Result<SocketAddr, BoxError>>),
96        Dialed(Option<(SocketAddr, Result<C, BoxError>)>),
97    }
98
99    loop {
100        if candidates_done && in_flight.is_empty() {
101            break;
102        }
103
104        let event = if !candidates_done && in_flight.len() < max_in_flight {
105            if in_flight.is_empty() {
106                Event::Candidate(candidates.next().await)
107            } else {
108                tokio::select! {
109                    candidate = candidates.next() => Event::Candidate(candidate),
110                    dialed = in_flight.next() => Event::Dialed(dialed),
111                }
112            }
113        } else {
114            Event::Dialed(in_flight.next().await)
115        };
116
117        match event {
118            Event::Candidate(Some(Ok(addr))) => {
119                in_flight.push(async move { (addr, dial(addr).await) });
120            }
121            Event::Candidate(Some(Err(err))) => last_err = Some(err),
122            Event::Candidate(None) => candidates_done = true,
123            Event::Dialed(Some((addr, Ok(conn)))) => return Ok((addr, conn)),
124            Event::Dialed(Some((_addr, Err(err)))) => last_err = Some(err),
125            Event::Dialed(None) => {}
126        }
127    }
128
129    Err(last_err
130        .unwrap_or_else(|| BoxError::from_static_str("race_connect: no connection candidates")))
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    use rama_core::futures::stream;
138
139    fn addr(port: u16) -> SocketAddr {
140        SocketAddr::from(([127, 0, 0, 1], port))
141    }
142
143    #[tokio::test]
144    async fn race_connect_returns_a_success_skipping_failures() {
145        let candidates = stream::iter([addr(1), addr(2), addr(3)].map(Ok::<_, BoxError>));
146
147        let (won, conn) = race_connect(candidates, 3, |a: SocketAddr| async move {
148            if a.port() == 1 {
149                Err(BoxError::from_static_str("refused"))
150            } else {
151                Ok::<_, BoxError>(a.port())
152            }
153        })
154        .await
155        .unwrap();
156        assert_ne!(won.port(), 1);
157        assert_eq!(conn, won.port());
158    }
159
160    #[tokio::test]
161    async fn race_connect_all_failures_returns_last_error() {
162        let candidates = stream::iter([addr(1), addr(2)].map(Ok::<_, BoxError>));
163        let result = race_connect(candidates, 3, |_a| async move {
164            Err::<u16, BoxError>(BoxError::from_static_str("always refused"))
165        })
166        .await;
167        result.unwrap_err();
168    }
169}
rama_net/client/connector.rs

rama_net/client/
connector.rs
