frontend/instance/

otlp.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod trace_ingest;
pub mod trace_semconv;
pub mod trace_types;

use std::collections::HashSet;
use std::sync::Arc;

use api::v1::RowInsertRequests;
use api::v1::value::ValueData;
use async_trait::async_trait;
use auth::{
    OTLP_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget,
    PermissionTableTargets,
};
use client::Output;
use common_catalog::consts::{
    METRICS_EXEMPLARS_TABLE, trace_operations_table_name, trace_services_table_name,
};
use common_error::ext::BoxedError;
use common_query::prelude::{EXEMPLAR_METRIC_COLUMN, GREPTIME_PHYSICAL_TABLE};
use common_telemetry::{tracing, warn};
use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
use operator::insert::{Inserter, admit_row_insert_batches, admit_write};
use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
use pipeline::{GreptimePipelineParams, PipelineWay};
use servers::error::{self, AuthSnafu, Result as ServerResult};
use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtocolInterceptorRef};
use servers::otlp;
use servers::otlp::trace::span::TraceSpanGroup;
use servers::query_handler::{
    MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome,
};
use servers::request_memory_limiter::ServerMemoryLimiter;
use session::context::QueryContextRef;
use session::protocol_ctx::OtlpMetricCtx;
use session::write_progress::WriteProgress;
use snafu::{IntoError, ResultExt};
use table::requests::{
    OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM, SEMANTIC_PER_TABLE_INDEX_KEY,
    SEMANTIC_PIPELINE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_LOG, SIGNAL_TYPE_METRIC,
    SOURCE_OPENTELEMETRY, TABLE_DATA_MODEL, TABLE_DATA_MODEL_LOG_V1,
};

pub(crate) use self::trace_ingest::TraceAuxCache;
use self::trace_ingest::trace_conventions;
use crate::instance::Instance;
use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS};

fn trace_permission_targets(
    table_name: &str,
    groups: &[TraceSpanGroup],
    ctx: &QueryContextRef,
) -> PermissionTableTargets {
    let catalog = ctx.current_catalog();
    let schema = ctx.current_schema();
    if catalog.is_empty() || schema.is_empty() || table_name.is_empty() {
        return PermissionTableTargets::Unresolved;
    }

    let has_spans = groups.iter().any(|group| !group.spans.is_empty());
    if !has_spans {
        return PermissionTableTargets::resolved(Vec::new());
    }

    let mut targets = vec![PermissionTableTarget::new(catalog, &schema, table_name)];
    if groups
        .iter()
        .flat_map(|group| &group.spans)
        .any(|span| span.service_name.is_some())
    {
        targets.extend([
            PermissionTableTarget::new(catalog, &schema, trace_services_table_name(table_name)),
            PermissionTableTarget::new(catalog, &schema, trace_operations_table_name(table_name)),
        ]);
    }

    PermissionTableTargets::resolved(targets)
}

#[async_trait]
impl OpenTelemetryProtocolHandler for Instance {
    #[tracing::instrument(skip_all)]
    async fn metrics(
        &self,
        request: ExportMetricsServiceRequest,
        ctx: QueryContextRef,
        memory_limiter: Option<(ServerMemoryLimiter, u64)>,
    ) -> ServerResult<MetricsIngestOutcome> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        let mut metric_ctx = ctx
            .protocol_ctx()
            .get_otlp_metric_ctx()
            .cloned()
            .unwrap_or_default();
        metric_ctx.resource_info = self.otlp_resource_info;

        let mut prepared = otlp::metrics::prepare_metrics(request, &mut metric_ctx)?;
        // Authorize source tables before preflight can expose schemas or discard points.
        let targets = {
            let catalog = ctx.current_catalog();
            let schema = ctx.current_schema();
            let mut table_names = HashSet::new();
            prepared
                .requests
                .iter()
                .filter(|request| table_names.insert(request.table_name.as_str()))
                .map(|request| PermissionTableTarget::new(catalog, &schema, &request.table_name))
                .collect::<Vec<_>>()
        };
        self.check_table_permission(
            &ctx,
            PermissionReq::Action(OTLP_WRITE),
            PermissionTableTargets::resolved(targets),
        )
        .context(AuthSnafu)?;

        let preflight = self
            .inserter
            .validate_and_normalize_otlp_metrics(
                &mut prepared.requests,
                &ctx,
                metric_ctx.with_metric_engine,
                ctx.extension(PHYSICAL_TABLE_PARAM)
                    .unwrap_or(GREPTIME_PHYSICAL_TABLE),
            )
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        let rejections = prepared.point_rejections(preflight.rejections)?;
        let otlp::metrics::MetricsConversion {
            requests,
            rows,
            mut exemplars,
            semantic_index,
            resource_info,
            mut outcome,
        } = prepared.finish(rejections, &preflight.physical_tables)?;
        if outcome.rejected_data_points > 0 {
            warn!(
                "Rejected {} OTLP metrics data points: {}",
                outcome.rejected_data_points,
                outcome.error_message.as_deref().unwrap_or_default()
            );
        }
        if outcome.accepted_data_points == 0 {
            return Ok(outcome);
        }

        // The engine option only controls new tables; existing ordinary sources cannot bind exemplars.
        let before = exemplars.len();
        exemplars.retain(|exemplar| {
            preflight
                .physical_tables
                .contains_key(exemplar.metric.as_ref())
        });
        servers::metrics::EXEMPLARS_DROPPED_TOTAL
            .with_label_values(&["otlp"])
            .inc_by((before - exemplars.len()) as u64);
        if !exemplars.is_empty() {
            self.check_table_permission(
                &ctx,
                PermissionReq::Action(OTLP_WRITE),
                PermissionTableTargets::resolved(vec![PermissionTableTarget::new(
                    ctx.current_catalog(),
                    ctx.current_schema(),
                    METRICS_EXEMPLARS_TABLE,
                )]),
            )
            .context(AuthSnafu)?;
        }

        let ctx = if !exemplars.is_empty()
            && let Some((limiter, held_bytes)) = memory_limiter
        {
            let bytes =
                servers::exemplar::estimated_row_bytes(std::iter::once(exemplars.as_slice()))
                    .and_then(|bytes| u64::try_from(bytes).ok())
                    .ok_or_else(|| {
                        error::MemoryLimitExceededSnafu.into_error(
                            common_memory_manager::Error::MemoryLimitExceeded {
                                requested_bytes: u64::MAX,
                                limit_bytes: limiter.limit_bytes(),
                            },
                        )
                    })?;
            limiter.check_request_fits(bytes, held_bytes)?;
            // Holding body permits while waiting for expansion can stall concurrent requests.
            let reservation = Arc::new(if held_bytes > 0 {
                limiter.try_acquire(bytes).ok_or_else(|| {
                    error::MemoryLimitExceededSnafu.into_error(
                        common_memory_manager::Error::MemoryLimitExceeded {
                            requested_bytes: bytes,
                            limit_bytes: limiter.limit_bytes(),
                        },
                    )
                })?
            } else {
                limiter.acquire(bytes).await?
            });
            let mut ctx = ctx.fork();
            // Queued and detached writes retain the reservation after caller cancellation.
            ctx.set_write_progress(Arc::new(WriteProgress::with_reservation(reservation)));
            Arc::new(ctx)
        } else {
            ctx
        };
        let exemplar_rows = exemplars.len();
        let (exemplars, _) = servers::exemplar::into_rows(exemplars)?;
        let total_rows = rows.checked_add(exemplar_rows).ok_or_else(|| {
            error::InvalidOtlpMetricInputSnafu {
                reason: "combined metric and exemplar row count overflow",
            }
            .build()
        })?;
        let ctx = admit_write(total_rows as u64, &ctx)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        OTLP_METRICS_ROWS.inc_by(rows as u64);

        let ctx = {
            let mut c = (*ctx).clone();
            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
            c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
            // Per-table metric specifics + resource/scope lineage ride this
            // internal channel; the auto-create path folds them per schema and
            // table name.
            if let Some(index) = semantic_index.encode(&c.current_schema()) {
                c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index);
            }
            c.set_extension(OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM.to_string());
            Arc::new(c)
        };

        if exemplar_rows > 0 {
            let confirmed = Arc::new(ctx.with_confirmed_writes());
            let source_names = exemplars
                .inserts
                .iter()
                .filter_map(|request| request.rows.as_ref())
                .flat_map(|rows| {
                    let index = rows
                        .schema
                        .iter()
                        .position(|c| c.column_name == EXEMPLAR_METRIC_COLUMN);
                    rows.rows.iter().filter_map(move |row| {
                        match &row.values.get(index?)?.value_data {
                            Some(ValueData::StringValue(name)) => Some(name.as_str()),
                            _ => None,
                        }
                    })
                })
                .collect::<HashSet<_>>();
            let (source_requests, other): (Vec<_>, Vec<_>) = requests
                .inserts
                .into_iter()
                .partition(|request| source_names.contains(request.table_name.as_str()));
            let prepared = self
                .inserter
                .prepare_metric_row_inserts(
                    RowInsertRequests {
                        inserts: source_requests,
                    },
                    confirmed.clone(),
                    &self.statement_executor,
                    ctx.extension(PHYSICAL_TABLE_PARAM)
                        .unwrap_or(GREPTIME_PHYSICAL_TABLE)
                        .to_string(),
                )
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu);
            let source_ready = prepared.is_ok();
            let sources = prepared
                .as_ref()
                .map(|request| request.source_tables().cloned().collect::<Vec<_>>())
                .unwrap_or_default();
            let (metric_output, other_output, exemplar_output) = tokio::join!(
                async {
                    self.inserter
                        .handle_prepared_metric_inserts(prepared?, confirmed.clone())
                        .await
                        .map_err(BoxedError::new)
                        .context(error::ExecuteGrpcQuerySnafu)
                },
                async {
                    if other.is_empty() {
                        Ok(Output::new_with_affected_rows(0))
                    } else {
                        self.handle_otlp_metric_row_inserts(
                            RowInsertRequests { inserts: other },
                            confirmed.clone(),
                            &metric_ctx,
                        )
                        .await
                    }
                },
                async {
                    // Failed source preparation must not bind exemplars to a later replacement table.
                    if !source_ready {
                        return Ok(Output::new_with_affected_rows(0));
                    }
                    self.inserter
                        .handle_exemplar_inserts(
                            exemplars,
                            confirmed.clone(),
                            &self.statement_executor,
                            &sources,
                        )
                        .await
                        .map_err(BoxedError::new)
                        .context(error::ExecuteGrpcQuerySnafu)
                },
            );
            let mut progress = common_query::WriteProgress::default();
            for output in [metric_output, other_output, exemplar_output] {
                match output {
                    Ok(output) => {
                        outcome.write_cost += output.meta.cost;
                        progress.merge(output.write_progress().await);
                    }
                    Err(error) => {
                        progress.merge(common_query::WriteProgress::failed(BoxedError::new(error)))
                    }
                }
            }
            progress
                .into_result()
                .context(error::ExecuteGrpcQuerySnafu)?;
        } else {
            let output = self
                .handle_otlp_metric_row_inserts(requests, ctx.clone(), &metric_ctx)
                .await?;
            outcome.write_cost = output.meta.cost;
        }

        // Derived enrichment follows the accepted metric submission, which may
        // still be queued in asynchronous mode. Failures remain warning-only
        // to avoid retrying metric data the server already accepted.
        if let Some(resource_info) = resource_info {
            let written = match self.check_row_insert_permission(
                &resource_info,
                &ctx,
                PermissionReq::Action(OTLP_WRITE),
            ) {
                Ok(_) => self
                    .handle_row_inserts(resource_info, ctx, false, false)
                    .await
                    .map_err(BoxedError::new)
                    .map_err(|e| e.to_string()),
                Err(e) => Err(e.to_string()),
            };
            match written {
                Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost,
                Err(e) => {
                    OTLP_RESOURCE_INFO_WRITE_ERRORS.inc();
                    warn!("Failed to write the OTLP resource descriptor table: {e}");
                    outcome.error_message.get_or_insert(format!(
                        "metric data was accepted, but writing the resource \
                         descriptor table `{}` failed: {e}",
                        otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME
                    ));
                }
            }
        }

        Ok(outcome)
    }

    #[tracing::instrument(skip_all)]
    async fn traces(
        &self,
        _pipeline_handler: PipelineHandlerRef,
        request: ExportTraceServiceRequest,
        pipeline: PipelineWay,
        _pipeline_params: GreptimePipelineParams,
        table_name: String,
        ctx: QueryContextRef,
    ) -> ServerResult<TraceIngestOutcome> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        // `schema_url` is consumed by `parse`, so derive conventions first.
        let conventions = trace_conventions(&request);
        let spans = otlp::trace::span::parse(request);
        let targets = trace_permission_targets(&table_name, &spans, &ctx);
        self.check_table_permission(&ctx, PermissionReq::Action(OTLP_WRITE), targets)
            .context(AuthSnafu)?;
        // Count the external spans once, before chunking, retries, or derived tables.
        let rows = spans.iter().map(|group| group.spans.len() as u64).sum();
        let ctx = admit_write(rows, &ctx)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;
        self.ingest_trace_spans(&pipeline, table_name, spans, &conventions, ctx)
            .await
    }

    #[tracing::instrument(skip_all)]
    async fn logs(
        &self,
        pipeline_handler: PipelineHandlerRef,
        request: ExportLogsServiceRequest,
        pipeline: PipelineWay,
        pipeline_params: GreptimePipelineParams,
        table_name: String,
        ctx: QueryContextRef,
    ) -> ServerResult<Vec<Output>> {
        self.plugins
            .get::<PermissionCheckerRef>()
            .as_ref()
            .check_permission(ctx.current_user(), PermissionReq::Action(OTLP_WRITE))
            .context(AuthSnafu)?;

        let interceptor_ref = self
            .plugins
            .get::<OpenTelemetryProtocolInterceptorRef<servers::error::Error>>();
        interceptor_ref.pre_execute(ctx.clone())?;
        let ctx = Arc::new(ctx.fork());

        // `as_req_iter` clones this ctx into each `temp_ctx`, so identity set here
        // reaches the context that drives table auto-create.
        let ctx = {
            let mut c = (*ctx).clone();
            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_LOG);
            c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY);
            // Only applied when the table is auto-created. The v1 rows are built only
            // when the table didn't exist at read time; pipeline-defined tables have no
            // built-in schema.
            if matches!(pipeline, PipelineWay::OtlpLogDirect(_)) {
                c.set_extension(TABLE_DATA_MODEL, TABLE_DATA_MODEL_LOG_V1);
                c.set_extension(SEMANTIC_PIPELINE, TABLE_DATA_MODEL_LOG_V1);
                // TODO: declare the `service` entity on `service_name` once the cost of
                // scanning log tables during entity-graph derivation is measured.
            }
            Arc::new(c)
        };

        let opt_req = otlp::logs::to_grpc_insert_requests(
            request,
            pipeline,
            pipeline_params,
            table_name,
            &ctx,
            pipeline_handler,
        )
        .await?;

        let mut batches = opt_req.as_req_iter(ctx).collect::<Vec<_>>();
        for (temp_ctx, requests) in &batches {
            self.check_row_insert_permission(requests, temp_ctx, PermissionReq::Action(OTLP_WRITE))
                .context(AuthSnafu)?;
        }
        admit_row_insert_batches(&mut batches)
            .await
            .map_err(BoxedError::new)
            .context(error::ExecuteGrpcQuerySnafu)?;

        let mut outputs = Vec::with_capacity(batches.len());
        for (temp_ctx, requests) in batches {
            let cnt = requests
                .inserts
                .iter()
                .filter_map(|r| r.rows.as_ref().map(|r| r.rows.len()))
                .sum::<usize>();

            let o = self
                .handle_log_inserts(requests, temp_ctx)
                .await
                .inspect(|_| OTLP_LOGS_ROWS.inc_by(cnt as u64))
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?;
            outputs.push(o);
        }

        Ok(outputs)
    }
}

impl Instance {
    /// Inserts converted OTLP metrics through the eligible batching or ordinary write path.
    /// The caller must check table permissions and admit the complete request first.
    pub async fn handle_otlp_metric_row_inserts(
        &self,
        requests: RowInsertRequests,
        ctx: QueryContextRef,
        metric_ctx: &OtlpMetricCtx,
    ) -> ServerResult<Output> {
        let physical_table = ctx
            .extension(PHYSICAL_TABLE_PARAM)
            .unwrap_or(GREPTIME_PHYSICAL_TABLE)
            .to_string();
        // The bulk path converts each request's time index unit to the
        // destination table's unit during batch alignment, so no pre-gate
        // alignment is needed here.
        let batcher = self
            .logical_batcher()
            .filter(|_| ctx.logical_batching_enabled() && metric_ctx.with_metric_engine);
        let batcher = if batcher.is_some()
            && self
                .inserter
                .can_batch_metric_rows(&requests, &ctx, &physical_table)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)?
        {
            batcher
        } else {
            None
        };

        // OTLP tables have one sample field in both the ordinary and physical paths.
        if let Some(batcher) = batcher {
            let (mut output, cost) = batcher
                .submit_output_with(requests, ctx.clone(), |mut requests| {
                    let ctx = ctx.clone();
                    async move {
                        Inserter::meter_row_inserts(&mut requests, &ctx)
                            .await
                            .map_err(BoxedError::new)
                            .context(error::ExecuteGrpcQuerySnafu)
                    }
                })
                .await?;
            output.meta.cost = cost as _;
            Ok(output)
        } else if !metric_ctx.with_metric_engine {
            self.handle_row_inserts(requests, ctx.clone(), false, true)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)
        } else {
            self.handle_metric_row_inserts(requests, ctx.clone(), physical_table)
                .await
                .map_err(BoxedError::new)
                .context(error::ExecuteGrpcQuerySnafu)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use auth::{PermissionTableTarget, PermissionTableTargets};
    use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
    use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value};
    use opentelemetry_proto::tonic::resource::v1::Resource;
    use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans, Span};
    use session::context::QueryContext;

    use super::trace_permission_targets;

    #[test]
    fn test_trace_permission_targets() {
        let request = ExportTraceServiceRequest {
            resource_spans: vec![ResourceSpans {
                resource: Some(Resource {
                    attributes: vec![KeyValue {
                        key: "service.name".to_string(),
                        value: Some(AnyValue {
                            value: Some(any_value::Value::StringValue("frontend".to_string())),
                        }),
                        ..Default::default()
                    }],
                    ..Default::default()
                }),
                scope_spans: vec![ScopeSpans {
                    spans: vec![Span::default()],
                    ..Default::default()
                }],
                ..Default::default()
            }],
        };
        let groups = servers::otlp::trace::span::parse(request);
        let ctx = Arc::new(QueryContext::with("greptime", "public"));

        assert_eq!(
            PermissionTableTargets::Resolved(vec![
                PermissionTableTarget::new("greptime", "public", "traces"),
                PermissionTableTarget::new("greptime", "public", "traces_services"),
                PermissionTableTarget::new("greptime", "public", "traces_operations"),
            ]),
            trace_permission_targets("traces", &groups, &ctx)
        );
        assert_eq!(
            PermissionTableTargets::Unresolved,
            trace_permission_targets("", &groups, &ctx)
        );
    }
}