/build/source/nativelink-util/src/connection_manager.rs

Source
// Copyright 2024 Trace Machina, Inc. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::pin::Pin;
use core::task::{Context, Poll};
use core::time::Duration;
use std::collections::VecDeque;
use std::sync::Arc;

use futures::Future;
use futures::stream::{FuturesUnordered, StreamExt, unfold};
use nativelink_config::stores::Retry;
use nativelink_error::{Code, Error, make_err};
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
use tonic::transport::{Channel, Endpoint, channel};
use tracing::{debug, error, info, warn};

use crate::background_spawn;
use crate::metrics::record_connection_acquired;
use crate::retry::{self, Retrier, RetryResult};

/// A helper utility that enables management of a suite of connections to an
/// upstream gRPC endpoint using Tonic.
#[derive(Debug)]
pub struct ConnectionManager {
    // The channel to request connections from the worker.
    worker_tx: mpsc::Sender<(String, oneshot::Sender<Connection>)>,
}

/// The index into `ConnectionManagerWorker::endpoints`.
type EndpointIndex = usize;
/// The identifier for a given connection to a given Endpoint, used to identify
/// when a particular connection has failed or becomes available.
type ConnectionIndex = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct ChannelIdentifier {
    /// The index into `ConnectionManagerWorker::endpoints` that established this
    /// Channel.
    endpoint_index: EndpointIndex,
    /// A unique identifier for this particular connection to the Endpoint.
    connection_index: ConnectionIndex,
}

/// The requests that can be made from a Connection to the
/// `ConnectionManagerWorker` such as informing it that it's been dropped or that
/// an error occurred.
enum ConnectionRequest {
    /// Notify that a Connection was dropped, if it was dropped while the
    /// connection was still pending, then return the pending Channel to be
    /// added back to the available channels.
    Dropped(Option<EstablishedChannel>),
    /// Notify that a Connection was established, return the Channel to the
    /// available channels.
    Connected(EstablishedChannel),
    /// Notify that there was a transport error on the given Channel, the bool
    /// specifies whether the connection was in the process of being established
    /// or not (i.e. whether it's been returned to available channels yet).
    Error((ChannelIdentifier, bool)),
}

/// The result of a Future that connects to a given Endpoint.  This is a tuple
/// of the index into the `ConnectionManagerWorker::endpoints` that this
/// connection is for, the iteration of the connection and the result of the
/// connection itself.
type IndexedChannel = Result<EstablishedChannel, (ChannelIdentifier, Error)>;

/// A channel that has been established to an endpoint with some metadata around
/// it to allow identification of the Channel if it errors in order to correctly
/// remove it.
#[derive(Debug, Clone)]
struct EstablishedChannel {
    /// The Channel itself that the meta data relates to.
    channel: Channel,
    /// The identifier of the channel in the worker.
    identifier: ChannelIdentifier,
}

/// The context of the worker used to manage all of the connections.  This
/// handles reconnecting to endpoints on errors and multiple connections to a
/// given endpoint.
struct ConnectionManagerWorker {
    /// The endpoints to establish Channels and the identifier of the last
    /// connection attempt to that endpoint.
    endpoints: Vec<(ConnectionIndex, Endpoint)>,
    /// The channel used to communicate between a Connection and the worker.
    connection_tx: mpsc::UnboundedSender<ConnectionRequest>,
    /// Gates the maximum number of in-flight `Connection` objects.
    /// Was an explicit `usize` counter; now an `Arc<Semaphore>` so the
    /// `OwnedSemaphorePermit` held by each `Connection` releases on
    /// drop (RAII), instead of relying on a `ConnectionRequest::Dropped`
    /// round-trip that could be lost on tonic transport errors or task
    /// aborts.
    available_connections: Arc<Semaphore>,

    /// Whether `available_connections` reflects a configured limit. When the
    /// pool is unlimited its permit count is a sentinel, not a measurement.
    bounded_connections: bool,
    /// Channels that are currently being connected.
    connecting_channels: FuturesUnordered<Pin<Box<dyn Future<Output = IndexedChannel> + Send>>>,
    /// Connected channels that are available for use.
    available_channels: VecDeque<EstablishedChannel>,
    /// Requests for a Channel when available - (reason, request)
    waiting_connections: VecDeque<(String, oneshot::Sender<Connection>)>,
    /// The retry configuration for connecting to an Endpoint, on failure will
    /// restart the retrier after a 1 second delay.
    retrier: Retrier,
}

/// The maximum number of queued requests to obtain a connection from the
/// worker before applying back pressure to the requestor.  It makes sense to
/// keep this small since it has to wait for a response anyway.
const WORKER_BACKLOG: usize = 8;

impl ConnectionManager {
    /// Create a connection manager that creates a balance list between a given
    /// set of Endpoints.  This will restrict the number of concurrent requests
    /// and automatically re-connect upon transport error.
    pub fn new(
        endpoints: impl IntoIterator<Item = Endpoint>,
        mut connections_per_endpoint: usize,
        mut max_concurrent_requests: usize,
        retry: Retry,
        jitter_fn: retry::JitterFn,
    ) -> Self {
        let (worker_tx, worker_rx) = mpsc::channel(WORKER_BACKLOG);
        // The connection messages always come from sync contexts (e.g. drop)
        // and therefore, we'd end up spawning for them if this was bounded
        // which defeats the object since there would be no backpressure
        // applied. Therefore it makes sense for this to be unbounded.
        let (connection_tx, connection_rx) = mpsc::unbounded_channel();
        let endpoints = endpoints
            .into_iter()
            .map(|endpoint| (0, endpoint))
            .collect();

        // Zero means unlimited, which becomes a semaphore holding
        // Semaphore::MAX_PERMITS. That is a sentinel, not a free-slot count,
        // so remember which case this is and report no figure when unbounded.
        let bounded_connections = max_concurrent_requests != 0;
        if max_concurrent_requests == 0 {
            max_concurrent_requests = Semaphore::MAX_PERMITS;
        } else {
            max_concurrent_requests = max_concurrent_requests.min(Semaphore::MAX_PERMITS);
        }
        if connections_per_endpoint == 0 {
            connections_per_endpoint = 1;
        }3
        let worker = ConnectionManagerWorker {
            endpoints,
            available_connections: Arc::new(Semaphore::new(max_concurrent_requests)),
            bounded_connections,
            connection_tx,
            connecting_channels: FuturesUnordered::new(),
            available_channels: VecDeque::new(),
            waiting_connections: VecDeque::new(),
            retrier: Retrier::new(
                Arc::new(|duration| Box::pin0(tokio::time::sleep0(duration0))),
                jitter_fn,
                retry,
            )
            .with_client("connection"),
        };
        background_spawn!("connection_manager_worker_spawn", async move {47
            worker
                .service_requests(connections_per_endpoint, worker_rx, connection_rx)
                .await;
        });
        Self { worker_tx }
    }

    /// Get a Connection that can be used as a `tonic::Channel`, except it
    /// performs some additional counting to reconnect on error and restrict
    /// the number of concurrent connections.
    pub async fn connection(&self, reason: String) -> Result<Connection, Error>0 {
        let (tx, rx) = oneshot::channel();
        self.worker_tx
            .send((reason, tx))
            .await
            .map_err(|err| make_err!0(Code::Unavailable0, "Requesting a new connection: {err:?}"))?0;
        rx.await
            .map_err(|err| make_err!0(Code::Unavailable0, "Waiting for a new connection: {err:?}"))
    }
}

impl ConnectionManagerWorker {
    async fn service_requests(
        mut self,
        connections_per_endpoint: usize,
        mut worker_rx: mpsc::Receiver<(String, oneshot::Sender<Connection>)>,
        mut connection_rx: mpsc::UnboundedReceiver<ConnectionRequest>,
    ) {
        // Make the initial set of connections, connection failures will be
        // handled in the same way as future transport failures, so no need to
        // do anything special.
        for endpoint_index in 0..self.endpoints.len() {
            for _ in 0..connections_per_endpoint47 {
                self.connect_endpoint(endpoint_index, None);
            }
        }

        // The main worker loop, when select resolves one of its arms the other
        // ones are cancelled, therefore it's important that they maintain no
        // state while `await`-ing.  This is enforced through the use of
        // non-async functions to do all of the work.
        loop {
            tokio::select! {
                request278 = worker_rx.recv() => {
                    let Some((reason271, request271)) = request else {
                        // The ConnectionManager was dropped, shut down the
                        // worker.
                        break;
                    };
                    self.handle_worker(reason, request);
                }
                maybe_request298 = connection_rx.recv() => {
                    if let Some(request) = maybe_request {
                        self.handle_connection(request);
                    }0
                }
                maybe_connection_result50 = self.connect_next() => {
                    if let Some(connection_result) = maybe_connection_result {
                        self.handle_connected(connection_result);
                    }0
                }
            }
        }
    }

    async fn connect_next(&mut self) -> Option<IndexedChannel>0 {
        if self.connecting_channels.is_empty() {
            // Make this Future never resolve, we will get cancelled by the
            // select if there's some change in state to `self` and can re-enter
            // and evaluate `connecting_channels` again.
            futures::future::pending::<()>().await;
        }
        self.connecting_channels.next().await
    }

    // This must never be made async otherwise the select may cancel it.
    fn handle_connected(&mut self, connection_result: IndexedChannel) {
        match connection_result {
            Ok(established_channel) => {
                self.available_channels.push_back(established_channel);
                self.maybe_available_connection();
            }
            // When the retrier runs out of attempts start again from the
            // beginning of the retry period.  Never want to be in a
            // situation where we give up on an Endpoint forever.
            Err((identifier, _)) => {
                self.connect_endpoint(identifier.endpoint_index, Some(identifier.connection_index));
            }
        }
    }

    fn connect_endpoint(&mut self, endpoint_index: usize, connection_index: Option<usize>) {
        let Some((current_connection_index, endpoint)) = self.endpoints.get_mut(endpoint_index)
        else {
            // Unknown endpoint, this should never happen.
            error!(?endpoint_index, "Connection to unknown endpoint requested");
            return;
        };
        let is_backoff = connection_index.is_some();
        let connection_index = connection_index.unwrap_or_else(|| {
            *current_connection_index += 1;
            *current_connection_index
        });
        if is_backoff {
            warn!(
                ?connection_index,
                endpoint = ?endpoint.uri(),
                "Connection failed, reconnecting"
            );
        } else {
            info!(
                ?connection_index,
                endpoint = ?endpoint.uri(),
                "Creating new connection"
            );
        }
        let identifier = ChannelIdentifier {
            endpoint_index,
            connection_index,
        };
        let connection_stream = unfold(endpoint.clone(), move |endpoint| async move {
            let result50 = endpoint.connect().await.map_err50(|err| {0
                make_err!(
                    Code::Unavailable,
                    "Failed to connect to {:?}: {err:?}",
                    endpoint.uri()
                )
            });
            Some((
                result.map_or_else(RetryResult::Retry, RetryResult::Ok),
                endpoint,
            ))
        });
        let retrier = self.retrier.clone();
        self.connecting_channels.push(Box::pin(async move {
            if is_backoff {
                // Just in case the retry config is 0, then we need to
                // introduce some delay so we aren't in a hard loop.
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
            retrier.retry(connection_stream).await.map_or_else50(
                |err| Err((identifier, err)),
                |channel| {
                    Ok(EstablishedChannel {
                        channel,
                        identifier,
                    })
                },
            )
        }));
    }

    /// Free slots, or `None` when the pool is unlimited and the permit count
    /// is a sentinel rather than a measurement.
    fn free_connection_slots(&self) -> Option<usize> {
        self.bounded_connections
            .then(|| self.available_connections215.available_permits215())
    }

    // This must never be made async otherwise the select may cancel it.
    fn handle_worker(&mut self, reason: String, tx: oneshot::Sender<Connection>) {
        let maybe_permit = self.available_connections.clone().try_acquire_owned().ok();
        if let Some(permit270) = maybe_permit
            && let Some(channel157) = self.available_channels.pop_front()
        {
            debug!(reason, "ConnectionManager: request running");
            record_connection_acquired("grpc", self.free_connection_slots(), false);
            self.provide_channel(channel, tx, permit);
        } else {
            debug!(
                available_permits = self.available_connections.available_permits(),
                available_channels = self.available_channels.len(),
                waiting_connections = self.waiting_connections.len(),
                reason,
                "ConnectionManager: no connection available, request queued",
            );
            record_connection_acquired("grpc", self.free_connection_slots(), true);
            self.waiting_connections.push_back((reason, tx));
        }
    }

    fn provide_channel(
        &self,
        channel: EstablishedChannel,
        tx: oneshot::Sender<Connection>,
        permit: OwnedSemaphorePermit,
    ) {
        drop(tx.send(Connection {
            tx: self.connection_tx.clone(),
            pending_channel: Some(channel.channel.clone()),
            channel,
            _permit: permit,
        }));
    }

    fn maybe_available_connection(&mut self) {
        while !self.waiting_connections.is_empty() && !self.available_channels.is_empty()123 {
            let Some(permit) = self.available_connections.clone().try_acquire_owned().ok() else {
                break;
            };
            let Some(channel) = self.available_channels.pop_front() else {
                drop(permit);
                break;
            };
            let Some((reason, tx)) = self.waiting_connections.pop_front() else {
                self.available_channels.push_front(channel);
                drop(permit);
                break;
            };
            debug!(reason, "ConnectionManager: channel available, running");
            self.provide_channel(channel, tx, permit);
        }
    }

    // This must never be made async otherwise the select may cancel it.
    fn handle_connection(&mut self, request: ConnectionRequest) {
        match request {
            ConnectionRequest::Dropped(maybe_channel) => {
                if let Some(channel209) = maybe_channel {
                    self.available_channels.push_back(channel);
                }33
                self.maybe_available_connection();
            }
            ConnectionRequest::Connected(channel) => {
                self.available_channels.push_back(channel);
                self.maybe_available_connection();
            }
            // Handle a transport error on a connection by making it unavailable
            // for use and establishing a new connection to the endpoint.
            ConnectionRequest::Error((identifier, was_pending)) => {
                let should_reconnect = if was_pending {
                    true
                } else {
                    let original_length = self.available_channels.len();
                    self.available_channels
                        .retain(|channel| channel.identifier != identifier);
                    // Only reconnect if it wasn't already disconnected.
                    original_length != self.available_channels.len()
                };
                if should_reconnect {
                    self.connect_endpoint(identifier.endpoint_index, None);
                }
            }
        }
    }
}

/// An instance of this is obtained for every communication with the gGRPC
/// service.  This handles the permit for limiting concurrency, and also
/// re-connecting the underlying channel on error.  It depends on users
/// reporting all errors.
/// NOTE: This should never be cloneable because its lifetime is linked to the
///       semaphore permit it carries — `_permit` is released exactly once,
///       when the `Connection` drops.
#[derive(Debug)]
pub struct Connection {
    /// Communication with `ConnectionManagerWorker` to inform about transport
    /// errors and when the Connection is dropped.
    tx: mpsc::UnboundedSender<ConnectionRequest>,
    /// If set, the Channel that will be returned to the worker when connection
    /// completes (success or failure) or when the Connection is dropped if that
    /// happens before connection completes.
    pending_channel: Option<Channel>,
    /// The identifier to send to `tx`.
    channel: EstablishedChannel,
    _permit: OwnedSemaphorePermit,
}

impl Drop for Connection {
    fn drop(&mut self) {
        let pending_channel = self
            .pending_channel
            .take()
            .map(|channel| EstablishedChannel {
                channel,
                identifier: self.channel.identifier,
            });
        drop(self.tx.send(ConnectionRequest::Dropped(pending_channel)));
    }
}

/// A wrapper around the `channel::ResponseFuture` that forwards errors to the `tx`.
pub struct ResponseFuture {
    /// The wrapped future that actually does the work.
    inner: channel::ResponseFuture,
    /// Communication with `ConnectionManagerWorker` to inform about transport
    /// errors.
    connection_tx: mpsc::UnboundedSender<ConnectionRequest>,
    /// The identifier to send to `connection_tx` on a transport error.
    identifier: ChannelIdentifier,
}

impl core::fmt::Debug for ResponseFuture {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ResponseFuture")
            .field("inner", &self.inner)
            .field("connection_tx", &self.connection_tx)
            .field("identifier", &self.identifier)
            .finish()
    }
}

/// This is mostly copied from `tonic::transport::channel` except it wraps it
/// to allow messaging about connection success and failure.
impl tonic::codegen::Service<tonic::codegen::http::Request<tonic::body::Body>> for Connection {
    type Response = tonic::codegen::http::Response<tonic::body::Body>;
    type Error = tonic::transport::Error;
    type Future = ResponseFuture;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let result = self.channel.channel.poll_ready(cx);
        if let Poll::Ready(result) = &result {
            match result {
                Ok(()) => {
                    if let Some(pending_channel) = self.pending_channel.take() {
                        drop(
                            self.tx
                                .send(ConnectionRequest::Connected(EstablishedChannel {
                                    channel: pending_channel,
                                    identifier: self.channel.identifier,
                                })),
                        );
                    }0
                }
                Err(err) => {
                    debug!(?err, "Error while creating connection on channel");
                    drop(self.tx.send(ConnectionRequest::Error((
                        self.channel.identifier,
                        self.pending_channel.take().is_some(),
                    ))));
                }
            }
        }
        result
    }

    fn call(&mut self, request: tonic::codegen::http::Request<tonic::body::Body>) -> Self::Future {
        ResponseFuture {
            inner: self.channel.channel.call(request),
            connection_tx: self.tx.clone(),
            identifier: self.channel.identifier,
        }
    }
}

/// This is mostly copied from `tonic::transport::channel` except it wraps it
/// to allow messaging about connection failure.
impl Future for ResponseFuture {
    type Output =
        Result<tonic::codegen::http::Response<tonic::body::Body>, tonic::transport::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let result = Pin::new(&mut self.inner).poll(cx);
        if let Poll::Ready(Err(_)) = &result {
            drop(
                self.connection_tx
                    .send(ConnectionRequest::Error((self.identifier, false))),
            );
        }
        result
    }
}

Coverage Report

Created: 2026-11-13 23:29

Line	Count	Source
1		// Copyright 2024 Trace Machina, Inc. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use core::pin::Pin;
16		use core::task::{Context, Poll};
17		use core::time::Duration;
18		use std::collections::VecDeque;
19		use std::sync::Arc;
20
21		use futures::Future;
22		use futures::stream::{FuturesUnordered, StreamExt, unfold};
23		use nativelink_config::stores::Retry;
24		use nativelink_error::{Code, Error, make_err};
25		use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
26		use tonic::transport::{Channel, Endpoint, channel};
27		use tracing::{debug, error, info, warn};
28
29		use crate::background_spawn;
30		use crate::metrics::record_connection_acquired;
31		use crate::retry::{self, Retrier, RetryResult};
32
33		/// A helper utility that enables management of a suite of connections to an
34		/// upstream gRPC endpoint using Tonic.
35		#[derive(Debug)]
36		pub struct ConnectionManager {
37		// The channel to request connections from the worker.
38		worker_tx: mpsc::Sender<(String, oneshot::Sender<Connection>)>,
39		}
40
41		/// The index into `ConnectionManagerWorker::endpoints`.
42		type EndpointIndex = usize;
43		/// The identifier for a given connection to a given Endpoint, used to identify
44		/// when a particular connection has failed or becomes available.
45		type ConnectionIndex = usize;
46
47		#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
48		struct ChannelIdentifier {
49		/// The index into `ConnectionManagerWorker::endpoints` that established this
50		/// Channel.
51		endpoint_index: EndpointIndex,
52		/// A unique identifier for this particular connection to the Endpoint.
53		connection_index: ConnectionIndex,
54		}
55
56		/// The requests that can be made from a Connection to the
57		/// `ConnectionManagerWorker` such as informing it that it's been dropped or that
58		/// an error occurred.
59		enum ConnectionRequest {
60		/// Notify that a Connection was dropped, if it was dropped while the
61		/// connection was still pending, then return the pending Channel to be
62		/// added back to the available channels.
63		Dropped(Option<EstablishedChannel>),
64		/// Notify that a Connection was established, return the Channel to the
65		/// available channels.
66		Connected(EstablishedChannel),
67		/// Notify that there was a transport error on the given Channel, the bool
68		/// specifies whether the connection was in the process of being established
69		/// or not (i.e. whether it's been returned to available channels yet).
70		Error((ChannelIdentifier, bool)),
71		}
72
73		/// The result of a Future that connects to a given Endpoint. This is a tuple
74		/// of the index into the `ConnectionManagerWorker::endpoints` that this
75		/// connection is for, the iteration of the connection and the result of the
76		/// connection itself.
77		type IndexedChannel = Result<EstablishedChannel, (ChannelIdentifier, Error)>;
78
79		/// A channel that has been established to an endpoint with some metadata around
80		/// it to allow identification of the Channel if it errors in order to correctly
81		/// remove it.
82		#[derive(Debug, Clone)]
83		struct EstablishedChannel {
84		/// The Channel itself that the meta data relates to.
85		channel: Channel,
86		/// The identifier of the channel in the worker.
87		identifier: ChannelIdentifier,
88		}
89
90		/// The context of the worker used to manage all of the connections. This
91		/// handles reconnecting to endpoints on errors and multiple connections to a
92		/// given endpoint.
93		struct ConnectionManagerWorker {
94		/// The endpoints to establish Channels and the identifier of the last
95		/// connection attempt to that endpoint.
96		endpoints: Vec<(ConnectionIndex, Endpoint)>,
97		/// The channel used to communicate between a Connection and the worker.
98		connection_tx: mpsc::UnboundedSender<ConnectionRequest>,
99		/// Gates the maximum number of in-flight `Connection` objects.
100		/// Was an explicit `usize` counter; now an `Arc<Semaphore>` so the
101		/// `OwnedSemaphorePermit` held by each `Connection` releases on
102		/// drop (RAII), instead of relying on a `ConnectionRequest::Dropped`
103		/// round-trip that could be lost on tonic transport errors or task
104		/// aborts.
105		available_connections: Arc<Semaphore>,
106
107		/// Whether `available_connections` reflects a configured limit. When the
108		/// pool is unlimited its permit count is a sentinel, not a measurement.
109		bounded_connections: bool,
110		/// Channels that are currently being connected.
111		connecting_channels: FuturesUnordered<Pin<Box<dyn Future<Output = IndexedChannel> + Send>>>,
112		/// Connected channels that are available for use.
113		available_channels: VecDeque<EstablishedChannel>,
114		/// Requests for a Channel when available - (reason, request)
115		waiting_connections: VecDeque<(String, oneshot::Sender<Connection>)>,
116		/// The retry configuration for connecting to an Endpoint, on failure will
117		/// restart the retrier after a 1 second delay.
118		retrier: Retrier,
119		}
120
121		/// The maximum number of queued requests to obtain a connection from the
122		/// worker before applying back pressure to the requestor. It makes sense to
123		/// keep this small since it has to wait for a response anyway.
124		const WORKER_BACKLOG: usize = 8;
125
126		impl ConnectionManager {
127		/// Create a connection manager that creates a balance list between a given
128		/// set of Endpoints. This will restrict the number of concurrent requests
129		/// and automatically re-connect upon transport error.
130	51	pub fn new(
131	51	endpoints: impl IntoIterator<Item = Endpoint>,
132	51	mut connections_per_endpoint: usize,
133	51	mut max_concurrent_requests: usize,
134	51	retry: Retry,
135	51	jitter_fn: retry::JitterFn,
136	51	) -> Self {
137	51	let (worker_tx, worker_rx) = mpsc::channel(WORKER_BACKLOG);
138		// The connection messages always come from sync contexts (e.g. drop)
139		// and therefore, we'd end up spawning for them if this was bounded
140		// which defeats the object since there would be no backpressure
141		// applied. Therefore it makes sense for this to be unbounded.
142	51	let (connection_tx, connection_rx) = mpsc::unbounded_channel();
143	51	let endpoints = endpoints
144	51	.into_iter()
145	51	.map(\|endpoint\| (0, endpoint))
146	51	.collect();
147
148		// Zero means unlimited, which becomes a semaphore holding
149		// Semaphore::MAX_PERMITS. That is a sentinel, not a free-slot count,
150		// so remember which case this is and report no figure when unbounded.
151	51	let bounded_connections = max_concurrent_requests != 0;
152	51	if max_concurrent_requests == 0 {
153	48	max_concurrent_requests = Semaphore::MAX_PERMITS;
154	48	} else {
155	3	max_concurrent_requests = max_concurrent_requests.min(Semaphore::MAX_PERMITS);
156	3	}
157	51	if connections_per_endpoint == 0 {
158	48	connections_per_endpoint = 1;
159	48	}3
160	51	let worker = ConnectionManagerWorker {
161	51	endpoints,
162	51	available_connections: Arc::new(Semaphore::new(max_concurrent_requests)),
163	51	bounded_connections,
164	51	connection_tx,
165	51	connecting_channels: FuturesUnordered::new(),
166	51	available_channels: VecDeque::new(),
167	51	waiting_connections: VecDeque::new(),
168	51	retrier: Retrier::new(
169	51	Arc::new(\|duration\| Box::pin0 ( tokio::time::sleep0 ( duration0 ))),
170	51	jitter_fn,
171	51	retry,
172		)
173	51	.with_client("connection"),
174		};
175	51	background_spawn!("connection_manager_worker_spawn", async move {47
176	47	worker
177	47	.service_requests(connections_per_endpoint, worker_rx, connection_rx)
178	47	.await;
179	7	});
180	51	Self { worker_tx }
181	51	}
182
183		/// Get a Connection that can be used as a `tonic::Channel`, except it
184		/// performs some additional counting to reconnect on error and restrict
185		/// the number of concurrent connections.
186	271	pub async fn connection(&self, reason: String) -> Result<Connection, Error>0 {
187	271	let (tx, rx) = oneshot::channel();
188	271	self.worker_tx
189	271	.send((reason, tx))
190	271	.await
191	271	.map_err(\|err\| make_err!0 ( Code::Unavailable0 , "Requesting a new connection: {err:?}")) ?0 ;
192	271	rx.await
193	263	.map_err(\|err\| make_err!0 ( Code::Unavailable0 , "Waiting for a new connection: {err:?}"))
194	263	}
195		}
196
197		impl ConnectionManagerWorker {
198	0	async fn service_requests(
199	0	mut self,
200	0	connections_per_endpoint: usize,
201	0	mut worker_rx: mpsc::Receiver<(String, oneshot::Sender<Connection>)>,
202	0	mut connection_rx: mpsc::UnboundedReceiver<ConnectionRequest>,
203	47	) {
204		// Make the initial set of connections, connection failures will be
205		// handled in the same way as future transport failures, so no need to
206		// do anything special.
207	47	for endpoint_index in 0..self.endpoints.len() {
208	51	for _ in 0..connections_per_endpoint47 {
209	51	self.connect_endpoint(endpoint_index, None);
210	51	}
211		}
212
213		// The main worker loop, when select resolves one of its arms the other
214		// ones are cancelled, therefore it's important that they maintain no
215		// state while `await`-ing. This is enforced through the use of
216		// non-async functions to do all of the work.
217		loop {
218	666	tokio::select! {
219	666	request278 = worker_rx.recv() => {
220	278	let Some(( reason271 , request271 )) = request else {
221		// The ConnectionManager was dropped, shut down the
222		// worker.
223	7	break;
224		};
225	271	self.handle_worker(reason, request);
226		}
227	666	maybe_request298 = connection_rx.recv() => {
228	298	if let Some(request) = maybe_request {
229	298	self.handle_connection(request);
230	298	}0
231		}
232	666	maybe_connection_result50 = self.connect_next() => {
233	50	if let Some(connection_result) = maybe_connection_result {
234	50	self.handle_connected(connection_result);
235	50	}0
236		}
237		}
238		}
239	7	}
240
241	508	async fn connect_next(&mut self) -> Option<IndexedChannel>0 {
242	508	if self.connecting_channels.is_empty() {
243		// Make this Future never resolve, we will get cancelled by the
244		// select if there's some change in state to `self` and can re-enter
245		// and evaluate `connecting_channels` again.
246	415	futures::future::pending::<()>().await;
247	93	}
248	93	self.connecting_channels.next().await
249	50	}
250
251		// This must never be made async otherwise the select may cancel it.
252	50	fn handle_connected(&mut self, connection_result: IndexedChannel) {
253	50	match connection_result {
254	50	Ok(established_channel) => {
255	50	self.available_channels.push_back(established_channel);
256	50	self.maybe_available_connection();
257	50	}
258		// When the retrier runs out of attempts start again from the
259		// beginning of the retry period. Never want to be in a
260		// situation where we give up on an Endpoint forever.
261	0	Err((identifier, _)) => {
262	0	self.connect_endpoint(identifier.endpoint_index, Some(identifier.connection_index));
263	0	}
264		}
265	50	}
266
267	51	fn connect_endpoint(&mut self, endpoint_index: usize, connection_index: Option<usize>) {
268	51	let Some((current_connection_index, endpoint)) = self.endpoints.get_mut(endpoint_index)
269		else {
270		// Unknown endpoint, this should never happen.
271	0	error!(?endpoint_index, "Connection to unknown endpoint requested");
272	0	return;
273		};
274	51	let is_backoff = connection_index.is_some();
275	51	let connection_index = connection_index.unwrap_or_else(\|\| {
276	51	*current_connection_index += 1;
277	51	*current_connection_index
278	51	});
279	51	if is_backoff {
280	0	warn!(
281		?connection_index,
282	0	endpoint = ?endpoint.uri(),
283		"Connection failed, reconnecting"
284		);
285		} else {
286	51	info!(
287		?connection_index,
288	51	endpoint = ?endpoint.uri(),
289		"Creating new connection"
290		);
291		}
292	51	let identifier = ChannelIdentifier {
293	51	endpoint_index,
294	51	connection_index,
295	51	};
296	51	let connection_stream = unfold(endpoint.clone(), move \|endpoint\| async move {
297	51	let result50 = endpoint.connect().await. map_err50 (\|err\| {0
298	0	make_err!(
299	0	Code::Unavailable,
300		"Failed to connect to {:?}: {err:?}",
301	0	endpoint.uri()
302		)
303	0	});
304	50	Some((
305	50	result.map_or_else(RetryResult::Retry, RetryResult::Ok),
306	50	endpoint,
307	50	))
308	101	});
309	51	let retrier = self.retrier.clone();
310	51	self.connecting_channels.push(Box::pin(async move {
311	51	if is_backoff {
312		// Just in case the retry config is 0, then we need to
313		// introduce some delay so we aren't in a hard loop.
314	0	tokio::time::sleep(Duration::from_secs(1)).await;
315	51	}
316	51	retrier.retry(connection_stream).await. map_or_else50 (
317	0	\|err\| Err((identifier, err)),
318	50	\|channel\| {
319	50	Ok(EstablishedChannel {
320	50	channel,
321	50	identifier,
322	50	})
323	50	},
324		)
325	50	}));
326	51	}
327
328		/// Free slots, or `None` when the pool is unlimited and the permit count
329		/// is a sentinel rather than a measurement.
330	271	fn free_connection_slots(&self) -> Option<usize> {
331	271	self.bounded_connections
332	271	.then(\|\| self.available_connections215 . available_permits215 ())
333	271	}
334
335		// This must never be made async otherwise the select may cancel it.
336	271	fn handle_worker(&mut self, reason: String, tx: oneshot::Sender<Connection>) {
337	271	let maybe_permit = self.available_connections.clone().try_acquire_owned().ok();
338	271	if let Some( permit270 ) = maybe_permit
339	270	&& let Some( channel157 ) = self.available_channels.pop_front()
340		{
341	157	debug!(reason, "ConnectionManager: request running");
342	157	record_connection_acquired("grpc", self.free_connection_slots(), false);
343	157	self.provide_channel(channel, tx, permit);
344		} else {
345	114	debug!(
346	114	available_permits = self.available_connections.available_permits(),
347	114	available_channels = self.available_channels.len(),
348	114	waiting_connections = self.waiting_connections.len(),
349		reason,
350		"ConnectionManager: no connection available, request queued",
351		);
352	114	record_connection_acquired("grpc", self.free_connection_slots(), true);
353	114	self.waiting_connections.push_back((reason, tx));
354		}
355	271	}
356
357	271	fn provide_channel(
358	271	&self,
359	271	channel: EstablishedChannel,
360	271	tx: oneshot::Sender<Connection>,
361	271	permit: OwnedSemaphorePermit,
362	271	) {
363	271	drop(tx.send(Connection {
364	271	tx: self.connection_tx.clone(),
365	271	pending_channel: Some(channel.channel.clone()),
366	271	channel,
367	271	_permit: permit,
368	271	}));
369	271	}
370
371	348	fn maybe_available_connection(&mut self) {
372	462	while !self.waiting_connections.is_empty() && !self.available_channels.is_empty()123 {
373	114	let Some(permit) = self.available_connections.clone().try_acquire_owned().ok() else {
374	0	break;
375		};
376	114	let Some(channel) = self.available_channels.pop_front() else {
377	0	drop(permit);
378	0	break;
379		};
380	114	let Some((reason, tx)) = self.waiting_connections.pop_front() else {
381	0	self.available_channels.push_front(channel);
382	0	drop(permit);
383	0	break;
384		};
385	114	debug!(reason, "ConnectionManager: channel available, running");
386	114	self.provide_channel(channel, tx, permit);
387		}
388	348	}
389
390		// This must never be made async otherwise the select may cancel it.
391	298	fn handle_connection(&mut self, request: ConnectionRequest) {
392	298	match request {
393	242	ConnectionRequest::Dropped(maybe_channel) => {
394	242	if let Some( channel209 ) = maybe_channel {
395	209	self.available_channels.push_back(channel);
396	209	}33
397	242	self.maybe_available_connection();
398		}
399	56	ConnectionRequest::Connected(channel) => {
400	56	self.available_channels.push_back(channel);
401	56	self.maybe_available_connection();
402	56	}
403		// Handle a transport error on a connection by making it unavailable
404		// for use and establishing a new connection to the endpoint.
405	0	ConnectionRequest::Error((identifier, was_pending)) => {
406	0	let should_reconnect = if was_pending {
407	0	true
408		} else {
409	0	let original_length = self.available_channels.len();
410	0	self.available_channels
411	0	.retain(\|channel\| channel.identifier != identifier);
412		// Only reconnect if it wasn't already disconnected.
413	0	original_length != self.available_channels.len()
414		};
415	0	if should_reconnect {
416	0	self.connect_endpoint(identifier.endpoint_index, None);
417	0	}
418		}
419		}
420	298	}
421		}
422
423		/// An instance of this is obtained for every communication with the gGRPC
424		/// service. This handles the permit for limiting concurrency, and also
425		/// re-connecting the underlying channel on error. It depends on users
426		/// reporting all errors.
427		/// NOTE: This should never be cloneable because its lifetime is linked to the
428		/// semaphore permit it carries — `_permit` is released exactly once,
429		/// when the `Connection` drops.
430		#[derive(Debug)]
431		pub struct Connection {
432		/// Communication with `ConnectionManagerWorker` to inform about transport
433		/// errors and when the Connection is dropped.
434		tx: mpsc::UnboundedSender<ConnectionRequest>,
435		/// If set, the Channel that will be returned to the worker when connection
436		/// completes (success or failure) or when the Connection is dropped if that
437		/// happens before connection completes.
438		pending_channel: Option<Channel>,
439		/// The identifier to send to `tx`.
440		channel: EstablishedChannel,
441		_permit: OwnedSemaphorePermit,
442		}
443
444		impl Drop for Connection {
445	271	fn drop(&mut self) {
446	271	let pending_channel = self
447	271	.pending_channel
448	271	.take()
449	271	.map(\|channel\| EstablishedChannel {
450	215	channel,
451	215	identifier: self.channel.identifier,
452	215	});
453	271	drop(self.tx.send(ConnectionRequest::Dropped(pending_channel)));
454	271	}
455		}
456
457		/// A wrapper around the `channel::ResponseFuture` that forwards errors to the `tx`.
458		pub struct ResponseFuture {
459		/// The wrapped future that actually does the work.
460		inner: channel::ResponseFuture,
461		/// Communication with `ConnectionManagerWorker` to inform about transport
462		/// errors.
463		connection_tx: mpsc::UnboundedSender<ConnectionRequest>,
464		/// The identifier to send to `connection_tx` on a transport error.
465		identifier: ChannelIdentifier,
466		}
467
468		impl core::fmt::Debug for ResponseFuture {
469	0	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
470	0	f.debug_struct("ResponseFuture")
471	0	.field("inner", &self.inner)
472	0	.field("connection_tx", &self.connection_tx)
473	0	.field("identifier", &self.identifier)
474	0	.finish()
475	0	}
476		}
477
478		/// This is mostly copied from `tonic::transport::channel` except it wraps it
479		/// to allow messaging about connection success and failure.
480		impl tonic::codegen::Service<tonic::codegen::http::Request<tonic::body::Body>> for Connection {
481		type Response = tonic::codegen::http::Response<tonic::body::Body>;
482		type Error = tonic::transport::Error;
483		type Future = ResponseFuture;
484
485	56	fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
486	56	let result = self.channel.channel.poll_ready(cx);
487	56	if let Poll::Ready(result) = &result {
488	56	match result {
489		Ok(()) => {
490	56	if let Some(pending_channel) = self.pending_channel.take() {
491	56	drop(
492	56	self.tx
493	56	.send(ConnectionRequest::Connected(EstablishedChannel {
494	56	channel: pending_channel,
495	56	identifier: self.channel.identifier,
496	56	})),
497	56	);
498	56	}0
499		}
500	0	Err(err) => {
501	0	debug!(?err, "Error while creating connection on channel");
502	0	drop(self.tx.send(ConnectionRequest::Error((
503	0	self.channel.identifier,
504	0	self.pending_channel.take().is_some(),
505	0	))));
506		}
507		}
508	0	}
509	56	result
510	56	}
511
512	56	fn call(&mut self, request: tonic::codegen::http::Request<tonic::body::Body>) -> Self::Future {
513	56	ResponseFuture {
514	56	inner: self.channel.channel.call(request),
515	56	connection_tx: self.tx.clone(),
516	56	identifier: self.channel.identifier,
517	56	}
518	56	}
519		}
520
521		/// This is mostly copied from `tonic::transport::channel` except it wraps it
522		/// to allow messaging about connection failure.
523		impl Future for ResponseFuture {
524		type Output =
525		Result<tonic::codegen::http::Response<tonic::body::Body>, tonic::transport::Error>;
526
527	171	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
528	171	let result = Pin::new(&mut self.inner).poll(cx);
529	56	if let Poll::Ready(Err(_)) = &result {
530	0	drop(
531	0	self.connection_tx
532	0	.send(ConnectionRequest::Error((self.identifier, false))),
533	0	);
534	171	}
535	171	result
536	171	}
537		}