servers/http/
opentsdb.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::collections::HashMap;
16use std::sync::Arc;
17
18use axum::body::Bytes;
19use axum::extract::{Query, State};
20use axum::http::StatusCode as HttpStatusCode;
21use axum::{Extension, Json};
22use common_error::ext::ErrorExt;
23use serde::{Deserialize, Serialize};
24use session::context::{Channel, QueryContext};
25use snafu::ResultExt;
26
27use crate::error::{self, Result};
28use crate::opentsdb::codec::DataPoint;
29use crate::query_handler::OpentsdbProtocolHandlerRef;
30
31#[derive(Serialize, Deserialize)]
32#[serde(untagged)]
33enum OneOrMany<T> {
34    One(T),
35    Vec(Vec<T>),
36}
37
38impl<T> From<OneOrMany<T>> for Vec<T> {
39    fn from(from: OneOrMany<T>) -> Self {
40        match from {
41            OneOrMany::One(val) => vec![val],
42            OneOrMany::Vec(vec) => vec,
43        }
44    }
45}
46
47#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
48pub struct DataPointRequest {
49    metric: String,
50    timestamp: i64,
51    value: f64,
52    tags: HashMap<String, String>,
53}
54
55impl From<DataPointRequest> for DataPoint {
56    fn from(request: DataPointRequest) -> Self {
57        let ts_millis = DataPoint::timestamp_to_millis(request.timestamp);
58
59        let tags = request.tags.into_iter().collect::<Vec<(String, String)>>();
60
61        DataPoint::new(request.metric, ts_millis, request.value, tags)
62    }
63}
64
65#[derive(Serialize, Deserialize, Debug)]
66#[serde(untagged)]
67pub enum OpentsdbPutResponse {
68    Empty,
69    Debug(OpentsdbDebuggingResponse),
70}
71
72// Please refer to the OpenTSDB documents of ["api/put"](http://opentsdb.net/docs/build/html/api_http/put.html)
73// for more details.
74#[axum_macros::debug_handler]
75pub async fn put(
76    State(opentsdb_handler): State<OpentsdbProtocolHandlerRef>,
77    Query(params): Query<HashMap<String, String>>,
78    Extension(mut ctx): Extension<QueryContext>,
79    body: Bytes,
80) -> Result<(HttpStatusCode, Json<OpentsdbPutResponse>)> {
81    let summary = params.contains_key("summary");
82    let details = params.contains_key("details");
83
84    let data_point_requests = parse_data_points(body).await?;
85    let data_points = data_point_requests
86        .iter()
87        .map(|point| point.clone().into())
88        .collect::<Vec<_>>();
89
90    ctx.set_channel(Channel::Opentsdb);
91    let mut ctx = Arc::new(ctx);
92
93    if summary || details {
94        opentsdb_handler
95            .preflight(&data_points, ctx.clone())
96            .await?;
97        ctx = operator::insert::admit_write(data_points.len() as u64, &ctx)
98            .await
99            .map_err(common_error::ext::BoxedError::new)
100            .context(error::ExecuteGrpcQuerySnafu)?;
101    }
102
103    let response = if !summary && !details {
104        opentsdb_handler
105            .exec_batch(data_points, ctx.clone())
106            .await?;
107        (HttpStatusCode::NO_CONTENT, Json(OpentsdbPutResponse::Empty))
108    } else {
109        let mut response = OpentsdbDebuggingResponse {
110            success: 0,
111            failed: 0,
112            errors: if details {
113                Some(Vec::with_capacity(data_points.len()))
114            } else {
115                None
116            },
117        };
118
119        for (data_point, request) in data_points.into_iter().zip(data_point_requests) {
120            let result = opentsdb_handler.exec(vec![data_point], ctx.clone()).await;
121            match result {
122                Ok(affected_rows) => response.on_success(affected_rows),
123                Err(e) => response.on_failed(request, e),
124            }
125        }
126        (
127            HttpStatusCode::OK,
128            Json(OpentsdbPutResponse::Debug(response)),
129        )
130    };
131    Ok(response)
132}
133
134async fn parse_data_points(body: Bytes) -> Result<Vec<DataPointRequest>> {
135    let data_points = serde_json::from_slice::<OneOrMany<DataPointRequest>>(&body[..])
136        .context(error::InvalidOpentsdbJsonRequestSnafu)?;
137    Ok(data_points.into())
138}
139
140#[derive(Serialize, Deserialize, Debug)]
141struct OpentsdbDetailError {
142    datapoint: DataPointRequest,
143    error: String,
144}
145
146#[derive(Serialize, Deserialize, Debug)]
147pub struct OpentsdbDebuggingResponse {
148    success: i32,
149    failed: i32,
150    #[serde(skip_serializing_if = "Option::is_none")]
151    errors: Option<Vec<OpentsdbDetailError>>,
152}
153
154impl OpentsdbDebuggingResponse {
155    fn on_success(&mut self, affected_rows: usize) {
156        self.success += affected_rows as i32;
157    }
158
159    fn on_failed(&mut self, datapoint: DataPointRequest, error: impl ErrorExt) {
160        self.failed += 1;
161
162        if let Some(details) = self.errors.as_mut() {
163            let error = OpentsdbDetailError {
164                datapoint,
165                error: error.output_msg(),
166            };
167            details.push(error);
168        };
169    }
170}
171
172#[cfg(test)]
173mod test {
174
175    use std::sync::atomic::{AtomicUsize, Ordering};
176
177    use session::context::QueryContextRef;
178
179    use crate::http::opentsdb::*;
180    use crate::query_handler::OpentsdbProtocolHandler;
181
182    #[derive(Default)]
183    struct CountingHandler {
184        direct: AtomicUsize,
185        batched: AtomicUsize,
186        preflight: AtomicUsize,
187    }
188
189    #[async_trait::async_trait]
190    impl OpentsdbProtocolHandler for CountingHandler {
191        async fn preflight(&self, _: &[DataPoint], _: QueryContextRef) -> Result<()> {
192            self.preflight.fetch_add(1, Ordering::Relaxed);
193            Ok(())
194        }
195
196        async fn exec(&self, points: Vec<DataPoint>, _: QueryContextRef) -> Result<usize> {
197            self.direct.fetch_add(1, Ordering::Relaxed);
198            Ok(points.len())
199        }
200
201        async fn exec_batch(&self, points: Vec<DataPoint>, _: QueryContextRef) -> Result<usize> {
202            self.batched.fetch_add(1, Ordering::Relaxed);
203            Ok(points.len())
204        }
205    }
206
207    #[tokio::test]
208    async fn test_put_batches_only_non_debug_requests() {
209        for params in [
210            HashMap::new(),
211            HashMap::from([("summary".to_string(), String::new())]),
212            HashMap::from([("details".to_string(), String::new())]),
213            HashMap::from([
214                ("summary".to_string(), String::new()),
215                ("details".to_string(), String::new()),
216            ]),
217        ] {
218            let debug = !params.is_empty();
219            let handler = Arc::new(CountingHandler::default());
220            let body = Bytes::from_static(
221                br#"[
222                {"metric":"cpu","timestamp":1000,"value":1,"tags":{"host":"a"}},
223                {"metric":"cpu","timestamp":1001,"value":2,"tags":{"host":"b"}}
224            ]"#,
225            );
226            let (status, Json(response)) = put(
227                State(handler.clone() as OpentsdbProtocolHandlerRef),
228                Query(params),
229                Extension(QueryContext::with("greptime", "public")),
230                body,
231            )
232            .await
233            .unwrap();
234            assert_eq!(
235                usize::from(debug) * 2,
236                handler.direct.load(Ordering::Relaxed)
237            );
238            assert_eq!(usize::from(!debug), handler.batched.load(Ordering::Relaxed));
239            assert_eq!(
240                usize::from(debug),
241                handler.preflight.load(Ordering::Relaxed)
242            );
243            if debug {
244                assert_eq!(HttpStatusCode::OK, status);
245                let OpentsdbPutResponse::Debug(response) = response else {
246                    panic!("expected debug response")
247                };
248                assert_eq!(2, response.success);
249                assert_eq!(0, response.failed);
250            } else {
251                assert_eq!(HttpStatusCode::NO_CONTENT, status);
252                assert!(matches!(response, OpentsdbPutResponse::Empty));
253            }
254        }
255    }
256
257    #[test]
258    fn test_into_opentsdb_data_point() {
259        let request = DataPointRequest {
260            metric: "hello".to_string(),
261            timestamp: 1234,
262            value: 1.0,
263            tags: HashMap::from([("foo".to_string(), "a".to_string())]),
264        };
265        let data_point: DataPoint = request.into();
266        assert_eq!(data_point.metric(), "hello");
267        assert_eq!(data_point.ts_millis(), 1234000);
268        assert_eq!(data_point.value(), 1.0);
269        assert_eq!(
270            data_point.tags(),
271            &vec![("foo".to_string(), "a".to_string())]
272        );
273    }
274
275    #[tokio::test]
276    async fn test_parse_data_points() {
277        let raw_data_point1 = r#"{
278                "metric": "sys.cpu.nice",
279                "timestamp": 1346846400,
280                "value": 18,
281                "tags": {
282                    "host": "web01",
283                    "dc": "lga"
284                }
285            }"#;
286        let data_point1 = serde_json::from_str::<DataPointRequest>(raw_data_point1).unwrap();
287
288        let raw_data_point2 = r#"{
289                "metric": "sys.cpu.nice",
290                "timestamp": 1346846400,
291                "value": 9,
292                "tags": {
293                    "host": "web02",
294                    "dc": "lga"
295                }
296            }"#;
297        let data_point2 = serde_json::from_str::<DataPointRequest>(raw_data_point2).unwrap();
298
299        let body = Bytes::from(raw_data_point1);
300        let data_points = parse_data_points(body).await.unwrap();
301        assert_eq!(data_points.len(), 1);
302        assert_eq!(data_points[0], data_point1);
303
304        let body = Bytes::from(format!("[{raw_data_point1},{raw_data_point2}]"));
305        let data_points = parse_data_points(body).await.unwrap();
306        assert_eq!(data_points.len(), 2);
307        assert_eq!(data_points[0], data_point1);
308        assert_eq!(data_points[1], data_point2);
309
310        let body = Bytes::from("");
311        let result = parse_data_points(body).await;
312        assert!(result.is_err());
313        let err = result.unwrap_err().output_msg();
314        assert!(err.contains("EOF while parsing a value at line 1 column 0"));
315
316        let body = Bytes::from("hello world");
317        let result = parse_data_points(body).await;
318        assert!(result.is_err());
319        let err = result.unwrap_err().output_msg();
320        assert!(err.contains("expected value at line 1 column 1"));
321    }
322}
servers/http/opentsdb.rs

servers/http/
opentsdb.rs
