agave_thread_manager/
config.rs

1use {
2    crate::{CoreAllocation, NativeConfig, RayonConfig, TokioConfig},
3    serde::{Deserialize, Serialize},
4    std::collections::HashMap,
5    toml::de::Error,
6};
7
8#[derive(Clone, Debug, Serialize, Deserialize)]
9#[serde(default)]
10pub struct ThreadManagerConfig {
11    /// Configurations for the native OS thread pools
12    pub native_configs: HashMap<String, NativeConfig>,
13    /// Mapping from lookup name to the config name for Native
14    pub native_runtime_mapping: HashMap<String, String>,
15
16    /// Configurations for the rayon runtimes
17    pub rayon_configs: HashMap<String, RayonConfig>,
18    /// Mapping from lookup name to the config name for Rayon
19    pub rayon_runtime_mapping: HashMap<String, String>,
20
21    /// Configurations for the Tokio runtimes
22    pub tokio_configs: HashMap<String, TokioConfig>,
23    /// Mapping from lookup name to the config name for Tokio
24    pub tokio_runtime_mapping: HashMap<String, String>,
25
26    pub default_core_allocation: CoreAllocation,
27}
28
29impl Default for ThreadManagerConfig {
30    // Default is mostly to keep Deserialize implementation happy. This is not intended to be a valid config
31    fn default() -> Self {
32        Self {
33            native_configs: HashMap::from([("default".to_owned(), NativeConfig::default())]),
34            native_runtime_mapping: HashMap::new(),
35            rayon_configs: HashMap::from([("default".to_owned(), RayonConfig::default())]),
36            rayon_runtime_mapping: HashMap::new(),
37            tokio_configs: HashMap::from([("default".to_owned(), TokioConfig::default())]),
38            tokio_runtime_mapping: HashMap::new(),
39            default_core_allocation: CoreAllocation::OsDefault,
40        }
41    }
42}
43
44impl ThreadManagerConfig {
45    /// Returns a valid "built-in" configuration for Agave that is safe to use in most cases.
46    /// Most operators will likely use this.
47    pub fn default_for_agave() -> Self {
48        let config_str = include_str!("default_config.toml");
49        toml::from_str(config_str).expect("Parsing the built-in config should never fail")
50    }
51
52    /// Loads configuration from a given str, filling in the gaps from what default_for_agave returns
53    pub fn from_toml_str(toml: &str) -> Result<Self, Error> {
54        let mut loaded: ThreadManagerConfig = toml::from_str(toml)?;
55        let mut result = Self::default_for_agave();
56        macro_rules! upsert {
57            ($name:ident) => {
58                result.$name.extend(loaded.$name.drain());
59            };
60        }
61        upsert!(native_configs);
62        upsert!(native_runtime_mapping);
63        upsert!(rayon_configs);
64        upsert!(rayon_runtime_mapping);
65        upsert!(tokio_configs);
66        upsert!(tokio_runtime_mapping);
67        result.default_core_allocation = loaded.default_core_allocation;
68        Ok(result)
69    }
70}