netget/client/prometheus_remote_write/
transport.rs

1use crate::server::prometheus_remote_write::codec::WriteBatch;
2use anyhow::{ensure, Context, Result};
3use bytes::Bytes;
4use http_body_util::{BodyExt, Full, Limited};
5use hyper::{
6    client::conn::http1,
7    header::{AUTHORIZATION, CONNECTION, CONTENT_ENCODING, CONTENT_TYPE, HOST, USER_AGENT},
8    Request, Uri,
9};
10use hyper_util::rt::TokioIo;
11use serde::Serialize;
12use std::time::Duration;
13pub const IO_TIMEOUT: Duration = Duration::from_secs(10);
14pub const MIN_BACKOFF: Duration = Duration::from_millis(100);
15pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
16pub const MAX_RESPONSE_BYTES: usize = 64 * 1024;
17pub const MAX_RESPONSE_HEADERS: usize = 64;
18pub const MAX_RESPONSE_HEADER_BYTES: usize = 32 * 1024;
19#[derive(Clone)]
20pub struct Origin {
21    pub authority: String,
22    pub connect_addr: String,
23}
24impl Origin {
25    pub fn parse(value: &str) -> Result<Self> {
26        ensure!(
27            !value.contains(['#', '@']),
28            "HTTP origin cannot contain credentials/fragment"
29        );
30        let value = if value.starts_with("http://") {
31            value.to_owned()
32        } else {
33            format!("http://{value}")
34        };
35        let uri: Uri = value.parse().context("invalid HTTP origin")?;
36        ensure!(
37            uri.scheme_str() == Some("http")
38                && uri.query().is_none()
39                && matches!(uri.path(), "" | "/"),
40            "cleartext HTTP origin only, without path/query"
41        );
42        let authority = uri.authority().context("HTTP host required")?;
43        ensure!(authority.as_str().len() <= 1024, "HTTP origin byte limit");
44        Ok(Self {
45            authority: authority.as_str().into(),
46            connect_addr: format!(
47                "{}:{}",
48                authority.host(),
49                authority.port_u16().unwrap_or(9090)
50            ),
51        })
52    }
53}
54#[derive(Clone)]
55pub struct Config {
56    pub origin: Origin,
57    pub token: Option<String>,
58    pub path: String,
59    pub retry_429: bool,
60}
61#[derive(Clone, Debug, Serialize)]
62pub struct WriteResponse {
63    pub series_count: usize,
64    pub sample_count: usize,
65    pub status: u16,
66    pub accepted: bool,
67    pub attempts: u64,
68    pub durable_storage_confirmed: bool,
69}
70pub async fn write(config: Config, batch: WriteBatch, body: Vec<u8>) -> Result<WriteResponse> {
71    let mut attempts = 0u64;
72    let mut backoff = MIN_BACKOFF;
73    loop {
74        attempts = attempts.saturating_add(1);
75        match write_once(&config, &body).await {
76            Ok(status) if (500..600).contains(&status) || (status == 429 && config.retry_429) => {}
77            Ok(status) => {
78                return Ok(WriteResponse {
79                    series_count: batch.series.len(),
80                    sample_count: batch.series.iter().map(|s| s.samples.len()).sum(),
81                    status,
82                    accepted: (200..300).contains(&status),
83                    attempts,
84                    durable_storage_confirmed: false,
85                })
86            }
87            // A transport failure can occur after remote acceptance. Retrying the
88            // identical batch is intentional; no exactly-once/dedup claim is made.
89            Err(_) => {}
90        }
91        // The caller owns this future. Disconnect/removal drops both the current
92        // exchange and backoff; one batch remains bounded in volatile memory only.
93        tokio::time::sleep(backoff).await;
94        backoff = backoff.saturating_mul(2).min(MAX_BACKOFF);
95    }
96}
97pub async fn write_once(config: &Config, body: &[u8]) -> Result<u16> {
98    tokio::time::timeout(IO_TIMEOUT, async {
99        let socket = tokio::net::TcpStream::connect(&config.origin.connect_addr).await?;
100        let (mut sender, connection) = http1::Builder::new()
101            .max_headers(MAX_RESPONSE_HEADERS)
102            .max_buf_size(MAX_RESPONSE_HEADER_BYTES)
103            .handshake(TokioIo::new(socket))
104            .await?;
105        let mut request = Request::builder()
106            .method("POST")
107            .uri(&config.path)
108            .header(HOST, &config.origin.authority)
109            .header(CONNECTION, "close")
110            .header(CONTENT_TYPE, "application/x-protobuf")
111            .header(CONTENT_ENCODING, "snappy")
112            .header(USER_AGENT, concat!("netget/", env!("CARGO_PKG_VERSION")))
113            .header("X-Prometheus-Remote-Write-Version", "0.1.0");
114        if let Some(token) = &config.token {
115            request = request.header(AUTHORIZATION, format!("Bearer {token}"));
116        }
117        let request = request.body(Full::new(Bytes::copy_from_slice(body)))?;
118        let exchange = async {
119            let response = sender.send_request(request).await?;
120            let (parts, body) = response.into_parts();
121            ensure!(
122                parts.headers.len() <= MAX_RESPONSE_HEADERS
123                    && parts
124                        .headers
125                        .iter()
126                        .map(|(k, v)| k.as_str().len() + v.as_bytes().len() + 4)
127                        .sum::<usize>()
128                        <= MAX_RESPONSE_HEADER_BYTES,
129                "response header count/byte limit"
130            );
131            // The v1 response body is reserved and MUST be ignored, including
132            // successful nonempty/binary bodies. Bounded draining validates HTTP
133            // framing without treating text, encodings or Retry-After as schema.
134            Limited::new(body, MAX_RESPONSE_BYTES)
135                .collect()
136                .await
137                .map_err(|_| anyhow::anyhow!("invalid/oversized response body"))?;
138            Ok::<_, anyhow::Error>(parts.status.as_u16())
139        };
140        tokio::pin!(connection);
141        tokio::pin!(exchange);
142        tokio::select! {
143            result = &mut exchange => result,
144            result = &mut connection => { result.context("HTTP driver failed")?; exchange.await }
145        }
146    })
147    .await
148    .context("remote write exchange deadline exceeded")?
149}
netget/client/prometheus_remote_write/transport.rs

netget/client/prometheus_remote_write/
transport.rs
