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.catalog_manager().clone(),
52                self.table_flownode_set_cache().clone(),
53                true,
54                Arc::new(LogicalTables(Arc::downgrade(self))),
55                options_batcher.pending_rows_flush_interval,
56                options_batcher.max_batch_rows,
57                options_batcher.max_concurrent_flushes,
58                options_batcher.worker_channel_capacity,
59                options_batcher.max_inflight_requests,
60                options_batcher.flow_notification_queue_capacity,
61            )
62        });
63    }
64
65    pub(crate) fn logical_batcher(&self) -> Option<&Arc<LogicalTablePendingRowsBatcher>> {
66        self.logical_batcher.get().and_then(Option::as_ref)
67    }
68}
69
70// The instance owns the shared batcher; schema preparation must not keep that
71// owner alive through a reference cycle.
72struct LogicalTables(Weak<Instance>);
73
74#[async_trait]
75impl PendingRowsSchemaAlterer for LogicalTables {
76    async fn create_tables_if_missing_batch(
77        &self,
78        catalog: &str,
79        schema: &str,
80        tables: &[(&str, &[ColumnSchema])],
81        with_metric_engine: bool,
82        ctx: QueryContextRef,
83    ) -> Result<()> {
84        self.0
85            .upgrade()
86            .context(BatcherChannelClosedSnafu)?
87            .create_tables_if_missing_batch(catalog, schema, tables, with_metric_engine, ctx)
88            .await
89    }
90
91    async fn add_missing_prom_tag_columns_batch(
92        &self,
93        catalog: &str,
94        schema: &str,
95        tables: &[(&str, &[String])],
96        ctx: QueryContextRef,
97    ) -> Result<()> {
98        self.0
99            .upgrade()
100            .context(BatcherChannelClosedSnafu)?
101            .add_missing_prom_tag_columns_batch(catalog, schema, tables, ctx)
102            .await
103    }
104}