dynamo_runtime/component/
service.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use derive_getters::Dissolve;
17use std::collections::HashMap;
18use std::sync::Mutex;
19
20use super::*;
21
22pub use super::endpoint::EndpointStats;
23pub type StatsHandler =
24    Box<dyn FnMut(String, EndpointStats) -> serde_json::Value + Send + Sync + 'static>;
25pub type EndpointStatsHandler =
26    Box<dyn FnMut(EndpointStats) -> serde_json::Value + Send + Sync + 'static>;
27
28pub const PROJECT_NAME: &str = "Dynamo";
29
30#[derive(Educe, Builder, Dissolve)]
31#[educe(Debug)]
32#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
33pub struct ServiceConfig {
34    #[builder(private)]
35    component: Component,
36
37    /// Description
38    #[builder(default)]
39    description: Option<String>,
40}
41
42impl ServiceConfigBuilder {
43    /// Create the [`Component`]'s service and store it in the registry.
44    pub async fn create(self) -> Result<Component> {
45        let (component, description) = self.build_internal()?.dissolve();
46
47        let version = "0.0.1".to_string();
48
49        let service_name = component.service_name();
50        log::debug!("component: {component}; creating, service_name: {service_name}");
51
52        let description = description.unwrap_or(format!(
53            "{PROJECT_NAME} component {} in namespace {}",
54            component.name, component.namespace
55        ));
56
57        let stats_handler_registry: Arc<Mutex<HashMap<String, EndpointStatsHandler>>> =
58            Arc::new(Mutex::new(HashMap::new()));
59
60        let stats_handler_registry_clone = stats_handler_registry.clone();
61
62        let mut guard = component.drt.component_registry.inner.lock().await;
63
64        if guard.services.contains_key(&service_name) {
65            return Err(anyhow::anyhow!("Service already exists"));
66        }
67
68        // create service on the secondary runtime
69        let builder = component.drt.nats_client.client().service_builder();
70
71        tracing::debug!("Starting service: {}", service_name);
72        let service_builder = builder
73            .description(description)
74            .stats_handler(move |name, stats| {
75                log::trace!("stats_handler: {name}, {stats:?}");
76                let mut guard = stats_handler_registry.lock().unwrap();
77                match guard.get_mut(&name) {
78                    Some(handler) => handler(stats),
79                    None => serde_json::Value::Null,
80                }
81            });
82        tracing::debug!("Got builder");
83        let service = service_builder
84            .start(service_name.clone(), version)
85            .await
86            .map_err(|e| anyhow::anyhow!("Failed to start service: {e}"))?;
87
88        tracing::debug!("Service started TEMP");
89
90        // new copy of service_name as the previous one is moved into the task above
91        let service_name = component.service_name();
92
93        // insert the service into the registry
94        guard.services.insert(service_name.clone(), service);
95
96        // insert the stats handler into the registry
97        guard
98            .stats_handlers
99            .insert(service_name, stats_handler_registry_clone);
100
101        // drop the guard to unlock the mutex
102        drop(guard);
103
104        tracing::debug!("create done");
105        Ok(component)
106    }
107}
108
109impl ServiceConfigBuilder {
110    pub(crate) fn from_component(component: Component) -> Self {
111        Self::default().component(component)
112    }
113}