Skip to main content

servers/
exemplar.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
15//! Normalized metric exemplars shared by Remote Write and OTLP conversion.
16
17pub mod query;
18
19use std::collections::{BTreeMap, BTreeSet};
20use std::sync::Arc;
21
22use api::v1::value::ValueData;
23use api::v1::{ColumnDataType, Row, RowInsertRequests, Value};
24use common_catalog::consts::METRICS_EXEMPLARS_TABLE;
25use common_grpc::precision::Precision;
26use common_query::prelude::{
27    EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, greptime_timestamp, greptime_value,
28};
29use datatypes::json;
30use metric_engine::row_modifier::TsidGenerator;
31use snafu::{OptionExt, ensure};
32use store_api::metric_engine_consts::{
33    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
34};
35use store_api::storage::consts::is_internal_column;
36
37use crate::error::{InvalidPromRemoteRequestSnafu, Result};
38use crate::row_writer::{self, MultiTableData};
39
40/// One observation belonging to a query-visible metric series.
41/// Metric names and inherited labels are shared until row construction.
42#[derive(Debug, Clone, PartialEq)]
43pub struct Exemplar {
44    pub metric: Arc<str>,
45    pub series_labels: Arc<BTreeMap<String, String>>,
46    pub timestamp: i64,
47    pub value: f64,
48    pub labels: BTreeMap<String, String>,
49}
50
51/// Exemplars grouped by the source series' optional database routing control.
52pub type ExemplarRequests = BTreeMap<Option<String>, Vec<Exemplar>>;
53
54/// Bounds row capacity after each database's metrics share one exemplar schema.
55pub fn estimated_row_bytes<'a>(
56    requests: impl IntoIterator<Item = &'a [Exemplar]>,
57) -> Option<usize> {
58    requests.into_iter().try_fold(0usize, |total, exemplars| {
59        let mut tags = BTreeSet::new();
60        let mut payload = 0usize;
61        // Shared series data is copied into each output row.
62        for exemplar in exemplars {
63            tags.extend(exemplar.series_labels.keys());
64            payload = payload.checked_add(exemplar.metric.len().checked_mul(4)?)?;
65            for (name, value) in exemplar.series_labels.iter().chain(&exemplar.labels) {
66                // Include owned strings, label-map entries and JSON2 encoding scratch.
67                payload = payload.checked_add(
68                    name.len()
69                        .checked_add(value.len())?
70                        .checked_mul(4)?
71                        .checked_add(256)?,
72                )?;
73            }
74        }
75        // Include the source table ID attached after conversion, and Vec growth.
76        let columns = tags.len().checked_add(6)?;
77        let values = exemplars
78            .len()
79            .checked_mul(columns)?
80            .checked_mul(2 * std::mem::size_of::<Value>())?;
81        let rows = exemplars
82            .len()
83            .checked_mul(2 * std::mem::size_of::<Row>() + std::mem::size_of::<Exemplar>())?;
84        total
85            .checked_add(values)?
86            .checked_add(rows)?
87            .checked_add(columns.checked_mul(512)?)?
88            .checked_add(payload)
89    })
90}
91
92/// Validates labels with the protocol's name policy, preserving duplicates as errors.
93pub fn labels<I>(
94    labels: I,
95    exemplar_owned: bool,
96    validate_name: impl Fn(&[u8]) -> bool,
97) -> Result<BTreeMap<String, String>>
98where
99    I: IntoIterator<Item = (String, String)>,
100{
101    let mut result = BTreeMap::new();
102    let mut length = 0usize;
103    for (name, value) in labels {
104        ensure!(
105            validate_name(name.as_bytes()),
106            InvalidPromRemoteRequestSnafu {
107                msg: format!("invalid exemplar label name: {name}"),
108            }
109        );
110        if exemplar_owned {
111            ensure!(
112                name != json::JSON2_REMAINDER_FIELD_NAME,
113                InvalidPromRemoteRequestSnafu {
114                    msg: "exemplar label name is reserved by JSON2",
115                }
116            );
117        } else {
118            ensure!(
119                ![
120                    greptime_timestamp(),
121                    greptime_value(),
122                    EXEMPLAR_METRIC_COLUMN,
123                    EXEMPLAR_LABELS_COLUMN,
124                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
125                    DATA_SCHEMA_TSID_COLUMN_NAME,
126                ]
127                .contains(&name.as_str())
128                    && !is_internal_column(&name),
129                InvalidPromRemoteRequestSnafu {
130                    msg: format!("inherited exemplar label collides with reserved column: {name}"),
131                }
132            );
133        }
134        length = length
135            .checked_add(name.chars().count())
136            .and_then(|n| n.checked_add(value.chars().count()))
137            .context(InvalidPromRemoteRequestSnafu {
138                msg: "exemplar label length overflow",
139            })?;
140        ensure!(
141            !exemplar_owned || length <= 128,
142            InvalidPromRemoteRequestSnafu {
143                msg: "exemplar label names and values exceed 128 characters",
144            }
145        );
146        ensure!(
147            result.insert(name.clone(), value).is_none(),
148            InvalidPromRemoteRequestSnafu {
149                msg: format!("duplicate exemplar label: {name}"),
150            }
151        );
152    }
153    Ok(result)
154}
155
156/// Converts Remote Write milliseconds without losing overflow information.
157pub fn timestamp_from_millis(timestamp: i64) -> Result<i64> {
158    timestamp
159        .checked_mul(1_000_000)
160        .context(InvalidPromRemoteRequestSnafu {
161            msg: "exemplar timestamp is outside the signed nanosecond range",
162        })
163}
164
165/// Builds ordinary-table rows and retains source targets for authorization.
166pub fn into_rows(exemplars: Vec<Exemplar>) -> Result<(RowInsertRequests, BTreeSet<String>)> {
167    let mut tables = MultiTableData::default();
168    let mut sources = BTreeSet::new();
169    for exemplar in exemplars {
170        ensure!(
171            !exemplar.metric.is_empty() && exemplar.metric.as_ref() != METRICS_EXEMPLARS_TABLE,
172            InvalidPromRemoteRequestSnafu {
173                msg: "invalid or reserved exemplar source metric name"
174            }
175        );
176        sources.insert(exemplar.metric.to_string());
177        let table = tables.get_or_default_table_data(
178            METRICS_EXEMPLARS_TABLE,
179            exemplar.series_labels.len() + 6,
180            1,
181        );
182        let mut row = table.alloc_one_row();
183        let tsid = TsidGenerator::from_sorted_tags(
184            exemplar
185                .series_labels
186                .iter()
187                .map(|(name, value)| (name.as_str(), value.as_str())),
188        );
189        row_writer::write_by_schema(
190            table,
191            std::iter::once((
192                api::v1::helper::tag_column_schema(
193                    DATA_SCHEMA_TSID_COLUMN_NAME,
194                    ColumnDataType::Uint64,
195                ),
196                Some(ValueData::U64Value(tsid)),
197            )),
198            &mut row,
199        )?;
200        row_writer::write_tag(table, EXEMPLAR_METRIC_COLUMN, &exemplar.metric, &mut row)?;
201        row_writer::write_tags(
202            table,
203            exemplar
204                .series_labels
205                .iter()
206                .map(|(name, value)| (name.as_str(), value.clone())),
207            &mut row,
208        )?;
209        row_writer::write_fields(
210            table,
211            std::iter::once((
212                greptime_value().to_string(),
213                ColumnDataType::Float64,
214                Some(ValueData::F64Value(exemplar.value)),
215            )),
216            &mut row,
217        )?;
218        row_writer::write_ts_to_nanos(
219            table,
220            greptime_timestamp(),
221            Some(exemplar.timestamp),
222            Precision::Nanosecond,
223            &mut row,
224        )?;
225        row_writer::write_by_schema(
226            table,
227            std::iter::once((
228                row_writer::build_json2_column_schema(EXEMPLAR_LABELS_COLUMN)?,
229                Some(row_writer::encode_json2(exemplar.labels)?),
230            )),
231            &mut row,
232        )?;
233        table.add_row(row);
234    }
235    Ok((tables.into_row_insert_requests().0, sources))
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::prom_remote_write::validation::validate_label_name;
242
243    #[test]
244    fn test_exemplar_labels_and_timestamp_validation() {
245        assert!(
246            labels(
247                [(String::from("traceID"), String::from("abc"))],
248                true,
249                validate_label_name
250            )
251            .is_ok()
252        );
253        assert!(
254            labels(
255                [
256                    (String::from("job"), String::new()),
257                    (String::from("job"), String::new())
258                ],
259                false,
260                validate_label_name,
261            )
262            .is_err()
263        );
264        assert!(
265            labels(
266                [(EXEMPLAR_METRIC_COLUMN.into(), "x".into())],
267                false,
268                validate_label_name
269            )
270            .is_err()
271        );
272        assert!(labels([("id".into(), "x".repeat(127))], true, validate_label_name).is_err());
273        for name in ["", "request.method", "请求"] {
274            assert!(labels([(name.into(), "x".into())], true, validate_label_name).is_err());
275        }
276        for name in [
277            DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
278            DATA_SCHEMA_TSID_COLUMN_NAME,
279        ] {
280            assert!(labels([(name.into(), "x".into())], false, validate_label_name).is_err());
281        }
282        assert_eq!(timestamp_from_millis(-1).unwrap(), -1_000_000);
283        assert!(timestamp_from_millis(i64::MAX).is_err());
284    }
285
286    #[test]
287    fn test_exemplar_tsid_uses_only_source_labels() {
288        let exemplar = Exemplar {
289            metric: "cpu".into(),
290            series_labels: BTreeMap::from([("job".into(), "api".into())]).into(),
291            timestamp: 1,
292            value: 2.0,
293            labels: BTreeMap::from([("trace_id".into(), "a".into())]),
294        };
295        let mut other = exemplar.clone();
296        other.metric = "memory".into();
297        other.labels.insert("trace_id".into(), "b".into());
298        let (requests, _) = into_rows(vec![exemplar, other]).unwrap();
299        let rows = requests.inserts[0].rows.as_ref().unwrap();
300        let index = rows
301            .schema
302            .iter()
303            .position(|c| c.column_name == DATA_SCHEMA_TSID_COLUMN_NAME)
304            .unwrap();
305        let expected = ValueData::U64Value(TsidGenerator::from_sorted_tags(
306            [("job", "api")].into_iter(),
307        ));
308        assert!(
309            rows.rows
310                .iter()
311                .all(|row| row.values[index].value_data.as_ref() == Some(&expected))
312        );
313    }
314}