refactor: centralize console worker connections

console/src/app.rs

@@ -1,3 +1,4 @@
+use crate::connector::WorkerConnector;
 use crate::state::{ClusterViewState, SortColumn, SortDirection, View, WorkerViewState};
 use crate::worker::{ConnectionStatus, WorkerConn, discover_cluster_workers};
 use datafusion::common::HashSet;
@@ -23,6 +24,8 @@ pub(crate) struct App {
     pub(crate) current_throughput: f64,
     /// Seed URL for worker discovery via `GetClusterWorkers`.
     seed_url: Url,
+    /// Transport policy shared by discovery and every per-worker connection.
+    connector: WorkerConnector,
     /// Last time we ran worker discovery.
     last_discovery: Option<Instant>,
 }
@@ -44,7 +47,7 @@ const DISCOVERY_INTERVAL: Duration = Duration::from_secs(5);
 
 impl App {
     /// Create a new App that discovers workers via `GetClusterWorkers` on the seed URL.
-    pub(crate) fn new(seed_url: Url) -> Self {
+    pub(crate) fn new(seed_url: Url, connector: WorkerConnector) -> Self {
         App {
             workers: Vec::new(),
             active_query_count: 0,
@@ -60,6 +63,7 @@ impl App {
             prev_output_rows_time: None,
             current_throughput: 0.0,
             seed_url,
+            connector,
             last_discovery: None,
         }
     }
@@ -112,7 +116,8 @@ impl App {
 
         self.last_discovery = Some(Instant::now());
 
-        let discovered_urls = match discover_cluster_workers(&self.seed_url).await {
+        let discovered_urls = match discover_cluster_workers(&self.connector, &self.seed_url).await
+        {
             Ok(urls) => urls,
             Err(_) => return,
         };
@@ -123,7 +128,8 @@ impl App {
         // Add new workers
         for url in &discovered_urls {
             if !known_urls.contains(url) {
-                self.workers.push(WorkerConn::new(url.clone()));
+                self.workers
+                    .push(WorkerConn::new(url.clone(), self.connector.clone()));
             }
         }
 

console/src/connector.rs

@@ -0,0 +1,67 @@
+use datafusion_distributed::grpc::ObservabilityServiceClient;
+use tonic::transport::{Channel, Error};
+use url::Url;
+
+/// Opens gRPC channels to the workers the console monitors.
+///
+/// Per-worker polling ([`crate::worker::WorkerConn`]) and cluster discovery
+/// ([`crate::worker::discover_cluster_workers`]) both dial workers, and both must agree on how a
+/// worker URL becomes a connection. Holding that decision in one cheap, cloneable value keeps the
+/// two call sites from drifting apart when the transport changes.
+#[derive(Clone, Debug, Default)]
+pub(crate) struct WorkerConnector;
+
+impl WorkerConnector {
+    /// Opens a channel to a worker.
+    ///
+    /// Connecting alone does not prove the worker serves the observability API, so callers that
+    /// care about liveness follow up with their own `Ping`.
+    pub(crate) async fn connect(
+        &self,
+        url: &Url,
+    ) -> Result<ObservabilityServiceClient<Channel>, Error> {
+        ObservabilityServiceClient::connect(url.to_string()).await
+    }
+
+    /// Turns a worker URL reported by `GetClusterWorkers` into a URL this connector can dial.
+    ///
+    /// Workers describe themselves through their own `WorkerResolver`, so the reported form is
+    /// whatever that implementation chose and is not guaranteed to be directly dialable.
+    pub(crate) fn worker_url(&self, reported: &str) -> Result<Url, url::ParseError> {
+        Url::parse(reported)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::error::Error;
+    use tokio::net::TcpListener;
+
+    #[test]
+    fn worker_url_keeps_a_reported_url_as_is() -> Result<(), Box<dyn Error>> {
+        let connector = WorkerConnector;
+        let url = connector.worker_url("http://10.0.0.1:9001")?;
+        assert_eq!(url.as_str(), "http://10.0.0.1:9001/");
+        Ok(())
+    }
+
+    #[test]
+    fn worker_url_rejects_a_value_that_is_not_a_url() {
+        let connector = WorkerConnector;
+        assert!(connector.worker_url("not a url").is_err());
+    }
+
+    #[tokio::test]
+    async fn connect_fails_when_nothing_is_listening() -> Result<(), Box<dyn Error>> {
+        // Binding and dropping a listener yields a port that is very unlikely to be reused.
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let port = listener.local_addr()?.port();
+        drop(listener);
+
+        let connector = WorkerConnector;
+        let url = Url::parse(&format!("http://127.0.0.1:{port}"))?;
+        assert!(connector.connect(&url).await.is_err());
+        Ok(())
+    }
+}

console/src/main.rs

@@ -1,10 +1,12 @@
 mod app;
+mod connector;
 mod input;
 mod state;
 mod ui;
 mod worker;
 
 use app::App;
+use connector::WorkerConnector;
 use crossterm::event::{self, Event};
 use ratatui::DefaultTerminal;
 use std::time::{Duration, Instant};
@@ -35,7 +37,7 @@ async fn main() -> color_eyre::Result<()> {
     let seed_url = Url::parse(&format!("http://localhost:{}", args.port)).expect("valid URL");
 
     let poll_interval = Duration::from_millis(args.poll_interval);
-    let mut app = App::new(seed_url);
+    let mut app = App::new(seed_url, WorkerConnector);
 
     let mut terminal = ratatui::init();
     terminal.clear()?;

console/src/worker.rs

@@ -1,3 +1,4 @@
+use crate::connector::WorkerConnector;
 use datafusion::common::{HashMap, HashSet};
 use datafusion_distributed::grpc::{
     GetClusterWorkersRequest, GetTaskProgressRequest, ObservabilityServiceClient, PingRequest,
@@ -17,6 +18,7 @@ pub(crate) const METRIC_HISTORY_LEN: usize = 300;
 /// Tracks connection and task state for a single worker.
 pub(crate) struct WorkerConn {
     pub(crate) url: Url,
+    connector: WorkerConnector,
     client: Option<ObservabilityServiceClient<Channel>>,
     pub(crate) connection_status: ConnectionStatus,
     pub(crate) tasks: Vec<TaskProgress>,
@@ -63,9 +65,10 @@ pub(crate) enum ConnectionStatus {
 
 impl WorkerConn {
     /// Create a new WorkerConn in the initial Connecting state.
-    pub(crate) fn new(url: Url) -> Self {
+    pub(crate) fn new(url: Url, connector: WorkerConnector) -> Self {
         Self {
             url,
+            connector,
             client: None,
             connection_status: ConnectionStatus::Connecting,
             tasks: Vec::new(),
@@ -89,7 +92,7 @@ impl WorkerConn {
     pub(crate) async fn try_connect(&mut self) {
         self.last_reconnect_attempt = Some(Instant::now());
 
-        match ObservabilityServiceClient::connect(self.url.to_string()).await {
+        match self.connector.connect(&self.url).await {
             Ok(mut client) => match client.ping(PingRequest {}).await {
                 Ok(_) => {
                     self.client = Some(client);
@@ -368,8 +371,12 @@ fn push_history(buf: &mut VecDeque<u64>, value: u64) {
 }
 
 /// Connects to a seed worker and calls `GetClusterWorkers` to discover all worker URLs.
-pub(crate) async fn discover_cluster_workers(seed_url: &Url) -> Result<Vec<Url>, String> {
-    let mut client = ObservabilityServiceClient::connect(seed_url.to_string())
+pub(crate) async fn discover_cluster_workers(
+    connector: &WorkerConnector,
+    seed_url: &Url,
+) -> Result<Vec<Url>, String> {
+    let mut client = connector
+        .connect(seed_url)
         .await
         .map_err(|e| format!("Failed to connect to seed worker {seed_url}: {e}"))?;
 
@@ -386,8 +393,8 @@ pub(crate) async fn discover_cluster_workers(seed_url: &Url) -> Result<Vec<Url>,
     let urls = response
         .into_inner()
         .worker_urls
-        .into_iter()
-        .filter_map(|s| Url::parse(&s).ok())
+        .iter()
+        .filter_map(|reported| connector.worker_url(reported).ok())
         .collect();
 
     Ok(urls)