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 common_error::ext::BoxedError;
31use common_memory_manager::MemoryGuard;
32use lazy_static::lazy_static;
33use object_pool::Pool;
34use snafu::{IntoError, ResultExt};
35
36use crate::error;
37use crate::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest};
38use crate::prom_remote_write::row_builder::TablesBuilder;
39use crate::prom_remote_write::validation::PromValidationMode;
40use crate::prom_store::{
41    ChargedBuffer, MAX_DECOMPRESSED_REQUEST_SIZE, snappy_decompress_limited,
42    zstd_decompress_limited,
43};
44use crate::request_memory_limiter::ServerMemoryLimiter;
45use crate::request_memory_metrics::RequestMemoryMetrics;
46
47/// Prometheus remote write protocol versions, also used as the `version` label
48/// of the remote write metrics.
49pub const REMOTE_WRITE_V1_VERSION: &str = "1.0";
50pub const REMOTE_WRITE_V2_VERSION: &str = "2.0";
51
52lazy_static! {
53    static ref PROM_WRITE_REQUEST_POOL: Pool<PromWriteRequest<'static>> =
54        Pool::new(256, PromWriteRequest::default);
55}
56
57/// Decompresses a remote-write body with the codec selected by `is_zstd`,
58/// bounded by the hard decoded-size cap and charged to the aggregate
59/// request-memory limiter.
60///
61/// Due to vmagent's limitation there is a chance that vmagent sends the
62/// content-encoding header wrong, so decoding falls back to the other codec
63/// when the first attempt fails on malformed input. Size-limit and quota
64/// errors are final: the body either expanded beyond the cap under the
65/// declared codec, or it could not be admitted, and the other codec cannot
66/// produce a different outcome.
67///
68/// see <https://github.com/VictoriaMetrics/VictoriaMetrics/issues/5301>
69/// see <https://github.com/GreptimeTeam/greptimedb/issues/3929>
70pub(crate) async fn decompress_remote_write_body(
71    is_zstd: bool,
72    body: &[u8],
73    limiter: &ServerMemoryLimiter,
74    body_reserved_bytes: u64,
75) -> crate::error::Result<ChargedBuffer> {
76    let decompress = |is_zstd: bool| async move {
77        if is_zstd {
78            zstd_decompress_limited(
79                body,
80                MAX_DECOMPRESSED_REQUEST_SIZE,
81                limiter,
82                body_reserved_bytes,
83            )
84            .await
85        } else {
86            snappy_decompress_limited(
87                body,
88                MAX_DECOMPRESSED_REQUEST_SIZE,
89                limiter,
90                body_reserved_bytes,
91            )
92            .await
93        }
94    };
95
96    match decompress(is_zstd).await {
97        Ok(buf) => Ok(buf),
98        Err(e) => {
99            if matches!(
100                e,
101                crate::error::Error::DecompressSnappyPromRemoteRequest { .. }
102                    | crate::error::Error::DecompressZstdPromRemoteRequest { .. }
103            ) {
104                decompress(!is_zstd).await
105            } else {
106                Err(e)
107            }
108        }
109    }
110}
111
112pub async fn decode_remote_write_request(
113    is_zstd: bool,
114    body: Bytes,
115    prom_validation_mode: PromValidationMode,
116    processor: &mut PromSeriesProcessor,
117    limiter: &ServerMemoryLimiter,
118) -> crate::error::Result<(
119    TablesBuilder<'static>,
120    Vec<MemoryGuard<RequestMemoryMetrics>>,
121)> {
122    let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED
123        .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION])
124        .start_timer();
125
126    // Holds the memory permits for the decompressed bytes; the caller must
127    // keep them alive as long as the returned `TablesBuilder`, which retains
128    // the decompressed buffer as its raw data.
129    let buf = decompress_remote_write_body(is_zstd, &body[..], limiter, 0).await?;
130    // Decompression copied the payload out, so the compressed body is no longer needed.
131    drop(body);
132    let (data, guards) = buf.into_parts();
133    let mut request = PROM_WRITE_REQUEST_POOL.pull(PromWriteRequest::default);
134
135    request
136        .decode(data, prom_validation_mode, processor)
137        .context(error::DecodePromRemoteRequestSnafu)?;
138    let tables = std::mem::take(&mut request.table_data);
139    Ok((tables, guards))
140}
141
142/// Maps memory wait timeouts to 5xx responses so Remote Write senders retry.
143pub(crate) fn map_memory_acquire_timeout(error: error::Error) -> error::Error {
144    match error {
145        error::Error::MemoryLimitExceeded {
146            source: source @ common_memory_manager::Error::MemoryAcquireTimeout { .. },
147            ..
148        } => error::OtherSnafu.into_error(BoxedError::new(source)),
149        error => error,
150    }
151}
servers/prom_remote_write/mod.rs

servers/prom_remote_write/
mod.rs
