service.rs 16.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// 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.
Ryan Olson's avatar
Ryan Olson committed
15
16
17
18
19
20
21

// TODO - refactor this entire module
//
// we want to carry forward the concept of live vs ready for the components
// we will want to associate the components cancellation token with the
// component's "service state"

22
23
24
use crate::{
    component::Component,
    error,
25
    metrics::{prometheus_names, prometheus_names::nats_service, MetricsRegistry},
26
27
28
29
30
    traits::*,
    transports::nats,
    utils::stream,
    DistributedRuntime, Result,
};
Ryan Olson's avatar
Ryan Olson committed
31
32
33
34
35

use async_nats::Message;
use async_stream::try_stream;
use bytes::Bytes;
use derive_getters::Dissolve;
Ryan Olson's avatar
Ryan Olson committed
36
use futures::stream::{StreamExt, TryStreamExt};
37
use prometheus;
Ryan Olson's avatar
Ryan Olson committed
38
39
40
41
42
43
44
45
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::time::Duration;

pub struct ServiceClient {
    nats_client: nats::Client,
}

impl ServiceClient {
46
    pub fn new(nats_client: nats::Client) -> Self {
Ryan Olson's avatar
Ryan Olson committed
47
48
49
50
        ServiceClient { nats_client }
    }
}

51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
/// ServiceSet contains a collection of services with their endpoints and metrics
///
/// Tree structure:
/// Structure:
/// - ServiceSet
///   - services: Vec<ServiceInfo>
///     - name: String
///     - id: String
///     - version: String
///     - started: String
///     - endpoints: Vec<EndpointInfo>
///       - name: String
///       - subject: String
///       - data: Option<NatsStatsMetrics>
///         - average_processing_time: f64
///         - last_error: String
///         - num_errors: u64
///         - num_requests: u64
///         - processing_time: u64
///         - queue_group: String
///         - data: serde_json::Value (custom stats)
Ryan Olson's avatar
Ryan Olson committed
72
#[derive(Debug, Clone, Serialize, Deserialize)]
Ryan Olson's avatar
Ryan Olson committed
73
74
75
76
pub struct ServiceSet {
    services: Vec<ServiceInfo>,
}

77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
/// This is a example JSON from `nats req '$SRV.STATS.dynamo_backend'`:
/// {
///   "type": "io.nats.micro.v1.stats_response",
///   "name": "dynamo_backend",
///   "id": "bdu7nA8tbhy9mEkxIWlkBA",
///   "version": "0.0.1",
///   "started": "2025-08-08T05:07:17.720783523Z",
///   "endpoints": [
///     {
///       "name": "dynamo_backend-generate-694d988806b92e39",
///       "subject": "dynamo_backend.generate-694d988806b92e39",
///       "num_requests": 0,
///       "num_errors": 0,
///       "processing_time": 0,
///       "average_processing_time": 0,
///       "last_error": "",
///       "data": {
///         "val": 10
///       },
///       "queue_group": "q"
///     }
///   ]
/// }
Ryan Olson's avatar
Ryan Olson committed
100
101
102
103
104
105
106
107
108
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceInfo {
    pub name: String,
    pub id: String,
    pub version: String,
    pub started: String,
    pub endpoints: Vec<EndpointInfo>,
}

109
/// Each endpoint has name, subject, num_requests, num_errors, processing_time, average_processing_time, last_error, queue_group, and data
Ryan Olson's avatar
Ryan Olson committed
110
111
112
113
114
#[derive(Debug, Clone, Serialize, Deserialize, Dissolve)]
pub struct EndpointInfo {
    pub name: String,
    pub subject: String,

115
    /// Extra fields that don't fit in EndpointInfo will be flattened into the Metrics struct.
Ryan Olson's avatar
Ryan Olson committed
116
    #[serde(flatten)]
117
    pub data: Option<NatsStatsMetrics>,
Ryan Olson's avatar
Ryan Olson committed
118
119
}

Ryan Olson's avatar
Ryan Olson committed
120
121
122
123
124
impl EndpointInfo {
    pub fn id(&self) -> Result<i64> {
        let id = self
            .subject
            .split('-')
125
            .next_back()
Ryan Olson's avatar
Ryan Olson committed
126
127
128
129
130
            .ok_or_else(|| error!("No id found in subject"))?;

        i64::from_str_radix(id, 16).map_err(|e| error!("Invalid id format: {}", e))
    }
}
131
132
133
134
135
136

// TODO: This is _really_ close to the async_nats::service::Stats object,
// but it's missing a few fields like "name", so use a temporary struct
// for easy deserialization. Ideally, this type already exists or can
// be exposed in the library somewhere.
/// Stats structure returned from NATS service API
137
/// https://github.com/nats-io/nats.rs/blob/main/async-nats/src/service/endpoint.rs
Ryan Olson's avatar
Ryan Olson committed
138
#[derive(Debug, Clone, Serialize, Deserialize, Dissolve)]
139
140
141
pub struct NatsStatsMetrics {
    // Standard NATS Stats Service API fields from $SRV.STATS.<service_name> requests
    pub average_processing_time: u64, // in nanoseconds according to nats-io
142
143
144
    pub last_error: String,
    pub num_errors: u64,
    pub num_requests: u64,
145
    pub processing_time: u64, // in nanoseconds according to nats-io
146
147
148
149
    pub queue_group: String,
    // Field containing custom stats handler data
    pub data: serde_json::Value,
}
Ryan Olson's avatar
Ryan Olson committed
150

151
impl NatsStatsMetrics {
Ryan Olson's avatar
Ryan Olson committed
152
    pub fn decode<T: for<'de> Deserialize<'de>>(self) -> Result<T> {
153
        serde_json::from_value(self.data).map_err(Into::into)
Ryan Olson's avatar
Ryan Olson committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    }
}

impl ServiceClient {
    pub async fn unary(
        &self,
        subject: impl Into<String>,
        payload: impl Into<Bytes>,
    ) -> Result<Message> {
        let response = self
            .nats_client
            .client()
            .request(subject.into(), payload.into())
            .await?;
        Ok(response)
    }

171
172
173
    pub async fn collect_services(
        &self,
        service_name: &str,
174
        timeout: Duration,
175
    ) -> Result<ServiceSet> {
Ryan Olson's avatar
Ryan Olson committed
176
        let sub = self.nats_client.scrape_service(service_name).await?;
177
178
        if timeout.is_zero() {
            tracing::warn!("collect_services: timeout is zero");
Ryan Olson's avatar
Ryan Olson committed
179
        }
180
181
        if timeout > Duration::from_secs(10) {
            tracing::warn!("collect_services: timeout is greater than 10 seconds");
Ryan Olson's avatar
Ryan Olson committed
182
        }
183
        let deadline = tokio::time::Instant::now() + timeout;
Ryan Olson's avatar
Ryan Olson committed
184

185
186
187
188
189
190
191
192
193
194
195
196
197
198
        let mut services = vec![];
        let mut s = stream::until_deadline(sub, deadline);
        while let Some(message) = s.next().await {
            if message.payload.is_empty() {
                // Expected while we wait for KV metrics in worker to start
                tracing::trace!(service_name, "collect_services: empty payload from nats");
                continue;
            }
            let info = serde_json::from_slice::<ServiceInfo>(&message.payload);
            match info {
                Ok(info) => services.push(info),
                Err(err) => {
                    let payload = String::from_utf8_lossy(&message.payload);
                    tracing::debug!(%err, service_name, %payload, "error decoding service info");
Ryan Olson's avatar
Ryan Olson committed
199
                }
200
201
            }
        }
Ryan Olson's avatar
Ryan Olson committed
202

Ryan Olson's avatar
Ryan Olson committed
203
        Ok(ServiceSet { services })
Ryan Olson's avatar
Ryan Olson committed
204
205
206
207
208
209
210
211
212
    }
}

impl ServiceSet {
    pub fn into_endpoints(self) -> impl Iterator<Item = EndpointInfo> {
        self.services
            .into_iter()
            .flat_map(|s| s.endpoints.into_iter())
    }
213
214
215
216
217

    /// Get a reference to the services in this ServiceSet
    pub fn services(&self) -> &[ServiceInfo] {
        &self.services
    }
Ryan Olson's avatar
Ryan Olson committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_service_set() {
        let services = vec![
            ServiceInfo {
                name: "service1".to_string(),
                id: "1".to_string(),
                version: "1.0".to_string(),
                started: "2021-01-01".to_string(),
                endpoints: vec![
                    EndpointInfo {
                        name: "endpoint1".to_string(),
                        subject: "subject1".to_string(),
237
238
                        data: Some(NatsStatsMetrics {
                            average_processing_time: 100_000, // 0.1ms = 100,000 nanoseconds
239
240
241
242
243
244
245
                            last_error: "none".to_string(),
                            num_errors: 0,
                            num_requests: 10,
                            processing_time: 100,
                            queue_group: "group1".to_string(),
                            data: serde_json::json!({"key": "value1"}),
                        }),
Ryan Olson's avatar
Ryan Olson committed
246
247
248
249
                    },
                    EndpointInfo {
                        name: "endpoint2-foo".to_string(),
                        subject: "subject2".to_string(),
250
251
                        data: Some(NatsStatsMetrics {
                            average_processing_time: 100_000, // 0.1ms = 100,000 nanoseconds
252
253
254
255
256
257
258
                            last_error: "none".to_string(),
                            num_errors: 0,
                            num_requests: 10,
                            processing_time: 100,
                            queue_group: "group1".to_string(),
                            data: serde_json::json!({"key": "value1"}),
                        }),
Ryan Olson's avatar
Ryan Olson committed
259
260
261
262
263
264
265
266
267
268
269
270
                    },
                ],
            },
            ServiceInfo {
                name: "service1".to_string(),
                id: "2".to_string(),
                version: "1.0".to_string(),
                started: "2021-01-01".to_string(),
                endpoints: vec![
                    EndpointInfo {
                        name: "endpoint1".to_string(),
                        subject: "subject1".to_string(),
271
272
                        data: Some(NatsStatsMetrics {
                            average_processing_time: 100_000, // 0.1ms = 100,000 nanoseconds
273
274
275
276
277
278
279
                            last_error: "none".to_string(),
                            num_errors: 0,
                            num_requests: 10,
                            processing_time: 100,
                            queue_group: "group1".to_string(),
                            data: serde_json::json!({"key": "value1"}),
                        }),
Ryan Olson's avatar
Ryan Olson committed
280
281
282
283
                    },
                    EndpointInfo {
                        name: "endpoint2-bar".to_string(),
                        subject: "subject2".to_string(),
284
285
                        data: Some(NatsStatsMetrics {
                            average_processing_time: 100_000, // 0.1ms = 100,000 nanoseconds
286
287
288
289
290
291
292
                            last_error: "none".to_string(),
                            num_errors: 0,
                            num_requests: 10,
                            processing_time: 100,
                            queue_group: "group1".to_string(),
                            data: serde_json::json!({"key": "value2"}),
                        }),
Ryan Olson's avatar
Ryan Olson committed
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
                    },
                ],
            },
        ];

        let service_set = ServiceSet { services };

        let endpoints: Vec<_> = service_set
            .into_endpoints()
            .filter(|e| e.name.starts_with("endpoint2"))
            .collect();

        assert_eq!(endpoints.len(), 2);
    }
}
308
309
310
311
312
313
314
315
316
317
318
319
320

/// Prometheus metrics for component service statistics (ordered to match NatsStatsMetrics)
///
/// ⚠️  IMPORTANT: These Prometheus Gauges are COPIES of NATS data, not live references!
///
/// How it works:
/// 1. NATS provides source data via NatsStatsMetrics
/// 2. Metrics callbacks read current NATS values and update these Prometheus Gauges
/// 3. Prometheus scrapes these Gauge values (snapshots, not live data)
///
/// Flow: NATS Service → NatsStatsMetrics (Counters) → Metrics Callback → Prometheus Gauge
/// Note: These are snapshots updated when execute_metrics_callbacks() is called.
#[derive(Debug, Clone)]
321
pub struct ComponentNatsServerPrometheusMetrics {
322
    /// Average processing time in milliseconds (maps to: average_processing_time)
323
    pub service_avg_processing_ms: prometheus::Gauge,
324
    /// Total errors across all endpoints (maps to: num_errors)
325
    pub service_total_errors: prometheus::IntGauge,
326
    /// Total requests across all endpoints (maps to: num_requests)
327
    pub service_total_requests: prometheus::IntGauge,
328
    /// Total processing time in milliseconds (maps to: processing_time)
329
    pub service_total_processing_ms: prometheus::IntGauge,
330
    /// Number of active services (derived from ServiceSet.services)
331
    pub service_active_services: prometheus::IntGauge,
332
    /// Number of active endpoints (derived from ServiceInfo.endpoints)
333
    pub service_active_endpoints: prometheus::IntGauge,
334
335
}

336
impl ComponentNatsServerPrometheusMetrics {
337
338
    /// Create new ComponentServiceMetrics using Component's DistributedRuntime's Prometheus constructors
    pub fn new(component: &Component) -> Result<Self> {
339
340
341
342
343
344
345
346
347
348
349
350
351
352
        let service_name = component.service_name();

        // Build labels: service_name first, then component's labels
        let mut labels_vec = vec![("service_name", service_name.as_str())];

        // Add component's labels (convert from (String, String) to (&str, &str))
        for (key, value) in component.labels() {
            labels_vec.push((key.as_str(), value.as_str()));
        }

        let labels: &[(&str, &str)] = &labels_vec;

        let service_avg_processing_ms = component.create_gauge(
            nats_service::AVG_PROCESSING_MS,
353
            "Average processing time across all component endpoints in milliseconds",
354
            labels,
355
356
        )?;

357
358
        let service_total_errors = component.create_intgauge(
            nats_service::TOTAL_ERRORS,
359
            "Total number of errors across all component endpoints",
360
            labels,
361
362
        )?;

363
364
        let service_total_requests = component.create_intgauge(
            nats_service::TOTAL_REQUESTS,
365
            "Total number of requests across all component endpoints",
366
            labels,
367
368
        )?;

369
370
        let service_total_processing_ms = component.create_intgauge(
            nats_service::TOTAL_PROCESSING_MS,
371
            "Total processing time across all component endpoints in milliseconds",
372
            labels,
373
374
        )?;

375
376
        let service_active_services = component.create_intgauge(
            nats_service::ACTIVE_SERVICES,
377
            "Number of active services in this component",
378
            labels,
379
380
        )?;

381
382
        let service_active_endpoints = component.create_intgauge(
            nats_service::ACTIVE_ENDPOINTS,
383
            "Number of active endpoints across all services",
384
            labels,
385
386
387
        )?;

        Ok(Self {
388
389
390
391
392
393
            service_avg_processing_ms,
            service_total_errors,
            service_total_requests,
            service_total_processing_ms,
            service_active_services,
            service_active_endpoints,
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
        })
    }

    /// Update metrics from scraped ServiceSet data
    pub fn update_from_service_set(&self, service_set: &ServiceSet) {
        // Variables ordered to match NatsStatsMetrics fields
        let mut processing_time_samples = 0u64; // for average_processing_time calculation
        let mut total_errors = 0u64; // maps to: num_errors
        let mut total_requests = 0u64; // maps to: num_requests
        let mut total_processing_time_nanos = 0u64; // maps to: processing_time (nanoseconds from NATS)
        let mut endpoint_count = 0u64; // for derived metrics

        let service_count = service_set.services().len() as i64;

        for service in service_set.services() {
            for endpoint in &service.endpoints {
                endpoint_count += 1;

                if let Some(ref stats) = endpoint.data {
                    total_errors += stats.num_errors;
                    total_requests += stats.num_requests;
                    total_processing_time_nanos += stats.processing_time;

                    if stats.num_requests > 0 {
                        processing_time_samples += 1;
                    }
                }
            }
        }

        // Update metrics (ordered to match NatsStatsMetrics fields)
        // Calculate average processing time in milliseconds (maps to: average_processing_time)
        if processing_time_samples > 0 && total_requests > 0 {
            let avg_time_nanos = total_processing_time_nanos as f64 / total_requests as f64;
            let avg_time_ms = avg_time_nanos / 1_000_000.0; // Convert nanoseconds to milliseconds
429
            self.service_avg_processing_ms.set(avg_time_ms);
430
        } else {
431
            self.service_avg_processing_ms.set(0.0);
432
433
        }

434
435
436
        self.service_total_errors.set(total_errors as i64); // maps to: num_errors
        self.service_total_requests.set(total_requests as i64); // maps to: num_requests
        self.service_total_processing_ms
437
            .set((total_processing_time_nanos / 1_000_000) as i64); // maps to: processing_time (converted to milliseconds)
438
439
        self.service_active_services.set(service_count); // derived from ServiceSet.services
        self.service_active_endpoints.set(endpoint_count as i64); // derived from ServiceInfo.endpoints
440
441
442
443
    }

    /// Reset all metrics to zero. Useful when no data is available or to clear stale values.
    pub fn reset_to_zeros(&self) {
444
445
446
447
448
449
        self.service_avg_processing_ms.set(0.0);
        self.service_total_errors.set(0);
        self.service_total_requests.set(0);
        self.service_total_processing_ms.set(0);
        self.service_active_services.set(0);
        self.service_active_endpoints.set(0);
450
451
    }
}