litellm_rs/utils/sys/
di.rs1#![allow(dead_code)] use crate::utils::error::{GatewayError, Result};
9use parking_lot::RwLock;
10use std::any::{Any, TypeId};
11use std::collections::HashMap;
12use std::sync::Arc;
13
14pub trait Injectable: Send + Sync + 'static {
16 fn type_name(&self) -> &'static str {
18 std::any::type_name::<Self>()
19 }
20}
21
22impl<T> Injectable for T where T: Send + Sync + 'static {}
24
25#[derive(Default)]
27pub struct ServiceContainer {
28 services: RwLock<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
29 singletons: RwLock<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
30}
31
32impl ServiceContainer {
33 pub fn new() -> Self {
35 Self::default()
36 }
37
38 pub fn register<T>(&self, service: T) -> Result<()>
40 where
41 T: Injectable,
42 {
43 let type_id = TypeId::of::<T>();
44 let service = Arc::new(service);
45
46 let mut services = self.services.write();
47 if services.contains_key(&type_id) {
48 return Err(GatewayError::Config(format!(
49 "Service {} is already registered",
50 std::any::type_name::<T>()
51 )));
52 }
53
54 services.insert(type_id, service);
55 Ok(())
56 }
57
58 pub fn register_singleton<T>(&self, service: T) -> Result<()>
60 where
61 T: Injectable,
62 {
63 let type_id = TypeId::of::<T>();
64 let service = Arc::new(service);
65
66 let mut singletons = self.singletons.write();
67 if singletons.contains_key(&type_id) {
68 return Err(GatewayError::Config(format!(
69 "Singleton {} is already registered",
70 std::any::type_name::<T>()
71 )));
72 }
73
74 singletons.insert(type_id, service);
75 Ok(())
76 }
77
78 pub fn register_factory<T, F>(&self, factory: F) -> Result<()>
80 where
81 T: Injectable,
82 F: Fn() -> T + Send + Sync + 'static,
83 {
84 let service = factory();
85 self.register(service)
86 }
87
88 pub fn get<T>(&self) -> Result<Arc<T>>
90 where
91 T: Injectable,
92 {
93 let type_id = TypeId::of::<T>();
94
95 {
97 let singletons = self.singletons.read();
98 if let Some(service) = singletons.get(&type_id) {
99 return service.clone().downcast::<T>().map_err(|_| {
100 GatewayError::Internal(format!(
101 "Failed to downcast singleton service {}",
102 std::any::type_name::<T>()
103 ))
104 });
105 }
106 }
107
108 {
110 let services = self.services.read();
111 if let Some(service) = services.get(&type_id) {
112 return service.clone().downcast::<T>().map_err(|_| {
113 GatewayError::Internal(format!(
114 "Failed to downcast service {}",
115 std::any::type_name::<T>()
116 ))
117 });
118 }
119 }
120
121 Err(GatewayError::Config(format!(
122 "Service {} is not registered",
123 std::any::type_name::<T>()
124 )))
125 }
126
127 pub fn try_get<T>(&self) -> Option<Arc<T>>
129 where
130 T: Injectable,
131 {
132 self.get().ok()
133 }
134
135 pub fn contains<T>(&self) -> bool
137 where
138 T: Injectable,
139 {
140 let type_id = TypeId::of::<T>();
141 let singletons = self.singletons.read();
142 let services = self.services.read();
143
144 singletons.contains_key(&type_id) || services.contains_key(&type_id)
145 }
146
147 pub fn remove<T>(&self) -> Result<()>
149 where
150 T: Injectable,
151 {
152 let type_id = TypeId::of::<T>();
153
154 {
155 let mut singletons = self.singletons.write();
156 if singletons.remove(&type_id).is_some() {
157 return Ok(());
158 }
159 }
160
161 {
162 let mut services = self.services.write();
163 if services.remove(&type_id).is_some() {
164 return Ok(());
165 }
166 }
167
168 Err(GatewayError::Config(format!(
169 "Service {} is not registered",
170 std::any::type_name::<T>()
171 )))
172 }
173
174 pub fn clear(&self) {
176 self.services.write().clear();
177 self.singletons.write().clear();
178 }
179
180 pub fn service_count(&self) -> usize {
182 let services = self.services.read();
183 let singletons = self.singletons.read();
184 services.len() + singletons.len()
185 }
186}
187
188pub struct ServiceBuilder {
190 container: ServiceContainer,
191}
192
193impl ServiceBuilder {
194 pub fn new() -> Self {
196 Self {
197 container: ServiceContainer::new(),
198 }
199 }
200
201 pub fn add_service<T>(self, service: T) -> Result<Self>
203 where
204 T: Injectable,
205 {
206 self.container.register(service)?;
207 Ok(self)
208 }
209
210 pub fn add_singleton<T>(self, service: T) -> Result<Self>
212 where
213 T: Injectable,
214 {
215 self.container.register_singleton(service)?;
216 Ok(self)
217 }
218
219 pub fn add_factory<T, F>(self, factory: F) -> Result<Self>
221 where
222 T: Injectable,
223 F: Fn() -> T + Send + Sync + 'static,
224 {
225 self.container.register_factory(factory)?;
226 Ok(self)
227 }
228
229 pub fn build(self) -> ServiceContainer {
231 self.container
232 }
233}
234
235impl Default for ServiceBuilder {
236 fn default() -> Self {
237 Self::new()
238 }
239}
240
241pub trait WithDependencies {
243 fn with_dependencies(self, container: &ServiceContainer) -> Result<Self>
245 where
246 Self: Sized;
247}
248
249use once_cell::sync::Lazy;
253
254static GLOBAL_CONTAINER: Lazy<ServiceContainer> = Lazy::new(ServiceContainer::new);
255
256pub fn global_container() -> &'static ServiceContainer {
258 &GLOBAL_CONTAINER
259}
260
261pub fn register_global<T>(service: T) -> Result<()>
263where
264 T: Injectable,
265{
266 global_container().register(service)
267}
268
269pub fn register_global_singleton<T>(service: T) -> Result<()>
271where
272 T: Injectable,
273{
274 global_container().register_singleton(service)
275}
276
277pub fn get_global<T>() -> Result<Arc<T>>
279where
280 T: Injectable,
281{
282 global_container().get()
283}