servers/prom_remote_write/
mod.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//! Prometheus remote write support.
16//!
17//! This module groups validation, row building, and protobuf decoding for
18//! the Prometheus remote write API.
19
20pub mod decode;
21pub(crate) mod row_builder;
22pub(crate) mod types;
23#[cfg(any(test, feature = "testing"))]
24pub mod v2;
25#[cfg(not(any(test, feature = "testing")))]
26pub(crate) mod v2;
27pub mod validation;
28
29use bytes::Bytes;
30use lazy_static::lazy_static;
31use object_pool::Pool;
32use snafu::ResultExt;
33
34use crate::error;
35use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest};
36use crate::prom_remote_write::row_builder::TablesBuilder;
37use crate::prom_remote_write::validation::PromValidationMode;
38use crate::prom_store::{snappy_decompress, zstd_decompress};
39
40/// Prometheus remote write protocol versions, also used as the `version` label
41/// of the remote write metrics.
42pub const REMOTE_WRITE_V1_VERSION: &str = "1.0";
43pub const REMOTE_WRITE_V2_VERSION: &str = "2.0";
44
45lazy_static! {
46    static ref PROM_WRITE_REQUEST_POOL: Pool<PromWriteRequest<'static>> =
47        Pool::new(256, PromWriteRequest::default);
48}
49
50pub fn try_decompress(is_zstd: bool, body: &[u8]) -> crate::error::Result<Vec<u8>> {
51    if is_zstd {
52        zstd_decompress(body)
53    } else {
54        snappy_decompress(body)
55    }
56}
57
58pub fn decode_remote_write_request(
59    is_zstd: bool,
60    body: Bytes,
61    prom_validation_mode: PromValidationMode,
62    processor: &mut PromSeriesProcessor,
63) -> crate::error::Result<TablesBuilder<'static>> {
64    let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED
65        .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION])
66        .start_timer();
67
68    // due to vmagent's limitation, there is a chance that vmagent is
69    // sending content type wrong so we have to apply a fallback with decoding
70    // the content in another method.
71    //
72    // see https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301
73    // see https://github.com/GreptimeTeam/greptimedb/issues/3929
74    let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) {
75        buf
76    } else {
77        // fallback to the other compression method
78        try_decompress(!is_zstd, &body[..])?
79    };
80
81    let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default);
82
83    request
84        .decode(buf, prom_validation_mode, processor)
85        .context(error::DecodePromRemoteRequestSnafu)?;
86    Ok(std::mem::take(&mut request.table_data))
87}
servers/prom_remote_write/mod.rs

servers/prom_remote_write/
mod.rs
