servers/http/
dyn_trace.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 axum::http::StatusCode;
16use axum::response::IntoResponse;
17use common_telemetry::{TRACE_RELOAD_HANDLE, get_or_init_tracer, info};
18
19use crate::error::{InvalidParameterSnafu, Result};
20
21#[axum_macros::debug_handler]
22pub async fn dyn_trace_handler(enable_str: String) -> Result<impl IntoResponse> {
23    let enable = enable_str.parse::<bool>().map_err(|e| {
24        InvalidParameterSnafu {
25            reason: format!("Invalid parameter \"enable\": {e:?}"),
26        }
27        .build()
28    })?;
29
30    let Some(trace_reload_handle) = TRACE_RELOAD_HANDLE.get() else {
31        return Ok((
32            StatusCode::SERVICE_UNAVAILABLE,
33            "trace reload handle is not initialized".to_string(),
34        ));
35    };
36
37    if enable {
38        let tracer = match get_or_init_tracer() {
39            Ok(tracer) => tracer,
40            Err(reason) => {
41                return Ok((StatusCode::SERVICE_UNAVAILABLE, reason.to_string()));
42            }
43        };
44
45        let trace_layer = tracing_opentelemetry::layer().with_tracer(tracer);
46        trace_reload_handle.reload(Some(trace_layer));
47        info!("trace enabled");
48        Ok((StatusCode::OK, "trace enabled".to_string()))
49    } else {
50        trace_reload_handle.reload(None);
51        info!("trace disabled");
52        Ok((StatusCode::OK, "trace disabled".to_string()))
53    }
54}