datafusion_distributed/worker/
worker_service.rs

1use crate::protocol::LocalWorkerContext;
2use crate::worker::{SingleWriteMultiRead, WorkerSessionBuilder};
3use crate::{DefaultSessionBuilder, TaskData, TaskKey};
4use datafusion::common::DataFusionError;
5use datafusion::execution::runtime_env::RuntimeEnv;
6use moka::future::Cache;
7use std::borrow::Cow;
8use std::sync::Arc;
9use std::time::Duration;
10use url::Url;
11
12const TASK_CACHE_TTI: Duration = Duration::from_mins(10);
13
14pub(crate) type ResultTaskData = Result<TaskData, Arc<DataFusionError>>;
15pub(crate) type TaskDataEntries = Cache<TaskKey, Arc<SingleWriteMultiRead<ResultTaskData>>>;
16
17#[derive(Clone)]
18pub struct Worker {
19    pub(super) runtime: Arc<RuntimeEnv>,
20    /// TTL-based cache for task execution data. Entries are automatically evicted after
21    /// TASK_CACHE_TTI seconds. This prevents memory leaks from abandoned or incomplete queries
22    /// while allowing concurrent access to task results across multiple partition requests.
23    pub(crate) task_data_entries: Arc<TaskDataEntries>,
24    pub(super) session_builder: Arc<dyn WorkerSessionBuilder + Send + Sync>,
25    pub(crate) max_message_size: Option<usize>,
26    pub(super) version: Cow<'static, str>,
27}
28
29impl Default for Worker {
30    fn default() -> Self {
31        let cache = Cache::builder().time_to_idle(TASK_CACHE_TTI).build();
32        Self {
33            runtime: Arc::new(RuntimeEnv::default()),
34            task_data_entries: Arc::new(cache),
35            session_builder: Arc::new(DefaultSessionBuilder),
36            max_message_size: Some(usize::MAX),
37            version: Cow::Borrowed(""),
38        }
39    }
40}
41
42impl Worker {
43    /// Builds a [Worker] with a custom [WorkerSessionBuilder]. Use this
44    /// method whenever you need to add custom stuff to the `SessionContext` that executes the query.
45    pub fn from_session_builder(
46        session_builder: impl WorkerSessionBuilder + Send + Sync + 'static,
47    ) -> Self {
48        Self {
49            session_builder: Arc::new(session_builder),
50            ..Default::default()
51        }
52    }
53
54    /// Sets a [RuntimeEnv] to be used in all the queries this [Worker] will handle during
55    /// its lifetime.
56    pub fn with_runtime_env(mut self, runtime_env: Arc<RuntimeEnv>) -> Self {
57        self.runtime = runtime_env;
58        self
59    }
60
61    /// Set the maximum message size for FlightData chunks.
62    ///
63    /// Defaults to `usize::MAX` to minimize chunking overhead for internal communication.
64    /// See [`FlightDataEncoderBuilder::with_max_flight_data_size`] for details.
65    ///
66    /// If you change this to a lower value, ensure you configure the server's
67    /// max_encoding_message_size and max_decoding_message_size to at least 2x this value
68    /// to allow for overhead. For most use cases, the default of `usize::MAX` is appropriate.
69    ///
70    /// [`FlightDataEncoderBuilder::with_max_flight_data_size`]: https://arrow.apache.org/rust/arrow_flight/encode/struct.FlightDataEncoderBuilder.html#structfield.max_flight_data_size
71    pub fn with_max_message_size(mut self, size: usize) -> Self {
72        self.max_message_size = Some(size);
73        self
74    }
75
76    /// Sets a version string reported by the `GetWorkerInfo` gRPC endpoint.
77    pub fn with_version(mut self, version: impl Into<Cow<'static, str>>) -> Self {
78        self.version = version.into();
79        self
80    }
81
82    /// Returns the version set by [Self::with_version].
83    pub fn version(&self) -> &str {
84        &self.version
85    }
86
87    /// Builds a [LocalWorkerContext] suitable to be injecting into a coordinating [SessionContext].
88    /// Having a [LocalWorkerContext] present in the coordinating [SessionContext] is not strictly
89    /// necessary, but it allows the planner to better colocate small stages near it, avoiding
90    /// unnecessary network hops.
91    pub fn to_local_worker_context(&self, self_url: Url) -> LocalWorkerContext {
92        LocalWorkerContext {
93            local_worker: self.clone(),
94            self_url,
95        }
96    }
97
98    /// Returns the number of cached task entries currently held by this worker.
99    #[cfg(any(test, feature = "integration"))]
100    pub async fn tasks_running(&self) -> usize {
101        // Use `run_pending_tasks()` to migigate inaccuracy from potential stale
102        // `entry_count()` task data.
103        self.task_data_entries.run_pending_tasks().await;
104        self.task_data_entries.entry_count() as usize
105    }
106}
datafusion_distributed/worker/worker_service.rs

datafusion_distributed/worker/
worker_service.rs
