From f97c4a172d26471343598bd12a9d8d29a46ab972 Mon Sep 17 00:00:00 2001 From: WenyXu Date: Wed, 16 Sep 2026 08:17:56 +0000 Subject: [PATCH] refactor: relocate the logical table batcher Signed-off-by: WenyXu --- src/frontend/src/instance/prom_store.rs | 2 +- src/frontend/src/server.rs | 8 +++-- src/servers/benches/flush_batch_physical.rs | 4 +-- src/servers/src/batcher.rs | 3 +- .../logical_table.rs} | 32 ++++++++++--------- src/servers/src/http.rs | 4 +-- src/servers/src/http/prom_store.rs | 12 +++---- src/servers/src/lib.rs | 1 - src/servers/src/prom_row_builder.rs | 2 +- tests-integration/src/test_util.rs | 4 +-- 10 files changed, 38 insertions(+), 34 deletions(-) rename src/servers/src/{pending_rows_batcher.rs => batcher/logical_table.rs} (99%) diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index e0b17b1a89e..9aa228715c0 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -42,11 +42,11 @@ use operator::insert::{ use operator::statement::StatementExecutor; use prost::Message; use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; +use servers::batcher::logical_table::PendingRowsSchemaAlterer; use servers::error::{self, AuthSnafu, Result as ServerResult}; use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, collect_plan_metrics}; use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef}; -use servers::pending_rows_batcher::PendingRowsSchemaAlterer; use servers::prom_store; use servers::query_handler::{ PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse, diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 257e5243e89..532b0d08b11 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -24,6 +24,9 @@ use common_base::Plugins; use common_config::Configurable; use common_telemetry::{info, warn}; use meta_client::MetaClientOptions; +use servers::batcher::logical_table::{ + LogicalTablePendingRowsBatcher, pending_rows_batch_sync_enabled, +}; use servers::error::Error as ServerError; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::FlightCraftRef; @@ -38,7 +41,6 @@ use servers::interceptor::LogIngestInterceptorRef; use servers::metrics_handler::MetricsHandler; use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef}; use servers::otel_arrow::OtelArrowServiceHandler; -use servers::pending_rows_batcher::{PendingRowsBatcher, pending_rows_batch_sync_enabled}; use servers::postgres::PostgresServer; use servers::request_memory_limiter::ServerMemoryLimiter; use servers::server::{Server, ServerHandlers}; @@ -132,7 +134,7 @@ where let prom_store = effective_prom_store_options(opts); if prom_store.enable { let pending_rows_batcher = if prom_store.with_metric_engine { - PendingRowsBatcher::try_new( + LogicalTablePendingRowsBatcher::try_new( self.instance.partition_manager().clone(), self.instance.node_manager().clone(), self.instance.catalog_manager().clone(), @@ -715,7 +717,7 @@ mod tests { #[test] fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() { - // Mirrors the conditions under which `PendingRowsBatcher::try_new` + // Mirrors the conditions under which `LogicalTablePendingRowsBatcher::try_new` // returns `None`; in these cases no request can wait for a pending-row // flush, so the timeout must not be raised. type KnobMutator = fn(&mut FrontendOptions); diff --git a/src/servers/benches/flush_batch_physical.rs b/src/servers/benches/flush_batch_physical.rs index 45a3f37c780..be992692d13 100644 --- a/src/servers/benches/flush_batch_physical.rs +++ b/src/servers/benches/flush_batch_physical.rs @@ -28,11 +28,11 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use partition::error::Result as PartitionResult; use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask}; -use servers::error::{self, Result}; -use servers::pending_rows_batcher::{ +use servers::batcher::logical_table::{ PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, RecordBatchWithTsIdx, TableBatch, flush_batch_physical, }; +use servers::error::{self, Result}; use store_api::storage::RegionId; use table::test_util::table_info::test_table_info; use tokio::runtime::Runtime; diff --git a/src/servers/src/batcher.rs b/src/servers/src/batcher.rs index c87ebd2eff4..10f8f558eb3 100644 --- a/src/servers/src/batcher.rs +++ b/src/servers/src/batcher.rs @@ -12,10 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Ordinary-table batching and shared notification components. +//! Logical-table and ordinary-table batching implementations. mod flow_notifier; mod flow_sender; +pub mod logical_table; pub mod table; #[cfg(test)] diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/batcher/logical_table.rs similarity index 99% rename from src/servers/src/pending_rows_batcher.rs rename to src/servers/src/batcher/logical_table.rs index 79ae645b98b..66c66322960 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/batcher/logical_table.rs @@ -346,7 +346,7 @@ fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { } /// Prometheus remote write pending rows batcher. -pub struct PendingRowsBatcher { +pub struct LogicalTablePendingRowsBatcher { workers: Arc>, flush_interval: Duration, flush_policy: TimingFlushPolicy, @@ -363,7 +363,7 @@ pub struct PendingRowsBatcher { shutdown: broadcast::Sender<()>, } -impl PendingRowsBatcher { +impl LogicalTablePendingRowsBatcher { #[allow(clippy::too_many_arguments)] pub fn try_new( partition_manager: PartitionRuleManagerRef, @@ -832,6 +832,7 @@ impl PendingRowsBatcher { Ok(aligned_batches) } + async fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker { let (tx, receiver) = self .workers @@ -872,7 +873,7 @@ impl PendingRowsBatcher { } } -impl Drop for PendingRowsBatcher { +impl Drop for LogicalTablePendingRowsBatcher { fn drop(&mut self) { let _ = self.shutdown.send(()); } @@ -1887,22 +1888,22 @@ mod tests { use tokio::sync::{Notify, Semaphore, broadcast, mpsc, oneshot}; use tokio::time::{advance, sleep}; - use crate::error; - use crate::metrics::FLOW_NOTIFICATION_DROPPED; - use crate::pending_rows_batcher::{ + use crate::batcher::logical_table::{ BatchKey, Error, FlushBatch, FlushLimiter, FlushRegionWrite, FlushTrigger, FlushWaiter, - Notifier, PendingBatch, PendingCore, PendingRowsBatcher, PhysicalFlushCatalogProvider, - PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, - PlannedRegionBatch, RecordBatchWithTsIdx, RequestLimiter, ResolvedRegionBatch, TableBatch, - TimingFlushPolicy, WorkerCommand, WorkerRegistry, batch_key_from_ctx, columns_taxonomy, - drain_batch, encode_region_write_requests, extract_timestamps, flush_batch, - flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp, + LogicalTablePendingRowsBatcher, Notifier, PendingBatch, PendingCore, + PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, + PhysicalTableMetadata, PlannedRegionBatch, RecordBatchWithTsIdx, RequestLimiter, + ResolvedRegionBatch, TableBatch, TimingFlushPolicy, WorkerCommand, WorkerRegistry, + batch_key_from_ctx, columns_taxonomy, drain_batch, encode_region_write_requests, extract_timestamps, + flush_batch, flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp, notify_flow_dirty_windows_after_flush, notify_waiters, plan_region_batches, remove_worker_if_same_channel, should_close_worker_on_idle_timeout, should_dispatch_concurrently, start_flow_notification_worker, start_worker, strip_partition_columns_from_batch, transform_logical_batches_to_physical, try_enqueue_flow_notification, }; + use crate::error; + use crate::metrics::FLOW_NOTIFICATION_DROPPED; use crate::prom_row_builder::rows_to_aligned_record_batch; fn mock_rows(row_count: usize, schema_name: &str) -> Rows { @@ -2127,7 +2128,7 @@ mod tests { #[test] fn test_flow_notification_queue_drops_when_full() { let (tx, mut rx) = Notifier::try_new(1).unwrap(); - let notification = |table_id| crate::pending_rows_batcher::FlowNotification { + let notification = |table_id| crate::batcher::logical_table::FlowNotification { table_id, timestamps: vec![table_id as i64], }; @@ -2334,7 +2335,8 @@ mod tests { ], }; - let (table_rows, total_rows) = PendingRowsBatcher::collect_non_empty_table_rows(requests); + let (table_rows, total_rows) = + LogicalTablePendingRowsBatcher::collect_non_empty_table_rows(requests); assert_eq!(2, total_rows); assert_eq!(1, table_rows.len()); @@ -2601,7 +2603,7 @@ mod tests { fn mock_flow_notification_sender( cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, - ) -> Notifier { + ) -> Notifier { let (tx, rx) = Notifier::try_new(16).unwrap(); start_flow_notification_worker(rx, cache, node_manager); tx diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 27411a5f8a6..fd82079c016 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -53,6 +53,7 @@ use tower_http::trace::TraceLayer; use self::authorize::AuthState; use self::result::table_result::TableResponse; +use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; use crate::elasticsearch; use crate::error::{ AddressBindSnafu, AlreadyStartedSnafu, Error, InternalIoSnafu, InvalidHeaderValueSnafu, Result, @@ -74,7 +75,6 @@ use crate::http::result::null_result::NullResponse; use crate::interceptor::LogIngestInterceptorRef; use crate::metrics::http_metrics_layer; use crate::metrics_handler::MetricsHandler; -use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prometheus_handler::PrometheusHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::query_handler::{ @@ -727,7 +727,7 @@ impl HttpServerBuilder { prom_store_with_metric_engine: bool, prom_validation_mode: PromValidationMode, experimental_enable_prometheus_native_histogram: bool, - pending_rows_batcher: Option>, + pending_rows_batcher: Option>, ) -> Self { let state = PromStoreState { prom_store_handler: handler, diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index cb7b6307c4b..0a7567eee7e 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -40,12 +40,12 @@ use table::requests::{ SOURCE_PROMETHEUS, }; +use crate::batcher::logical_table::LogicalTablePendingRowsBatcher; use crate::error::{self, InternalSnafu, PipelineSnafu, Result}; use crate::http::extractor::PipelineInfo; use crate::http::header::{ CONTENT_TYPE_PROTOBUF_STR, GREPTIME_DB_HEADER_METRICS, write_cost_header_map, }; -use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prom_remote_write::decode::PromSeriesProcessor; use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; @@ -75,7 +75,7 @@ pub struct PromStoreState { pub prom_store_with_metric_engine: bool, pub prom_validation_mode: PromValidationMode, pub experimental_enable_prometheus_native_histogram: bool, - pub pending_rows_batcher: Option>, + pub pending_rows_batcher: Option>, } #[derive(Debug, Serialize, Deserialize)] @@ -358,9 +358,9 @@ trait PromWriteBatcher: Send + Sync { } #[async_trait] -impl PromWriteBatcher for PendingRowsBatcher { +impl PromWriteBatcher for LogicalTablePendingRowsBatcher { async fn submit(&self, requests: RowInsertRequests, ctx: QueryContextRef) -> Result { - PendingRowsBatcher::submit(self, requests, ctx).await + LogicalTablePendingRowsBatcher::submit(self, requests, ctx).await } } @@ -387,7 +387,7 @@ async fn preflight_prometheus_rows( /// sample/histogram headers even when a later table write fails. async fn write_prometheus_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, - pending_rows_batcher: Option>, + pending_rows_batcher: Option>, prom_store_with_metric_engine: bool, mut batches: Vec, ) -> std::result::Result { @@ -453,7 +453,7 @@ async fn write_prometheus_rows_with_progress( async fn write_prometheus_v2_rows_with_progress( prom_store_handler: PromStoreProtocolHandlerRef, - pending_rows_batcher: Option>, + pending_rows_batcher: Option>, prom_store_with_metric_engine: bool, sample_batches: Vec, histogram_batches: Vec, diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index ac3ca064dfb..1106f89d3e9 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -38,7 +38,6 @@ pub mod mysql; pub mod opentsdb; pub mod otel_arrow; pub mod otlp; -pub mod pending_rows_batcher; mod pipeline; pub mod postgres; pub mod prom_remote_write; diff --git a/src/servers/src/prom_row_builder.rs b/src/servers/src/prom_row_builder.rs index aad12814683..8b4e35bcd32 100644 --- a/src/servers/src/prom_row_builder.rs +++ b/src/servers/src/prom_row_builder.rs @@ -35,9 +35,9 @@ use datatypes::data_type::DataType; use datatypes::prelude::ConcreteDataType; use snafu::{OptionExt, ResultExt, ensure}; +use crate::batcher::logical_table::RecordBatchWithTsIdx; use crate::error; use crate::error::Result; -use crate::pending_rows_batcher::RecordBatchWithTsIdx; /// Extract timestamp, field, and tag column names from a logical region schema. fn unzip_logical_region_schema( diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index e0a749b813f..562bd48cdb0 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -58,6 +58,7 @@ use object_store::config::{ use object_store::services::{Azblob, Gcs, Oss, S3}; use object_store::test_util::TempFolder; use object_store::{AzblobConnection, GcsConnection, ObjectStore, OssConnection, S3Connection}; +use servers::batcher::logical_table::LogicalTablePendingRowsBatcher; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcOptions, GrpcServer, GrpcServerConfig}; @@ -65,7 +66,6 @@ use servers::http::{HttpOptions, HttpServerBuilder}; use servers::metrics_handler::MetricsHandler; use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef}; use servers::otel_arrow::OtelArrowServiceHandler; -use servers::pending_rows_batcher::PendingRowsBatcher; use servers::postgres::PostgresServer; use servers::prom_remote_write::validation::PromValidationMode; use servers::query_handler::sql::SqlQueryHandler; @@ -867,7 +867,7 @@ pub fn build_test_prom_server( // Mirror the production wiring at `frontend::server`: build the batcher from the // instance's managers. A short flush interval keeps the test responsive. let pending_rows_batcher = if enable_batcher { - PendingRowsBatcher::try_new( + LogicalTablePendingRowsBatcher::try_new( frontend_ref.partition_manager().clone(), frontend_ref.node_manager().clone(), frontend_ref.catalog_manager().clone(),