servers/http/
otlp.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::sync::Arc;
16
17use axum::Extension;
18use axum::extract::State;
19use axum::http::{StatusCode, header};
20use axum::response::IntoResponse;
21use axum_extra::TypedHeader;
22use bytes::Bytes;
23use common_catalog::consts::{TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY};
24use common_telemetry::tracing;
25use headers::ContentType;
26use mime_guess::mime;
27use opentelemetry_proto::tonic::collector::logs::v1::{
28    ExportLogsServiceRequest, ExportLogsServiceResponse,
29};
30use opentelemetry_proto::tonic::collector::metrics::v1::{
31    ExportMetricsPartialSuccess, ExportMetricsServiceResponse,
32};
33use opentelemetry_proto::tonic::collector::trace::v1::{
34    ExportTracePartialSuccess, ExportTraceServiceRequest, ExportTraceServiceResponse,
35};
36use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
37use pipeline::PipelineWay;
38use prost::Message;
39use session::context::{Channel, QueryContext};
40use session::protocol_ctx::{MetricType, OtlpMetricCtx, ProtocolCtx};
41use snafu::prelude::*;
42
43use crate::error::{self, PipelineSnafu, Result};
44use crate::http::extractor::{
45    LogTableName, OtlpMetricOptions, PipelineInfo, SelectInfoWrapper, TraceTableName,
46};
47use crate::http::header::{CONTENT_TYPE_PROTOBUF, write_cost_header_map};
48use crate::http::memory_limit::BodyMemoryAccounting;
49use crate::metrics::METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED;
50use crate::query_handler::{
51    MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome,
52};
53use crate::request_memory_limiter::ServerMemoryLimiter;
54
55#[derive(Clone, prost::Message)]
56pub struct GoogleRpcStatus {
57    #[prost(int32, tag = "1")]
58    pub code: i32,
59    #[prost(string, tag = "2")]
60    pub message: String,
61}
62
63fn is_json_content_type(content_type: Option<&ContentType>) -> bool {
64    match content_type {
65        None => false,
66        Some(ct) => {
67            let mime: mime::Mime = ct.clone().into();
68            mime.subtype() == mime::JSON
69        }
70    }
71}
72
73fn content_type_to_string(content_type: Option<&TypedHeader<ContentType>>) -> String {
74    content_type
75        .map(|h| h.0.to_string())
76        .unwrap_or_else(|| "not specified".to_string())
77}
78
79#[derive(Clone)]
80pub struct OtlpState {
81    pub with_metric_engine: bool,
82    pub handler: OpenTelemetryProtocolHandlerRef,
83}
84
85#[axum_macros::debug_handler]
86#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "metrics"))]
87pub async fn metrics(
88    State(state): State<OtlpState>,
89    Extension(mut query_ctx): Extension<QueryContext>,
90    http_opts: OtlpMetricOptions,
91    content_type: Option<TypedHeader<ContentType>>,
92    extensions: http::Extensions,
93    bytes: Bytes,
94) -> Result<OtlpMetricsResponse> {
95    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
96        return error::UnsupportedJsonContentTypeSnafu {}.fail();
97    }
98
99    let db = query_ctx.get_db_string();
100    query_ctx.set_channel(Channel::Otlp);
101
102    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_METRICS_ELAPSED
103        .with_label_values(&[db.as_str()])
104        .start_timer();
105    let request = ExportMetricsServiceRequest::decode(bytes).with_context(|_| {
106        error::DecodeOtlpRequestSnafu {
107            content_type: content_type_to_string(content_type.as_ref()),
108        }
109    })?;
110
111    let OtlpState {
112        with_metric_engine,
113        handler,
114    } = state;
115
116    query_ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx {
117        promote_all_resource_attrs: http_opts.promote_all_resource_attrs,
118        resource_attrs: http_opts.resource_attrs,
119        promote_scope_attrs: http_opts.promote_scope_attrs,
120        with_metric_engine,
121        // set by the frontend from its config
122        resource_info: false,
123        metric_type: MetricType::Init,
124        metric_translation_strategy: http_opts.metric_translation_strategy,
125    }));
126    let query_ctx = Arc::new(query_ctx);
127
128    let memory_limiter = extensions
129        .get::<ServerMemoryLimiter>()
130        .cloned()
131        .map(|limiter| {
132            let held_bytes = extensions
133                .get::<BodyMemoryAccounting>()
134                .map_or(0, BodyMemoryAccounting::reserved_bytes);
135            (limiter, held_bytes)
136        });
137    match handler.metrics(request, query_ctx, memory_limiter).await {
138        Ok(outcome) => {
139            if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 {
140                Ok(OtlpMetricsResponse::Failure(outcome))
141            } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() {
142                Ok(OtlpMetricsResponse::PartialSuccess(outcome))
143            } else {
144                Ok(OtlpMetricsResponse::FullSuccess(outcome))
145            }
146        }
147        Err(error::Error::InvalidOtlpMetricInput { reason }) => {
148            Ok(OtlpMetricsResponse::Failure(MetricsIngestOutcome {
149                error_message: Some(reason),
150                ..Default::default()
151            }))
152        }
153        Err(error) => Err(error),
154    }
155}
156
157#[axum_macros::debug_handler]
158#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "traces"))]
159pub async fn traces(
160    State(state): State<OtlpState>,
161    TraceTableName(table_name): TraceTableName,
162    pipeline_info: PipelineInfo,
163    Extension(mut query_ctx): Extension<QueryContext>,
164    content_type: Option<TypedHeader<ContentType>>,
165    bytes: Bytes,
166) -> Result<OtlpTraceResponse> {
167    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
168        return error::UnsupportedJsonContentTypeSnafu {}.fail();
169    }
170
171    let db = query_ctx.get_db_string();
172    let table_name = table_name.unwrap_or_else(|| TRACE_TABLE_NAME.to_string());
173
174    query_ctx.set_channel(Channel::Otlp);
175    query_ctx.set_extension(TRACE_TABLE_NAME_SESSION_KEY, &table_name);
176
177    let query_ctx = Arc::new(query_ctx);
178    let _timer = crate::metrics::METRIC_HTTP_OPENTELEMETRY_TRACES_ELAPSED
179        .with_label_values(&[db.as_str()])
180        .start_timer();
181    let request = ExportTraceServiceRequest::decode(bytes).with_context(|_| {
182        error::DecodeOtlpRequestSnafu {
183            content_type: content_type_to_string(content_type.as_ref()),
184        }
185    })?;
186
187    let pipeline = PipelineWay::from_name_and_default(
188        pipeline_info.pipeline_name.as_deref(),
189        pipeline_info.pipeline_version.as_deref(),
190        None,
191    )
192    .context(PipelineSnafu)?;
193
194    let pipeline_params = pipeline_info.pipeline_params;
195
196    let OtlpState { handler, .. } = state;
197
198    // here we use nightly feature `trait_upcasting` to convert handler to
199    // pipeline_handler
200    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
201
202    handler
203        .traces(
204            pipeline_handler,
205            request,
206            pipeline,
207            pipeline_params,
208            table_name,
209            query_ctx,
210        )
211        .await
212        .map(|outcome| {
213            if outcome.accepted_spans == 0 && outcome.rejected_spans > 0 {
214                OtlpTraceResponse::Failure(outcome)
215            } else if outcome.rejected_spans > 0 || outcome.error_message.is_some() {
216                OtlpTraceResponse::PartialSuccess(outcome)
217            } else {
218                OtlpTraceResponse::FullSuccess(outcome)
219            }
220        })
221}
222
223#[axum_macros::debug_handler]
224#[tracing::instrument(skip_all, fields(protocol = "otlp", request_type = "logs"))]
225pub async fn logs(
226    State(state): State<OtlpState>,
227    Extension(mut query_ctx): Extension<QueryContext>,
228    pipeline_info: PipelineInfo,
229    LogTableName(tablename): LogTableName,
230    SelectInfoWrapper(select_info): SelectInfoWrapper,
231    content_type: Option<TypedHeader<ContentType>>,
232    bytes: Bytes,
233) -> Result<OtlpResponse<ExportLogsServiceResponse>> {
234    if is_json_content_type(content_type.as_ref().map(|h| &h.0)) {
235        return error::UnsupportedJsonContentTypeSnafu {}.fail();
236    }
237
238    let tablename = tablename.unwrap_or_else(|| "opentelemetry_logs".to_string());
239    let db = query_ctx.get_db_string();
240    query_ctx.set_channel(Channel::Otlp);
241    let query_ctx = Arc::new(query_ctx);
242    let _timer = METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED
243        .with_label_values(&[db.as_str()])
244        .start_timer();
245    let request = ExportLogsServiceRequest::decode(bytes).with_context(|_| {
246        error::DecodeOtlpRequestSnafu {
247            content_type: content_type_to_string(content_type.as_ref()),
248        }
249    })?;
250
251    let pipeline = PipelineWay::from_name_and_default(
252        pipeline_info.pipeline_name.as_deref(),
253        pipeline_info.pipeline_version.as_deref(),
254        Some(PipelineWay::OtlpLogDirect(Box::new(select_info))),
255    )
256    .context(PipelineSnafu)?;
257    let pipeline_params = pipeline_info.pipeline_params;
258
259    let OtlpState { handler, .. } = state;
260
261    // here we use nightly feature `trait_upcasting` to convert handler to
262    // pipeline_handler
263    let pipeline_handler: Arc<dyn PipelineHandler + Send + Sync> = handler.clone();
264    handler
265        .logs(
266            pipeline_handler,
267            request,
268            pipeline,
269            pipeline_params,
270            tablename,
271            query_ctx,
272        )
273        .await
274        .map(|o| OtlpResponse {
275            resp_body: ExportLogsServiceResponse {
276                partial_success: None,
277            },
278            write_cost: o.iter().map(|o| o.meta.cost).sum(),
279        })
280}
281
282pub struct OtlpResponse<T: Message> {
283    resp_body: T,
284    write_cost: usize,
285}
286
287pub enum OtlpMetricsResponse {
288    FullSuccess(MetricsIngestOutcome),
289    PartialSuccess(MetricsIngestOutcome),
290    Failure(MetricsIngestOutcome),
291}
292
293impl IntoResponse for OtlpMetricsResponse {
294    fn into_response(self) -> axum::response::Response {
295        match self {
296            OtlpMetricsResponse::FullSuccess(outcome) => {
297                let mut header_map = write_cost_header_map(outcome.write_cost);
298                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
299                let body = ExportMetricsServiceResponse {
300                    partial_success: None,
301                };
302                (header_map, body.encode_to_vec()).into_response()
303            }
304            OtlpMetricsResponse::PartialSuccess(outcome) => {
305                let mut header_map = write_cost_header_map(outcome.write_cost);
306                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
307                let body = ExportMetricsServiceResponse {
308                    partial_success: Some(ExportMetricsPartialSuccess {
309                        rejected_data_points: outcome.rejected_data_points,
310                        error_message: outcome.error_message.unwrap_or_default(),
311                    }),
312                };
313                (header_map, body.encode_to_vec()).into_response()
314            }
315            OtlpMetricsResponse::Failure(outcome) => {
316                let status = GoogleRpcStatus {
317                    code: tonic::Code::InvalidArgument as i32,
318                    message: outcome.error_message.unwrap_or_default(),
319                };
320                (
321                    StatusCode::BAD_REQUEST,
322                    [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())],
323                    status.encode_to_vec(),
324                )
325                    .into_response()
326            }
327        }
328    }
329}
330
331impl<T: Message> IntoResponse for OtlpResponse<T> {
332    fn into_response(self) -> axum::response::Response {
333        let mut header_map = write_cost_header_map(self.write_cost);
334        header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
335
336        (header_map, self.resp_body.encode_to_vec()).into_response()
337    }
338}
339
340pub enum OtlpTraceResponse {
341    FullSuccess(TraceIngestOutcome),
342    PartialSuccess(TraceIngestOutcome),
343    Failure(TraceIngestOutcome),
344}
345
346impl IntoResponse for OtlpTraceResponse {
347    fn into_response(self) -> axum::response::Response {
348        match self {
349            OtlpTraceResponse::FullSuccess(outcome) => {
350                let mut header_map = write_cost_header_map(outcome.write_cost);
351                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
352                let body = ExportTraceServiceResponse {
353                    partial_success: None,
354                };
355                (header_map, body.encode_to_vec()).into_response()
356            }
357            OtlpTraceResponse::PartialSuccess(outcome) => {
358                let mut header_map = write_cost_header_map(outcome.write_cost);
359                header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone());
360                let body = ExportTraceServiceResponse {
361                    partial_success: outcome.error_message.map(|error_message| {
362                        ExportTracePartialSuccess {
363                            rejected_spans: outcome.rejected_spans as i64,
364                            error_message,
365                        }
366                    }),
367                };
368                (header_map, body.encode_to_vec()).into_response()
369            }
370            OtlpTraceResponse::Failure(outcome) => {
371                let status = GoogleRpcStatus {
372                    code: 0,
373                    message: outcome.error_message.unwrap_or_default(),
374                };
375                (
376                    StatusCode::BAD_REQUEST,
377                    [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())],
378                    status.encode_to_vec(),
379                )
380                    .into_response()
381            }
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests;
servers/http/otlp.rs

servers/http/
otlp.rs
