Skip to main content

frontend/instance/
logical_batcher.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::{Arc, Weak};
16
17use api::v1::ColumnSchema;
18use async_trait::async_trait;
19use servers::batcher::BatchingProtocol;
20use servers::batcher::logical_table::{LogicalTablePendingRowsBatcher, PendingRowsSchemaAlterer};
21use servers::error::{BatcherChannelClosedSnafu, Result};
22use session::context::QueryContextRef;
23use snafu::OptionExt;
24
25use crate::frontend::FrontendOptions;
26use crate::instance::Instance;
27
28impl Instance {
29    /// Initializes the shared logical batcher for enabled ingestion endpoints.
30    /// Disabled endpoints leave it uninitialized so downstream routers can enable it.
31    pub fn init_logical_batcher(self: &Arc<Self>, options: &FrontendOptions) {
32        let options_batcher = options.logical_batcher_options();
33        let enabled = options_batcher
34            .protocols
35            .iter()
36            .any(|protocol| match protocol {
37                BatchingProtocol::Prom => options.prom_store.enable,
38                BatchingProtocol::Otlp => options.otlp.enable,
39                _ => false,
40            });
41        if !options.prom_store.with_metric_engine
42            || !enabled
43            || !options_batcher.pending_rows_batching_enabled()
44        {
45            return;
46        }
47        self.logical_batcher.get_or_init(|| {
48            LogicalTablePendingRowsBatcher::try_new(
49                self.partition_manager().clone(),
50                self.node_manager().clone(),
51                self.inserter.write_dispatcher().clone(),
52                self.catalog_manager().clone(),
53                self.table_flownode_set_cache().clone(),
54                true,
55                Arc::new(LogicalTables(Arc::downgrade(self))),
56                options_batcher.pending_rows_flush_interval,
57                options_batcher.max_batch_rows,
58                options_batcher.max_concurrent_flushes,
59                options_batcher.worker_channel_capacity,
60                options_batcher.max_inflight_requests,
61                options_batcher.flow_notification_queue_capacity,
62            )
63        });
64    }
65
66    pub(crate) fn logical_batcher(&self) -> Option<&Arc<LogicalTablePendingRowsBatcher>> {
67        self.logical_batcher.get().and_then(Option::as_ref)
68    }
69}
70
71// The instance owns the shared batcher; schema preparation must not keep that
72// owner alive through a reference cycle.
73struct LogicalTables(Weak<Instance>);
74
75#[async_trait]
76impl PendingRowsSchemaAlterer for LogicalTables {
77    async fn create_tables_if_missing_batch(
78        &self,
79        catalog: &str,
80        schema: &str,
81        tables: &[(&str, &[ColumnSchema])],
82        with_metric_engine: bool,
83        ctx: QueryContextRef,
84    ) -> Result<()> {
85        self.0
86            .upgrade()
87            .context(BatcherChannelClosedSnafu)?
88            .create_tables_if_missing_batch(catalog, schema, tables, with_metric_engine, ctx)
89            .await
90    }
91
92    async fn add_missing_prom_tag_columns_batch(
93        &self,
94        catalog: &str,
95        schema: &str,
96        tables: &[(&str, &[String])],
97        ctx: QueryContextRef,
98    ) -> Result<()> {
99        self.0
100            .upgrade()
101            .context(BatcherChannelClosedSnafu)?
102            .add_missing_prom_tag_columns_batch(catalog, schema, tables, ctx)
103            .await
104    }
105}