datafusion_distributed/protocol/grpc/
channel_resolver.rs

1use crate::protocol::grpc::worker_client::create_worker_client;
2use crate::{ChannelResolver, WorkerChannel};
3use async_trait::async_trait;
4use datafusion::common::{DataFusionError, config_datafusion_err, exec_datafusion_err};
5use futures::FutureExt;
6use futures::future::Shared;
7use std::sync::{Arc, LazyLock};
8use std::time::Duration;
9use tonic::body::Body;
10use tonic::codegen::BoxFuture;
11use tonic::transport::Channel;
12use tower::ServiceExt;
13use url::Url;
14
15// Unlike TaskContext, a DataFusion RuntimeEnv does not allow to introduce user-defined extensions.
16// For the default implementation of the ChannelResolvers, we cannot inject one DefaultChannelResolver
17// per TaskContext, as this holds reference to Tonic channels that must outlive a single TaskContext.
18//
19// The Tonic channels need to be established and reused under a whole RuntimeEnv scope, not a single
20// TaskContext, which forces us to put the default implementation in a static global variable that
21// stores and reuses tonic channels per RuntimeEnv's pointer address.
22pub(crate) static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock<
23    moka::sync::Cache<
24        /* Arc<RuntimeEnv> pointer address */ usize,
25        /* ChannelResolver that reuses built channels */ Arc<DefaultChannelResolver>,
26    >,
27> = LazyLock::new(|| moka::sync::Cache::builder().max_capacity(256).build());
28
29pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
30    http::Request<Body>,
31    http::Response<Body>,
32    tonic::transport::Error,
33>;
34
35type ChannelCacheValue = Shared<BoxFuture<BoxCloneSyncChannel, Arc<DataFusionError>>>;
36
37/// Default implementation of a [ChannelResolver] that connects to the workers given the URL once
38/// and stores the connection instance in a TTI cache.
39///
40/// Sane default over which other [ChannelResolver] can be built for better customization of the
41/// [WorkerServiceClient]s.
42#[derive(Clone)]
43pub struct DefaultChannelResolver {
44    cache: Arc<moka::sync::Cache<Url, ChannelCacheValue>>,
45}
46
47impl Default for DefaultChannelResolver {
48    fn default() -> Self {
49        Self {
50            cache: Arc::new(
51                moka::sync::Cache::builder()
52                    // Use an unrealistic max capacity, just in case there is a logic error on the
53                    // user part that produces an unreasonable amount of URLs.
54                    .max_capacity(64556)
55                    // If a channel has not been used in 5 mins, delete it.
56                    .time_to_idle(Duration::from_secs(5 * 60))
57                    .build(),
58            ),
59        }
60    }
61}
62
63impl DefaultChannelResolver {
64    /// Gets the cached [BoxCloneSyncChannel] for the given URL, or builds a new one.
65    pub async fn get_channel(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
66        let channel = self.cache.get_with_by_ref(url, move || {
67            let url = url.to_string();
68            async move {
69                let endpoint = Channel::from_shared(url.clone()).map_err(|err| {
70                    config_datafusion_err!(
71                        "Invalid URL '{url}' returned by WorkerResolver implementation: {err}"
72                    )
73                })?;
74                let mut channel = endpoint.connect().await.map_err(|err| {
75                    DataFusionError::Context(
76                        format!("{err:?}"),
77                        Box::new(exec_datafusion_err!(
78                            "Error connecting to Distributed DataFusion worker on '{url}': {err}"
79                        )),
80                    )
81                })?;
82                channel.ready().await.map_err(|err| {
83                    DataFusionError::Context(
84                        format!("{err:?}"),
85                        Box::new(exec_datafusion_err!(
86                            "Error waiting for Distributed DataFusion channel to be ready on '{url}': {err}"
87                        )),
88                    )
89                })?;
90                Ok(BoxCloneSyncChannel::new(channel))
91            }
92                .boxed()
93                .shared()
94        });
95
96        channel.await.map_err(|err| {
97            self.cache.invalidate(url);
98            DataFusionError::Shared(err)
99        })
100    }
101}
102
103#[async_trait]
104impl ChannelResolver for DefaultChannelResolver {
105    async fn get_worker_client_for_url(
106        &self,
107        url: &Url,
108    ) -> Result<Box<dyn WorkerChannel>, DataFusionError> {
109        self.get_channel(url).await.map(create_worker_client)
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use crate::Worker;
117    use datafusion::common::assert_contains;
118    use datafusion::common::runtime::SpawnedTask;
119    use std::error::Error;
120    use std::time::Instant;
121    use tokio::net::TcpListener;
122    use tonic::transport::Server;
123
124    #[tokio::test]
125    async fn fails_establishing_connection() -> Result<(), Box<dyn Error>> {
126        let (url, _guard) = spawn_http_localhost_worker().await?;
127        drop(_guard);
128        let channel_resolver = DefaultChannelResolver::default();
129        let err = channel_resolver.get_channel(&url).await.unwrap_err();
130        assert_contains!(err.to_string(), "tcp connect error");
131        Ok(())
132    }
133
134    #[tokio::test]
135    async fn can_establish_connection() -> Result<(), Box<dyn Error>> {
136        let (url, _guard) = spawn_http_localhost_worker().await?;
137        let channel_resolver = DefaultChannelResolver::default();
138        channel_resolver.get_channel(&url).await?;
139        Ok(())
140    }
141
142    #[tokio::test]
143    async fn channel_resolve_is_cached() -> Result<(), Box<dyn Error>> {
144        let (url, _guard) = spawn_http_localhost_worker().await?;
145        let channel_resolver = DefaultChannelResolver::default();
146
147        let start = Instant::now();
148        channel_resolver.get_channel(&url).await?;
149        let first_call = start.elapsed();
150
151        let start = Instant::now();
152        channel_resolver.get_channel(&url).await?;
153        let second_call = start.elapsed();
154
155        assert!(first_call > second_call);
156        Ok(())
157    }
158
159    async fn spawn_http_localhost_worker() -> Result<(Url, SpawnedTask<()>), Box<dyn Error>> {
160        let listener = TcpListener::bind("127.0.0.1:0").await?;
161
162        let port = listener
163            .local_addr()
164            .expect("Failed to get local address")
165            .port();
166
167        let task = SpawnedTask::spawn(async {
168            let worker = Worker::default();
169            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
170            if let Err(err) = Server::builder()
171                .add_service(worker.into_worker_server())
172                .serve_with_incoming(incoming)
173                .await
174            {
175                panic!("{err}")
176            }
177        });
178
179        Ok((Url::parse(&format!("http://127.0.0.1:{port}"))?, task))
180    }
181}
datafusion_distributed/protocol/grpc/channel_resolver.rs

datafusion_distributed/protocol/grpc/
channel_resolver.rs
