Skip to main content

frontend/service_config/
prom_store.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::num::NonZeroUsize;
16use std::time::Duration;
17
18use common_batcher::flush_policy::timing::TimingFlushPolicy;
19use serde::{Deserialize, Serialize};
20use servers::prom_remote_write::validation::PromValidationMode;
21use tokio::sync::Semaphore;
22
23#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
24pub struct PromStoreOptions {
25    pub enable: bool,
26    pub with_metric_engine: bool,
27    /// Validation mode while decoding Prometheus remote write requests.
28    #[serde(default)]
29    pub prom_validation_mode: PromValidationMode,
30    #[serde(default, with = "humantime_serde")]
31    pub pending_rows_flush_interval: Duration,
32    #[serde(default = "default_max_batch_rows")]
33    pub max_batch_rows: usize,
34    #[serde(default = "default_max_concurrent_flushes")]
35    pub max_concurrent_flushes: usize,
36    #[serde(default = "default_worker_channel_capacity")]
37    pub worker_channel_capacity: usize,
38    #[serde(default = "default_max_inflight_requests")]
39    pub max_inflight_requests: usize,
40    /// Maximum number of logical-table flow notifications waiting in the shared queue.
41    #[serde(default = "default_flow_notification_queue_capacity")]
42    pub flow_notification_queue_capacity: NonZeroUsize,
43}
44
45fn default_max_batch_rows() -> usize {
46    100_000
47}
48
49fn default_max_concurrent_flushes() -> usize {
50    256
51}
52
53fn default_worker_channel_capacity() -> usize {
54    65_536
55}
56
57fn default_max_inflight_requests() -> usize {
58    3000
59}
60
61fn default_flow_notification_queue_capacity() -> NonZeroUsize {
62    NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN)
63}
64
65impl PromStoreOptions {
66    /// Returns whether batching is enabled and its controls pass construction validation.
67    pub fn pending_rows_batching_enabled(&self) -> bool {
68        self.enable
69            && self.with_metric_engine
70            && TimingFlushPolicy::validate(self.pending_rows_flush_interval)
71            && self.max_batch_rows > 0
72            && [
73                self.max_concurrent_flushes,
74                self.worker_channel_capacity,
75                self.max_inflight_requests,
76                self.flow_notification_queue_capacity.get(),
77            ]
78            .into_iter()
79            .all(|capacity| (1..=Semaphore::MAX_PERMITS).contains(&capacity))
80    }
81}
82
83impl Default for PromStoreOptions {
84    fn default() -> Self {
85        Self {
86            enable: true,
87            with_metric_engine: true,
88            prom_validation_mode: PromValidationMode::Strict,
89            pending_rows_flush_interval: Duration::ZERO,
90            max_batch_rows: default_max_batch_rows(),
91            max_concurrent_flushes: default_max_concurrent_flushes(),
92            worker_channel_capacity: default_worker_channel_capacity(),
93            max_inflight_requests: default_max_inflight_requests(),
94            flow_notification_queue_capacity: default_flow_notification_queue_capacity(),
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use std::time::Duration;
102
103    use super::{PromStoreOptions, PromValidationMode};
104    use crate::service_config::prom_store::{
105        default_flow_notification_queue_capacity, default_max_batch_rows,
106        default_max_concurrent_flushes, default_max_inflight_requests,
107    };
108
109    #[test]
110    fn test_batching_capacity_boundaries() {
111        let max = tokio::sync::Semaphore::MAX_PERMITS;
112        let options = PromStoreOptions {
113            pending_rows_flush_interval: Duration::from_secs(5),
114            max_batch_rows: usize::MAX,
115            max_concurrent_flushes: max,
116            worker_channel_capacity: max,
117            max_inflight_requests: max,
118            flow_notification_queue_capacity: std::num::NonZeroUsize::new(max).unwrap(),
119            ..Default::default()
120        };
121        assert!(options.pending_rows_batching_enabled());
122        let cases = [
123            PromStoreOptions {
124                max_concurrent_flushes: max + 1,
125                ..options.clone()
126            },
127            PromStoreOptions {
128                worker_channel_capacity: max + 1,
129                ..options.clone()
130            },
131            PromStoreOptions {
132                max_inflight_requests: max + 1,
133                ..options.clone()
134            },
135            PromStoreOptions {
136                flow_notification_queue_capacity: std::num::NonZeroUsize::new(max + 1).unwrap(),
137                ..options
138            },
139        ];
140        for options in cases {
141            assert!(!options.pending_rows_batching_enabled());
142        }
143    }
144
145    #[test]
146    fn test_prom_store_options() {
147        let default = PromStoreOptions::default();
148        assert!(default.enable);
149        assert!(default.with_metric_engine);
150        assert_eq!(default.prom_validation_mode, PromValidationMode::Strict);
151        assert_eq!(default.pending_rows_flush_interval, Duration::ZERO);
152        assert_eq!(default.max_batch_rows, default_max_batch_rows());
153        assert_eq!(
154            default.max_concurrent_flushes,
155            default_max_concurrent_flushes()
156        );
157        assert_eq!(default.worker_channel_capacity, 65_536);
158        assert_eq!(
159            default.max_inflight_requests,
160            default_max_inflight_requests()
161        );
162        assert_eq!(
163            default.flow_notification_queue_capacity,
164            default_flow_notification_queue_capacity()
165        );
166    }
167}