rocketmq_observability/runtime_diagnostics/
service.rs

1// Copyright 2026 The RocketMQ Rust Authors
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;
16use std::time::Duration;
17
18use rocketmq_runtime::{
19    ChildServiceContext, RuntimeComponent, RuntimeDiagnosticsInputs, ShutdownDeadline, ShutdownReport,
20};
21use tokio::sync::OnceCell;
22
23use super::{read_token, serve, RuntimeDiagnosticsEndpointConfig, RuntimeDiagnosticsEndpointHandle};
24use crate::metrics::runtime::RuntimeMetricsRecorder;
25use crate::{ObservabilityError, TelemetryHandle};
26
27pub const RUNTIME_DIAGNOSTICS_MODE_ENV: &str = "ROCKETMQ_RUNTIME_DIAGNOSTICS_MODE";
28
29/// Explicitly selects diagnostic work, independently of compiled features.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub enum RuntimeDiagnosticsMode {
32    Disabled,
33    MetricsOnly { sample_interval: Duration },
34    EndpointOnly(RuntimeDiagnosticsEndpointConfig),
35    EndpointAndMetrics(RuntimeDiagnosticsEndpointConfig),
36}
37
38impl RuntimeDiagnosticsMode {
39    /// Loads explicit mode intent, retaining the legacy endpoint opt-in default.
40    ///
41    /// With no mode configured, an endpoint configuration selects both outputs;
42    /// no endpoint configuration selects `Disabled`. A no-op telemetry handle
43    /// never creates a sampler, including in either metrics-requesting mode.
44    ///
45    /// # Errors
46    ///
47    /// Returns a sanitized configuration error for an unknown mode or invalid
48    /// settings required by that mode.
49    pub fn from_env() -> Result<Self, ObservabilityError> {
50        match super::optional_env(RUNTIME_DIAGNOSTICS_MODE_ENV)?.as_deref() {
51            Some("disabled") => Ok(Self::Disabled),
52            Some("metrics_only") => Ok(Self::MetricsOnly {
53                sample_interval: super::parse_sample_interval()?,
54            }),
55            Some(mode @ ("endpoint_only" | "endpoint_and_metrics")) => {
56                let config = RuntimeDiagnosticsEndpointConfig::from_env()?.ok_or_else(|| {
57                    ObservabilityError::invalid_config(
58                        "runtime diagnostics endpoint mode requires bind and token settings",
59                    )
60                })?;
61                Ok(if mode == "endpoint_only" {
62                    Self::EndpointOnly(config)
63                } else {
64                    Self::EndpointAndMetrics(config)
65                })
66            }
67            Some(_) => Err(ObservabilityError::invalid_config("unknown runtime diagnostics mode")),
68            None => Ok(RuntimeDiagnosticsEndpointConfig::from_env()?.map_or(Self::Disabled, Self::EndpointAndMetrics)),
69        }
70    }
71
72    fn endpoint(&self) -> Option<&RuntimeDiagnosticsEndpointConfig> {
73        match self {
74            Self::EndpointOnly(config) | Self::EndpointAndMetrics(config) => Some(config),
75            Self::Disabled | Self::MetricsOnly { .. } => None,
76        }
77    }
78
79    fn sample_interval(&self) -> Option<Duration> {
80        match self {
81            Self::MetricsOnly { sample_interval } => Some(*sample_interval),
82            Self::EndpointAndMetrics(config) => Some(config.sample_interval),
83            Self::Disabled | Self::EndpointOnly(_) => None,
84        }
85    }
86}
87
88/// Bounded, synchronous access to diagnostics owned by the business component.
89///
90/// Implementations must read in-memory state without disk/network I/O, waiting
91/// for business completion, or unbounded scans. Retained data must not keep write
92/// authority alive after component shutdown. V2 keeps its existing convention:
93/// an empty schedule vector is omitted, including when it represents known zero.
94pub trait RuntimeDiagnosticsDataProvider: Send + Sync {
95    fn snapshot(&self) -> RuntimeDiagnosticsInputs;
96}
97
98impl<F> RuntimeDiagnosticsDataProvider for F
99where
100    F: Fn() -> RuntimeDiagnosticsInputs + Send + Sync,
101{
102    fn snapshot(&self) -> RuntimeDiagnosticsInputs {
103        self()
104    }
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub struct RuntimeDiagnosticsStarted {
109    pub endpoint: Option<RuntimeDiagnosticsEndpointHandle>,
110    pub sampler_active: bool,
111}
112
113struct StartedConfiguration {
114    mode: RuntimeDiagnosticsMode,
115    handle: RuntimeDiagnosticsStarted,
116}
117
118struct ServiceInner {
119    context: ChildServiceContext,
120    observed: ChildServiceContext,
121    component: RuntimeComponent,
122    telemetry: TelemetryHandle,
123    provider: Arc<dyn RuntimeDiagnosticsDataProvider>,
124    started: OnceCell<StartedConfiguration>,
125}
126
127/// One owned, idempotently initialized sampler and optional HTTP listener.
128///
129/// Clones share initialization. Keep one instance at the composition root;
130/// constructing another service creates a separate owner. Dropping this handle
131/// does not detach its tasks from the supplied parent. Shutdown cancels and
132/// awaits both outputs using the caller's deadline.
133#[derive(Clone)]
134pub struct RuntimeDiagnosticsService {
135    inner: Arc<ServiceInner>,
136}
137
138impl RuntimeDiagnosticsService {
139    pub fn new(
140        parent: &ChildServiceContext,
141        component: RuntimeComponent,
142        telemetry: TelemetryHandle,
143        provider: Arc<dyn RuntimeDiagnosticsDataProvider>,
144    ) -> Self {
145        Self {
146            inner: Arc::new(ServiceInner {
147                context: parent.component("runtime-diagnostics"),
148                observed: parent.clone(),
149                component,
150                telemetry,
151                provider,
152                started: OnceCell::new(),
153            }),
154        }
155    }
156
157    /// Starts exactly the requested outputs for this service instance.
158    ///
159    /// # Errors
160    ///
161    /// Returns a sanitized error for invalid sampling bounds, endpoint setup,
162    /// closed ownership, or reinitialization with a different configuration.
163    pub async fn start(&self, mode: RuntimeDiagnosticsMode) -> Result<RuntimeDiagnosticsStarted, ObservabilityError> {
164        if self.inner.context.task_group().lifecycle_state() != rocketmq_runtime::TaskGroupLifecycleState::Open {
165            return Err(ObservabilityError::invalid_config(
166                "runtime diagnostics owner is closed",
167            ));
168        }
169        if mode
170            .sample_interval()
171            .is_some_and(|interval| interval.is_zero() || interval > Duration::from_secs(300))
172        {
173            return Err(ObservabilityError::invalid_config(
174                "runtime diagnostics sample interval must be positive and at most 300 seconds",
175            ));
176        }
177        let started = self
178            .inner
179            .started
180            .get_or_try_init(|| self.initialize(mode.clone()))
181            .await?;
182        if started.mode != mode {
183            return Err(ObservabilityError::invalid_config(
184                "runtime diagnostics service already has a different configuration",
185            ));
186        }
187        Ok(started.handle)
188    }
189
190    async fn initialize(&self, mode: RuntimeDiagnosticsMode) -> Result<StartedConfiguration, ObservabilityError> {
191        let context = &self.inner.context;
192        let bound = if let Some(config) = mode.endpoint() {
193            read_token(context, config.token_file.clone()).await?;
194            let listener = tokio::net::TcpListener::bind(config.bind_addr)
195                .await
196                .map_err(|_| ObservabilityError::invalid_config("runtime diagnostics listener cannot bind"))?;
197            let local_addr = listener.local_addr().map_err(|_| {
198                ObservabilityError::invalid_config("runtime diagnostics listener address is unavailable")
199            })?;
200            Some((listener, local_addr, config.token_file.clone()))
201        } else {
202            None
203        };
204        let recorder = if mode.sample_interval().is_some() {
205            RuntimeMetricsRecorder::from_handle(&self.inner.telemetry, self.inner.component)
206        } else {
207            RuntimeMetricsRecorder::noop(self.inner.component)
208        };
209        let sample_interval = mode.sample_interval().filter(|_| recorder.is_enabled());
210        let sampler_active = sample_interval.is_some();
211        if let Some(sample_interval) = sample_interval {
212            let sample_context = self.inner.observed.clone();
213            let component = self.inner.component;
214            let provider = self.inner.provider.clone();
215            let metrics = recorder.clone();
216            context
217                .scheduled_tasks("sampler")
218                .schedule(
219                    rocketmq_runtime::ScheduledTaskConfig::fixed_delay("runtime-diagnostics.sample", sample_interval),
220                    rocketmq_runtime::ScheduledExecutionPolicy::serial(rocketmq_runtime::MissedTickPolicy::Skip),
221                    move || {
222                        let context = sample_context.clone();
223                        let provider = provider.clone();
224                        let metrics = metrics.clone();
225                        async move {
226                            metrics.record_snapshot_v2(&context.diagnostics_view_v2(component, provider.snapshot()));
227                        }
228                    },
229                )
230                .map_err(|_| ObservabilityError::invalid_config("runtime diagnostics sampler cannot start"))?;
231        }
232        let endpoint = if let Some((listener, local_addr, token_file)) = bound {
233            let cancellation = context.task_group().cancellation_token();
234            let endpoint_context = super::EndpointContext {
235                owner: context.clone(),
236                observed: self.inner.observed.clone(),
237                component: self.inner.component,
238                token_file,
239                metrics: recorder,
240                provider: self.inner.provider.clone(),
241            };
242            if context
243                .spawn_service("runtime-diagnostics.endpoint", async move {
244                    serve(listener, endpoint_context, cancellation).await;
245                })
246                .is_err()
247            {
248                context.task_group().shutdown(Duration::ZERO).await;
249                return Err(ObservabilityError::invalid_config(
250                    "runtime diagnostics endpoint cannot start",
251                ));
252            }
253            Some(RuntimeDiagnosticsEndpointHandle { local_addr })
254        } else {
255            None
256        };
257        Ok(StartedConfiguration {
258            mode,
259            handle: RuntimeDiagnosticsStarted {
260                endpoint,
261                sampler_active,
262            },
263        })
264    }
265
266    pub async fn shutdown_until(&self, deadline: ShutdownDeadline) -> ShutdownReport {
267        self.inner.context.task_group().shutdown_until(deadline).await
268    }
269}
rocketmq_observability/runtime_diagnostics/service.rs

rocketmq_observability/runtime_diagnostics/
service.rs
