From 513c8814259b6c8d68cdcfe1a812156a5fafac4c Mon Sep 17 00:00:00 2001 From: Edson Petry Date: Mon, 10 Aug 2026 14:21:35 -0400 Subject: [PATCH] refactor: centralize console worker connection creation --- console/src/app.rs | 12 +++++-- console/src/connector.rs | 67 ++++++++++++++++++++++++++++++++++++++++ console/src/main.rs | 4 ++- console/src/worker.rs | 19 ++++++++---- 4 files changed, 92 insertions(+), 10 deletions(-) create mode 100644 console/src/connector.rs diff --git a/console/src/app.rs b/console/src/app.rs index 1e874771e..6f0f2cb2a 100644 --- a/console/src/app.rs +++ b/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, } @@ -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())); } } diff --git a/console/src/connector.rs b/console/src/connector.rs new file mode 100644 index 000000000..006944bf9 --- /dev/null +++ b/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, 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::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> { + 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> { + // 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(()) + } +} diff --git a/console/src/main.rs b/console/src/main.rs index 731b504ba..706a4ff5e 100644 --- a/console/src/main.rs +++ b/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()?; diff --git a/console/src/worker.rs b/console/src/worker.rs index 4a23c5b90..05c001016 100644 --- a/console/src/worker.rs +++ b/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>, pub(crate) connection_status: ConnectionStatus, pub(crate) tasks: Vec, @@ -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, 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, String> { - let mut client = ObservabilityServiceClient::connect(seed_url.to_string()) +pub(crate) async fn discover_cluster_workers( + connector: &WorkerConnector, + seed_url: &Url, +) -> Result, 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, 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)