Line | Count | Source |
1 | // Copyright 2024 The NativeLink Authors. All rights reserved. | |
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 | ||
15 | use std::collections::HashMap; | |
16 | ||
17 | use nativelink_metric::{MetricsComponent, RootMetricsComponent}; | |
18 | use nativelink_util::store_trait::Store; | |
19 | use parking_lot::RwLock; | |
20 | ||
21 | #[derive(MetricsComponent)] | |
22 | pub struct StoreManager { | |
23 | #[metric] | |
24 | stores: RwLock<HashMap<String, Store>>, | |
25 | } | |
26 | ||
27 | impl StoreManager { | |
28 | 33 | pub fn new() -> StoreManager { |
29 | 33 | StoreManager { |
30 | 33 | stores: RwLock::new(HashMap::new()), |
31 | 33 | } |
32 | 33 | } |
33 | ||
34 | 42 | pub fn add_store(&self, name: &str, store: Store) { |
35 | 42 | let mut stores = self.stores.write(); |
36 | 42 | stores.insert(name.to_string(), store); |
37 | 42 | } |
38 | ||
39 | 52 | pub fn get_store(&self, name: &str) -> Option<Store> { |
40 | 52 | let stores = self.stores.read(); |
41 | 52 | if let Some(store) = stores.get(name) { |
42 | 52 | return Some(store.clone()); |
43 | 0 | } |
44 | 0 | None |
45 | 52 | } |
46 | } | |
47 | ||
48 | impl RootMetricsComponent for StoreManager {} | |
49 | ||
50 | impl Default for StoreManager { | |
51 | 0 | fn default() -> Self { |
52 | 0 | Self::new() |
53 | 0 | } |
54 | } |