From 9936a83384dca7e0871238c4f8bf3f6a714edf9c Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 7 Jul 2026 15:18:23 +0800 Subject: [PATCH] feat(query): add runtime provider interface (#8386) * feat(query): add runtime provider interface Signed-off-by: discord9 * docs(query): document runtime provider interface Signed-off-by: discord9 * feat(query): pass runtime builder to provider Signed-off-by: discord9 * fix(query): make runtime provider fallible Signed-off-by: discord9 * fix(query): propagate runtime provider errors Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/datanode/src/datanode.rs | 9 +- src/flow/src/server.rs | 11 +- src/frontend/src/instance/builder.rs | 5 +- src/query/src/query_engine.rs | 59 +++++++- src/query/src/query_engine/runtime.rs | 80 ++++++++++ src/query/src/query_engine/state.rs | 205 ++++++++++++++++++++++++-- 6 files changed, 340 insertions(+), 29 deletions(-) create mode 100644 src/query/src/query_engine/runtime.rs diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index dfba397a14..f72bb04da7 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -61,8 +61,8 @@ use tokio::sync::Notify; use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, - GetMetadataSnafu, MissingCacheSnafu, MissingNodeIdSnafu, OpenLogStoreSnafu, Result, - ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, + DataFusionSnafu, GetMetadataSnafu, MissingCacheSnafu, MissingNodeIdSnafu, OpenLogStoreSnafu, + Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver, @@ -432,7 +432,7 @@ impl DatanodeBuilder { ) -> Result { let opts: &DatanodeOptions = &self.opts; - let query_engine_factory = QueryEngineFactory::new_with_plugins( + let query_engine_factory = QueryEngineFactory::try_new_with_plugins( // query engine in datanode only executes plan with resolved table source. DummyCatalogManager::arc(), None, @@ -443,7 +443,8 @@ impl DatanodeBuilder { false, self.plugins.clone(), opts.query.clone(), - ); + ) + .context(DataFusionSnafu)?; let query_engine = query_engine_factory.query_engine(); let table_provider_factory = self diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index 97a61af90b..c44c9e94a3 100644 --- a/src/flow/src/server.rs +++ b/src/flow/src/server.rs @@ -57,8 +57,8 @@ use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef}; use crate::adapter::{FlowStreamingEngineRef, create_worker}; use crate::batching_mode::engine::BatchingEngine; use crate::error::{ - CacheRequiredSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu, - StartServerSnafu, UnexpectedSnafu, to_status_with_last_err, + CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, + ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err, }; use crate::heartbeat::HeartbeatTask; use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS}; @@ -384,7 +384,7 @@ impl FlownodeBuilder { pub async fn build(mut self) -> Result { // TODO(discord9): does this query engine need those? - let query_engine_factory = QueryEngineFactory::new_with_plugins( + let query_engine_factory = QueryEngineFactory::try_new_with_plugins( // query engine in flownode is only used for translate plan with resolved table source. self.catalog_manager.clone(), None, @@ -395,7 +395,10 @@ impl FlownodeBuilder { false, Default::default(), self.opts.query.clone(), - ); + ) + .context(DatafusionSnafu { + context: "Failed to build query engine", + })?; let manager = Arc::new( self.build_manager(query_engine_factory.query_engine()) .await?, diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 3fa7fca33f..37bfc3013b 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -45,7 +45,7 @@ use query::QueryEngineFactory; use query::region_query::RegionQueryHandlerFactoryRef; use snafu::{OptionExt, ResultExt}; -use crate::error::{self, ExternalSnafu, Result}; +use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result}; use crate::events::EventHandlerImpl; use crate::frontend::FrontendOptions; use crate::heartbeat::frontend_peer_addr; @@ -242,7 +242,7 @@ impl FrontendBuilder { let mut query_options = self.options.query.clone(); query_options.enable_per_region_metrics = self.options.logging.enable_per_region_metrics; - let query_engine = QueryEngineFactory::new_with_plugins( + let query_engine = QueryEngineFactory::try_new_with_plugins( self.catalog_manager.clone(), Some(partition_manager.clone()), Some(region_query_handler.clone()), @@ -253,6 +253,7 @@ impl FrontendBuilder { plugins.clone(), query_options, ) + .context(DataFusionSnafu)? .query_engine(); let frontend_peer_addr = frontend_peer_addr(&self.options); diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index 1415998486..c512b9c1e5 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -15,6 +15,7 @@ mod context; mod default_serializer; pub mod options; +pub mod runtime; mod state; use std::any::Any; use std::sync::Arc; @@ -41,6 +42,9 @@ use crate::error::Result; use crate::options::QueryOptions; use crate::planner::LogicalPlanner; pub use crate::query_engine::context::QueryEngineContext; +pub use crate::query_engine::runtime::{ + DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, QueryRuntimeProviderRef, +}; pub use crate::query_engine::state::QueryEngineState; use crate::region_query::RegionQueryHandlerRef; @@ -113,7 +117,28 @@ impl QueryEngineFactory { with_dist_planner: bool, options: QueryOptions, ) -> Self { - Self::new_with_plugins( + Self::try_new( + catalog_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + options, + ) + .expect("Failed to build query engine factory") + } + + pub fn try_new( + catalog_manager: CatalogManagerRef, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + options: QueryOptions, + ) -> datafusion::error::Result { + Self::try_new_with_plugins( catalog_manager, None, region_query_handler, @@ -138,7 +163,33 @@ impl QueryEngineFactory { plugins: Plugins, options: QueryOptions, ) -> Self { - let state = Arc::new(QueryEngineState::new( + Self::try_new_with_plugins( + catalog_manager, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + plugins, + options, + ) + .expect("Failed to build query engine factory") + } + + #[allow(clippy::too_many_arguments)] + pub fn try_new_with_plugins( + catalog_manager: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + plugins: Plugins, + options: QueryOptions, + ) -> datafusion::error::Result { + let state = Arc::new(QueryEngineState::try_new( catalog_manager, partition_rule_manager, region_query_handler, @@ -148,10 +199,10 @@ impl QueryEngineFactory { with_dist_planner, plugins.clone(), options, - )); + )?); let query_engine = Arc::new(DatafusionQueryEngine::new(state, plugins)); register_functions(&query_engine); - Self { query_engine } + Ok(Self { query_engine }) } pub fn query_engine(&self) -> QueryEngineRef { diff --git a/src/query/src/query_engine/runtime.rs b/src/query/src/query_engine/runtime.rs new file mode 100644 index 0000000000..3169db4e83 --- /dev/null +++ b/src/query/src/query_engine/runtime.rs @@ -0,0 +1,80 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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 std::sync::Arc; + +use datafusion::error::Result as DfResult; +use datafusion::execution::context::SessionConfig; +use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + +use crate::options::QueryOptions; +use crate::query_engine::state::MetricsMemoryPool; + +/// Reference-counted query runtime provider. +pub type QueryRuntimeProviderRef = Arc; + +/// Context for building query runtime components. +#[derive(Clone, Copy)] +#[non_exhaustive] +pub struct QueryRuntimeContext<'a> { + /// Query options used by the query engine. + pub query_options: &'a QueryOptions, + /// Resolved memory pool size in bytes. + pub resolved_memory_pool_size: usize, +} + +impl<'a> QueryRuntimeContext<'a> { + /// Creates a new query runtime context. + pub fn new(query_options: &'a QueryOptions, resolved_memory_pool_size: usize) -> Self { + Self { + query_options, + resolved_memory_pool_size, + } + } +} + +/// Provides DataFusion session and runtime setup for the query engine. +pub trait QueryRuntimeProvider: Send + Sync + 'static { + /// Configures the DataFusion session config before building the session state. + fn configure_session_config(&self, _ctx: QueryRuntimeContext<'_>, _config: &mut SessionConfig) { + } + + /// Builds the DataFusion runtime environment. + fn build_runtime_env( + &self, + _ctx: QueryRuntimeContext<'_>, + builder: RuntimeEnvBuilder, + ) -> DfResult> { + builder.build().map(Arc::new) + } +} + +/// Default query runtime provider. +#[derive(Debug, Default)] +pub struct DefaultQueryRuntimeProvider; + +impl DefaultQueryRuntimeProvider { + /// Creates a default DataFusion runtime environment builder. + pub fn runtime_env_builder(ctx: QueryRuntimeContext<'_>) -> RuntimeEnvBuilder { + if ctx.resolved_memory_pool_size > 0 { + RuntimeEnvBuilder::new().with_memory_pool(Arc::new(MetricsMemoryPool::new( + ctx.resolved_memory_pool_size, + ))) + } else { + RuntimeEnvBuilder::new() + } + } +} + +impl QueryRuntimeProvider for DefaultQueryRuntimeProvider {} diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index e86c5916b3..50898aaa58 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -38,7 +38,6 @@ use datafusion::execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, TrackConsumersPool, }; -use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::optimizer::PhysicalOptimizer; use datafusion::physical_optimizer::sanity_checker::SanityCheckPlan; @@ -80,6 +79,9 @@ use crate::optimizer::windowed_sort::WindowedSortPhysicalRule; use crate::options::QueryOptions as QueryOptionsNew; use crate::query_engine::DefaultSerializer; use crate::query_engine::options::QueryOptions; +use crate::query_engine::runtime::{ + DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProviderRef, +}; use crate::range_select::planner::RangeSelectPlanner; use crate::region_query::RegionQueryHandlerRef; @@ -118,18 +120,37 @@ impl QueryEngineState { plugins: Plugins, options: QueryOptionsNew, ) -> Self { + Self::try_new( + catalog_list, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + plugins, + options, + ) + .expect("Failed to build query engine state") + } + + #[allow(clippy::too_many_arguments)] + pub fn try_new( + catalog_list: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + plugins: Plugins, + options: QueryOptionsNew, + ) -> DfResult { let total_memory = get_total_memory_bytes().max(0) as u64; let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize; - let runtime_env = if memory_pool_size > 0 { - Arc::new( - RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(MetricsMemoryPool::new(memory_pool_size))) - .build() - .expect("Failed to build RuntimeEnv"), - ) - } else { - Arc::new(RuntimeEnv::default()) - }; + let runtime_provider = plugins + .get::() + .unwrap_or_else(|| Arc::new(DefaultQueryRuntimeProvider)); let mut session_config = SessionConfig::new().with_create_default_catalog_and_schema(false); if options.parallelism > 0 { session_config = session_config.with_target_partitions(options.parallelism); @@ -150,6 +171,11 @@ impl QueryEngineState { .execution .skip_physical_aggregate_schema_check = true; + let runtime_context = QueryRuntimeContext::new(&options, memory_pool_size); + runtime_provider.configure_session_config(runtime_context, &mut session_config); + let runtime_builder = DefaultQueryRuntimeProvider::runtime_env_builder(runtime_context); + let runtime_env = runtime_provider.build_runtime_env(runtime_context, runtime_builder)?; + // Apply extension rules let mut extension_rules = Vec::new(); @@ -246,7 +272,7 @@ impl QueryEngineState { let df_context = SessionContext::new_with_state(session_state); register_function_aliases(&df_context); - Self { + Ok(Self { df_context, catalog_manager: catalog_list, dyn_filter_registry_manager: Arc::new(DynFilterRegistryManager::default()), @@ -260,7 +286,7 @@ impl QueryEngineState { extension_rules, plugins, scalar_functions: Arc::new(RwLock::new(HashMap::new())), - } + }) } fn remove_physical_optimizer_rule( @@ -545,7 +571,7 @@ impl DfQueryPlanner { /// This wrapper intercepts all memory pool operations and updates /// Prometheus metrics for monitoring query memory usage and rejections. #[derive(Debug)] -struct MetricsMemoryPool { +pub(super) struct MetricsMemoryPool { inner: Arc>, } @@ -553,7 +579,7 @@ impl MetricsMemoryPool { // Number of top memory consumers to report in OOM error messages const TOP_CONSUMERS_TO_REPORT: usize = 5; - fn new(limit: usize) -> Self { + pub(super) fn new(limit: usize) -> Self { Self { inner: Arc::new(TrackConsumersPool::new( GreedyMemoryPool::new(limit), @@ -611,13 +637,25 @@ impl MemoryPool for MetricsMemoryPool { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicBool, Ordering}; + use common_base::Plugins; + use common_base::memory_limit::MemoryLimit; + use common_base::readable_size::ReadableSize; + use datafusion::error::DataFusionError; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryLimit as DfMemoryLimit}; + use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use session::context::QueryContext; use super::*; use crate::options::QueryOptions; + use crate::query_engine::runtime::{QueryRuntimeProvider, QueryRuntimeProviderRef}; fn new_query_engine_state() -> QueryEngineState { + new_query_engine_state_with(Plugins::default(), QueryOptions::default()) + } + + fn new_query_engine_state_with(plugins: Plugins, options: QueryOptions) -> QueryEngineState { QueryEngineState::new( catalog::memory::new_memory_catalog_manager().unwrap(), None, @@ -626,9 +664,146 @@ mod tests { None, None, false, + plugins, + options, + ) + } + + struct TestRuntimeProvider { + build_called: AtomicBool, + configure_called: AtomicBool, + } + + impl TestRuntimeProvider { + fn new() -> Self { + Self { + build_called: AtomicBool::new(false), + configure_called: AtomicBool::new(false), + } + } + } + + impl QueryRuntimeProvider for TestRuntimeProvider { + fn configure_session_config( + &self, + ctx: QueryRuntimeContext<'_>, + config: &mut SessionConfig, + ) { + assert_eq!(ctx.resolved_memory_pool_size, 1024); + self.configure_called.store(true, Ordering::SeqCst); + *config = config.clone().with_target_partitions(7); + } + + fn build_runtime_env( + &self, + ctx: QueryRuntimeContext<'_>, + builder: RuntimeEnvBuilder, + ) -> DfResult> { + assert_eq!(ctx.resolved_memory_pool_size, 1024); + self.build_called.store(true, Ordering::SeqCst); + builder + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2048))) + .build() + .map(Arc::new) + } + } + + struct ErrorRuntimeProvider; + + impl QueryRuntimeProvider for ErrorRuntimeProvider { + fn build_runtime_env( + &self, + _ctx: QueryRuntimeContext<'_>, + _builder: RuntimeEnvBuilder, + ) -> DfResult> { + Err(DataFusionError::Execution("runtime provider error".into())) + } + } + + #[test] + fn query_runtime_default_provider_keeps_bounded_memory_pool() { + let state = new_query_engine_state_with( Plugins::default(), + QueryOptions { + memory_pool_size: MemoryLimit::Size(ReadableSize(1024)), + ..Default::default() + }, + ); + + assert!(matches!( + state + .session_state() + .runtime_env() + .memory_pool + .memory_limit(), + DfMemoryLimit::Finite(1024) + )); + } + + #[test] + fn query_runtime_provider_from_plugins_builds_runtime_env() { + let plugins = Plugins::default(); + let provider = Arc::new(TestRuntimeProvider::new()); + plugins.insert::(provider.clone()); + + let state = new_query_engine_state_with( + plugins, + QueryOptions { + memory_pool_size: MemoryLimit::Size(ReadableSize(1024)), + ..Default::default() + }, + ); + + assert!(provider.build_called.load(Ordering::SeqCst)); + assert!(matches!( + state + .session_state() + .runtime_env() + .memory_pool + .memory_limit(), + DfMemoryLimit::Finite(2048) + )); + } + + #[test] + fn query_runtime_provider_from_plugins_configures_session_config() { + let plugins = Plugins::default(); + let provider = Arc::new(TestRuntimeProvider::new()); + plugins.insert::(provider.clone()); + + let state = new_query_engine_state_with( + plugins, + QueryOptions { + memory_pool_size: MemoryLimit::Size(ReadableSize(1024)), + ..Default::default() + }, + ); + + assert!(provider.configure_called.load(Ordering::SeqCst)); + assert_eq!(7, state.session_state().config().target_partitions()); + } + + #[test] + fn query_runtime_provider_error_is_returned_by_try_new() { + let plugins = Plugins::default(); + plugins.insert::(Arc::new(ErrorRuntimeProvider)); + + let err = QueryEngineState::try_new( + catalog::memory::new_memory_catalog_manager().unwrap(), + None, + None, + None, + None, + None, + false, + plugins, QueryOptions::default(), ) + .unwrap_err(); + + assert!( + matches!(err, DataFusionError::Execution(message) if message == "runtime provider error") + ); } #[test]