servers/
otel_arrow.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 auth::UserProviderRef;
18use common_error::ext::ErrorExt;
19use common_error::status_code::status_to_tonic_code;
20use common_telemetry::error;
21use futures::SinkExt;
22use otel_arrow_rust::Consumer;
23use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsService;
24use otel_arrow_rust::proto::opentelemetry::arrow::v1::{
25    BatchArrowRecords, BatchStatus, StatusCode as ArrowStatusCode,
26};
27use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
28use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric;
29use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx};
30use tonic::metadata::{Entry, MetadataValue};
31use tonic::service::Interceptor;
32use tonic::{Request, Response, Status, Streaming};
33
34use crate::error;
35use crate::grpc::context_auth;
36use crate::grpc::memory_limit::PreDecodeMemoryReservation;
37use crate::query_handler::{MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef};
38use crate::request_memory_limiter::ServerMemoryLimiter;
39
40const EXPONENTIAL_HISTOGRAM_UNSUPPORTED: &str = "OTel Arrow exponential histograms are unsupported because the Arrow wire format omits zero_threshold";
41
42pub struct OtelArrowServiceHandler<T> {
43    handler: T,
44    user_provider: Option<UserProviderRef>,
45}
46
47impl<T> OtelArrowServiceHandler<T> {
48    pub fn new(handler: T, user_provider: Option<UserProviderRef>) -> Self {
49        Self {
50            handler,
51            user_provider,
52        }
53    }
54}
55
56/// Removes unsupported Arrow histograms before ingestion, preserving other metrics.
57fn remove_exponential_histograms(request: &mut ExportMetricsServiceRequest) -> bool {
58    let mut has_data_points = false;
59    for scope in request
60        .resource_metrics
61        .iter_mut()
62        .flat_map(|resource| &mut resource.scope_metrics)
63    {
64        scope.metrics.retain(|item| {
65            if let Some(metric::Data::ExponentialHistogram(histogram)) = &item.data {
66                has_data_points |= !histogram.data_points.is_empty();
67                false
68            } else {
69                true
70            }
71        });
72    }
73    has_data_points
74}
75
76fn batch_status(
77    batch_id: i64,
78    outcome: MetricsIngestOutcome,
79    has_exponential_histogram_data_points: bool,
80) -> BatchStatus {
81    let status_code = if outcome.accepted_data_points == 0
82        && (outcome.rejected_data_points > 0 || has_exponential_histogram_data_points)
83    {
84        ArrowStatusCode::InvalidArgument
85    } else {
86        ArrowStatusCode::Ok
87    };
88    let status_message = if has_exponential_histogram_data_points {
89        EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string()
90    } else {
91        outcome.error_message.unwrap_or_default()
92    };
93    BatchStatus {
94        batch_id,
95        status_code: status_code as i32,
96        status_message,
97    }
98}
99
100#[async_trait::async_trait]
101impl ArrowMetricsService for OtelArrowServiceHandler<OpenTelemetryProtocolHandlerRef> {
102    type ArrowMetricsStream = futures::channel::mpsc::Receiver<Result<BatchStatus, Status>>;
103    async fn arrow_metrics(
104        &self,
105        request: Request<Streaming<BatchArrowRecords>>,
106    ) -> Result<Response<Self::ArrowMetricsStream>, Status> {
107        let (mut sender, receiver) = futures::channel::mpsc::channel(100);
108
109        let (headers, extensions, mut incoming_requests) = request.into_parts();
110
111        // OTEL Arrow is currently used only for external ingestion. Its service bypasses
112        // the frontend router middleware, so even the internal listener defaults to
113        // Channel::Grpc and remains rate-limited. Before using this path internally,
114        // propagate the server-owned Channel::Internal marker to this service.
115        let query_ctx =
116            context_auth::create_query_context_from_grpc_metadata(&headers, &extensions)?;
117        context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?;
118        let query_ctx = {
119            let mut ctx = query_ctx.fork();
120            ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx::default()));
121            Arc::new(ctx)
122        };
123
124        let handler = self.handler.clone();
125        let memory_limiter = extensions.get::<ServerMemoryLimiter>().cloned();
126        let decode_reservation = extensions.get::<PreDecodeMemoryReservation>().cloned();
127
128        // handles incoming requests
129        common_runtime::spawn_global(async move {
130            let mut consumer = Consumer::default();
131            while let Some(batch_res) = incoming_requests.message().await.transpose() {
132                let mut batch = match batch_res {
133                    Ok(batch) => batch,
134                    Err(e) => {
135                        error!(
136                            "Failed to receive batch from otel-arrow client, error: {}",
137                            e
138                        );
139                        let _ = sender.send(Err(e)).await;
140                        return;
141                    }
142                };
143                let batch_id = batch.batch_id;
144                let mut request = match consumer.consume_metrics_batches(&mut batch).map_err(|e| {
145                    error::HandleOtelArrowRequestSnafu {
146                        err_msg: e.to_string(),
147                    }
148                    .build()
149                }) {
150                    Ok(request) => request,
151                    Err(e) => {
152                        let _ = sender
153                            .send(Err(Status::new(
154                                status_to_tonic_code(e.status_code()),
155                                e.to_string(),
156                            )))
157                            .await;
158                        error!(e;
159                            "Failed to consume batch from otel-arrow client"
160                        );
161                        return;
162                    }
163                };
164                let has_exponential_histogram_data_points =
165                    remove_exponential_histograms(&mut request);
166                let memory_limiter = memory_limiter.clone().map(|limiter| {
167                    let held_bytes = decode_reservation
168                        .as_ref()
169                        .map_or(0, PreDecodeMemoryReservation::reserved_bytes);
170                    (limiter, held_bytes)
171                });
172                let outcome = match handler
173                    .metrics(request, query_ctx.clone(), memory_limiter)
174                    .await
175                {
176                    Ok(outcome) => outcome,
177                    Err(error::Error::InvalidOtlpMetricInput { reason }) => {
178                        let _ = sender
179                            .send(Ok(BatchStatus {
180                                batch_id,
181                                status_code: ArrowStatusCode::InvalidArgument as i32,
182                                status_message: reason,
183                            }))
184                            .await;
185                        continue;
186                    }
187                    Err(e) => {
188                        let _ = sender
189                            .send(Err(Status::new(
190                                status_to_tonic_code(e.status_code()),
191                                e.to_string(),
192                            )))
193                            .await;
194                        error!(e; "Failed to ingest metrics from otel-arrow");
195                        return;
196                    }
197                };
198                let batch_status =
199                    batch_status(batch_id, outcome, has_exponential_histogram_data_points);
200                let _ = sender.send(Ok(batch_status)).await;
201            }
202        });
203        Ok(Response::new(receiver))
204    }
205}
206
207/// This serves as a workaround for otel-arrow collector's custom header.
208#[derive(Clone)]
209pub struct HeaderInterceptor;
210
211impl Interceptor for HeaderInterceptor {
212    fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
213        if let Ok(Entry::Occupied(mut e)) = request.metadata_mut().entry("grpc-encoding") {
214            // This works as a workaround to handle customized compression type (zstdarrow*) in otel-arrow.
215            if e.get().as_bytes().starts_with(b"zstdarrow") {
216                e.insert(MetadataValue::from_static("zstd"));
217            }
218        }
219        Ok(request)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn removes_arrow_exponential_histograms_and_preserves_other_metrics() {
229        use otel_arrow_rust::proto::opentelemetry::metrics::v1::{
230            ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, ResourceMetrics,
231            ScopeMetrics,
232        };
233
234        let mut request = ExportMetricsServiceRequest {
235            resource_metrics: vec![ResourceMetrics {
236                scope_metrics: vec![ScopeMetrics {
237                    metrics: vec![
238                        Metric {
239                            data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram {
240                                data_points: vec![ExponentialHistogramDataPoint::default()],
241                                ..Default::default()
242                            })),
243                            ..Default::default()
244                        },
245                        Metric {
246                            data: Some(metric::Data::Gauge(Gauge::default())),
247                            ..Default::default()
248                        },
249                    ],
250                    ..Default::default()
251                }],
252                ..Default::default()
253            }],
254        };
255        assert!(remove_exponential_histograms(&mut request));
256        let metrics = &request.resource_metrics[0].scope_metrics[0].metrics;
257        assert_eq!(metrics.len(), 1);
258        assert!(matches!(metrics[0].data, Some(metric::Data::Gauge(_))));
259        assert!(!remove_exponential_histograms(&mut request));
260    }
261
262    #[test]
263    fn batch_status_explains_arrow_exponential_histogram_limit() {
264        for accepted_data_points in [0, 1] {
265            let status = batch_status(
266                7,
267                MetricsIngestOutcome {
268                    accepted_data_points,
269                    ..Default::default()
270                },
271                true,
272            );
273            assert_eq!(7, status.batch_id);
274            assert_eq!(
275                if accepted_data_points == 0 {
276                    ArrowStatusCode::InvalidArgument
277                } else {
278                    ArrowStatusCode::Ok
279                } as i32,
280                status.status_code
281            );
282            assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message);
283        }
284    }
285}
servers/otel_arrow.rs

servers/
otel_arrow.rs
