reflex/lifecycle/
config.rs1use std::env;
2use std::path::PathBuf;
3use std::time::Duration;
4
5use super::error::LifecycleResult;
6
7pub const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 15 * 60;
9pub const DEFAULT_SNAPSHOT_FILENAME: &str = "snapshot.rkyv";
11pub const REAPER_CHECK_INTERVAL_SECS: u64 = 30;
13
14#[derive(Debug, Clone, PartialEq, Eq, Default)]
15pub enum CloudProviderType {
17 #[default]
18 Gcp,
20 Local,
22}
23
24impl std::str::FromStr for CloudProviderType {
25 type Err = String;
26
27 fn from_str(s: &str) -> Result<Self, Self::Err> {
28 match s.to_lowercase().as_str() {
29 "gcp" | "google" => Ok(Self::Gcp),
30 "local" => Ok(Self::Local),
31 _ => Err(format!("Unknown cloud provider: {}", s)),
32 }
33 }
34}
35
36#[derive(Debug, Clone)]
37pub struct LifecycleConfig {
39 pub gcs_bucket: String,
41 pub local_snapshot_path: PathBuf,
43 pub idle_timeout: Duration,
45 pub enable_instance_stop: bool,
47 pub cloud_provider: CloudProviderType,
49 pub cloud_region: Option<String>,
51}
52
53impl Default for LifecycleConfig {
54 fn default() -> Self {
55 Self {
56 gcs_bucket: String::new(),
57 local_snapshot_path: PathBuf::from("/mnt/nvme/reflex_data")
58 .join(DEFAULT_SNAPSHOT_FILENAME),
59 idle_timeout: Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS),
60 enable_instance_stop: true,
61 cloud_provider: CloudProviderType::default(),
62 cloud_region: None,
63 }
64 }
65}
66
67impl LifecycleConfig {
68 const ENV_GCS_BUCKET: &'static str = "REFLEX_GCS_BUCKET";
69 const ENV_LOCAL_SNAPSHOT_PATH: &'static str = "REFLEX_SNAPSHOT_PATH";
70 const ENV_IDLE_TIMEOUT_SECS: &'static str = "REFLEX_IDLE_TIMEOUT_SECS";
71 const ENV_ENABLE_INSTANCE_STOP: &'static str = "REFLEX_ENABLE_INSTANCE_STOP";
72 const ENV_CLOUD_PROVIDER: &'static str = "REFLEX_CLOUD_PROVIDER";
73 const ENV_CLOUD_REGION: &'static str = "REFLEX_CLOUD_REGION";
74
75 pub fn from_env() -> LifecycleResult<Self> {
77 let defaults = Self::default();
78 let gcs_bucket = env::var(Self::ENV_GCS_BUCKET).unwrap_or_default();
79 let local_snapshot_path = env::var(Self::ENV_LOCAL_SNAPSHOT_PATH)
80 .map(PathBuf::from)
81 .unwrap_or(defaults.local_snapshot_path);
82 let idle_timeout = env::var(Self::ENV_IDLE_TIMEOUT_SECS)
83 .ok()
84 .and_then(|s| s.parse::<u64>().ok())
85 .map(Duration::from_secs)
86 .unwrap_or(defaults.idle_timeout);
87 let enable_instance_stop = env::var(Self::ENV_ENABLE_INSTANCE_STOP)
88 .map(|s| s != "false" && s != "0")
89 .unwrap_or(defaults.enable_instance_stop);
90
91 let cloud_provider = env::var(Self::ENV_CLOUD_PROVIDER)
92 .ok()
93 .and_then(|s| s.parse().ok())
94 .unwrap_or(defaults.cloud_provider);
95
96 let cloud_region = env::var(Self::ENV_CLOUD_REGION).ok();
97
98 Ok(Self {
99 gcs_bucket,
100 local_snapshot_path,
101 idle_timeout,
102 enable_instance_stop,
103 cloud_provider,
104 cloud_region,
105 })
106 }
107
108 pub fn has_gcs_bucket(&self) -> bool {
110 !self.gcs_bucket.is_empty()
111 }
112
113 #[cfg(test)]
114 pub fn for_testing(gcs_bucket: &str, local_path: PathBuf) -> Self {
116 Self {
117 gcs_bucket: gcs_bucket.to_string(),
118 local_snapshot_path: local_path,
119 idle_timeout: Duration::from_secs(1),
120 enable_instance_stop: false,
121 cloud_provider: CloudProviderType::Local,
122 cloud_region: None,
123 }
124 }
125}