triton_distributed/component/
registry.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 super::{Component, Registry, Result};
17use async_once_cell::OnceCell;
18use std::{
19    collections::HashMap,
20    sync::{Arc, Weak},
21};
22use tokio::sync::Mutex;
23
24impl Default for Registry {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl Registry {
31    pub fn new() -> Self {
32        Self {
33            services: Arc::new(Mutex::new(HashMap::new())),
34        }
35    }
36}
37
38// impl ComponentRegistry {
39//     pub fn new() -> Self {
40//         Self {
41//             clients: Arc::new(Mutex::new(HashMap::new())),
42//         }
43//     }
44
45//     pub async fn get_or_create(&mut self, component: Component) -> Result<Arc<Client>> {
46//         // Lock the clients HashMap for thread-safe access
47//         let mut guard = self.clients.lock().await;
48
49//         // Check if the component already exists in the registry
50//         if let Some(weak) = guard.get(&component.slug()) {
51//             // Attempt to upgrade the Weak pointer
52//             if let Some(client) = weak.upgrade() {
53//                 return Ok(client);
54//             }
55//         }
56
57//         // Fallback: Create a new Client
58//         let client = component.client().await?;
59
60//         // Insert a Weak reference to the new client into the map
61//         guard.insert(component.slug(), Arc::downgrade(&client));
62
63//         Ok(client)
64//     }
65// }
66
67// #[derive(Clone)]
68// pub struct ServiceRegistry {
69//     clients: Arc<Mutex<HashMap<String, Arc<Service>>>>,
70// }
71
72// impl ServiceRegistry {
73//     pub fn new() -> Self {
74//         Self {
75//             clients: Arc::new(Mutex::new(HashMap::new())),
76//         }
77//     }
78
79//     pub async fn get_or_create(&mut self, component: Component) -> Result<Arc<Client>> {
80//         // Lock the clients HashMap for thread-safe access
81//         let mut guard = self.clients.lock().await;
82
83//         // Check if the component already exists in the registry
84//         if let Some(weak) = guard.get(&component.slug()) {
85//             // Attempt to upgrade the Weak pointer
86//             if let Some(client) = weak.upgrade() {
87//                 return Ok(client);
88//             }
89//         }
90
91//         // Fallback: Create a new Client
92//         let client = component.client().await?;
93
94//         // Insert a Weak reference to the new client into the map
95//         guard.insert(component.slug(), Arc::downgrade(&client));
96
97//         Ok(client)
98//     }
99// }
triton_distributed/component/registry.rs

triton_distributed/component/
registry.rs
