rocketmq_model/utils/
env_utils.rs

1// Copyright 2023 The RocketMQ Rust Authors
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
15use std::ffi::OsStr;
16
17use crate::common::mix_all::ROCKETMQ_HOME_ENV;
18
19/// Utility functions related to environment variables.
20pub struct EnvUtils;
21
22impl EnvUtils {
23    /// Gets the value of the specified environment variable.
24    ///
25    /// # Arguments
26    ///
27    /// * `key` - The name of the environment variable to retrieve.
28    ///
29    /// # Returns
30    ///
31    /// An `Option` containing the value of the environment variable, or `None` if the variable is
32    /// not set.
33    pub fn get_property<K: AsRef<OsStr>>(key: K) -> Option<String> {
34        std::env::var(key).ok()
35    }
36
37    /// Retrieves the value of the specified environment variable or returns a default value if the
38    /// variable is not set.
39    ///
40    /// # Arguments
41    /// * `key` - The name of the environment variable to retrieve.
42    /// * `default` - The default value to return if the environment variable is not set.
43    ///
44    /// # Returns
45    /// A `String` containing the value of the environment variable, or the default value.
46    pub fn get_property_or_default<K: AsRef<OsStr>>(key: K, default: impl Into<String>) -> String {
47        std::env::var(key).unwrap_or_else(|_| default.into())
48    }
49
50    /// Retrieves the value of the specified environment variable as an `i32`, or returns a default
51    /// value if the variable is not set or cannot be parsed.
52    ///
53    /// # Arguments
54    /// * `key` - The name of the environment variable to retrieve.
55    /// * `default` - The default value to return if the environment variable is not set or cannot
56    ///   be parsed.
57    ///
58    /// # Returns
59    /// An `i32` containing the value of the environment variable, or the default value.
60    pub fn get_property_as_i32<K: AsRef<OsStr>>(key: K, default: i32) -> i32 {
61        std::env::var(key)
62            .ok()
63            .and_then(|v| v.parse::<i32>().ok())
64            .unwrap_or(default)
65    }
66
67    /// Retrieves the value of the specified environment variable as a `bool`, or returns a default
68    /// value if the variable is not set or cannot be parsed.
69    ///
70    /// # Arguments
71    /// * `key` - The name of the environment variable to retrieve.
72    /// * `default` - The default value to return if the environment variable is not set or cannot
73    ///   be parsed.
74    ///
75    /// # Returns
76    /// A `bool` containing the value of the environment variable, or the default value.
77    ///
78    /// # Notes
79    /// The function considers the following values as `true`: `"true"`, `"1"`.
80    /// The function considers the following values as `false`: `"false"`, `"0"`.
81    pub fn get_property_as_bool<K: AsRef<OsStr>>(key: K, default: bool) -> bool {
82        std::env::var(key)
83            .ok()
84            .and_then(|v| {
85                let lower = v.to_lowercase();
86                match lower.as_str() {
87                    "true" | "1" => Some(true),
88                    "false" | "0" => Some(false),
89                    _ => None,
90                }
91            })
92            .unwrap_or(default)
93    }
94
95    /// Gets the value of the ROCKETMQ_HOME environment variable.
96    ///
97    /// If ROCKETMQ_HOME is not set, it defaults to the current directory without
98    /// mutating the process environment.
99    ///
100    /// # Returns
101    ///
102    /// The value of the ROCKETMQ_HOME environment variable as a `String`.
103    pub fn get_rocketmq_home() -> String {
104        std::env::var(ROCKETMQ_HOME_ENV).unwrap_or_else(|_| {
105            std::env::current_dir()
106                .ok()
107                .and_then(|path| path.into_os_string().into_string().ok())
108                .unwrap_or_else(|| ".".to_string())
109        })
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use std::ffi::OsString;
116    use std::sync::Mutex;
117
118    use super::*;
119
120    static ENV_LOCK: Mutex<()> = Mutex::new(());
121
122    struct EnvironmentRestore {
123        key: &'static str,
124        original: Option<OsString>,
125    }
126
127    impl Drop for EnvironmentRestore {
128        fn drop(&mut self) {
129            // SAFETY: ENV_LOCK serializes every environment mutation in this
130            // module, and the guard restores the original value before release.
131            unsafe {
132                match &self.original {
133                    Some(value) => std::env::set_var(self.key, value),
134                    None => std::env::remove_var(self.key),
135                }
136            }
137        }
138    }
139
140    fn with_environment<R>(key: &'static str, value: Option<&str>, test: impl FnOnce() -> R) -> R {
141        let _lock = ENV_LOCK.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
142        let restore = EnvironmentRestore {
143            key,
144            original: std::env::var_os(key),
145        };
146        // SAFETY: ENV_LOCK prevents concurrent mutations made by this module;
147        // EnvironmentRestore restores the original process value after the test.
148        unsafe {
149            match value {
150                Some(value) => std::env::set_var(key, value),
151                None => std::env::remove_var(key),
152            }
153        }
154        let result = test();
155        drop(restore);
156        result
157    }
158
159    #[test]
160    fn test_get_property_existing_variable() {
161        with_environment("ROCKETMQ_ENV_UTILS_EXISTING", Some("expected"), || {
162            assert_eq!(
163                EnvUtils::get_property("ROCKETMQ_ENV_UTILS_EXISTING"),
164                Some("expected".to_string())
165            );
166        });
167    }
168
169    #[test]
170    fn test_get_property_non_existing_variable() {
171        // Set up
172        with_environment("ROCKETMQ_ENV_UTILS_MISSING", None, || {
173            assert_eq!(EnvUtils::get_property("ROCKETMQ_ENV_UTILS_MISSING"), None);
174        });
175    }
176
177    #[test]
178    fn rocketmq_home_fallback_does_not_write_the_environment() {
179        with_environment(ROCKETMQ_HOME_ENV, None, || {
180            let result = EnvUtils::get_rocketmq_home();
181
182            assert_eq!(result, std::env::current_dir().unwrap().to_string_lossy());
183            assert_eq!(std::env::var_os(ROCKETMQ_HOME_ENV), None);
184        });
185    }
186
187    #[test]
188    fn retrieves_env_variable_value() {
189        with_environment("ROCKETMQ_ENV_UTILS_STRING", Some("test_value"), || {
190            assert_eq!(
191                EnvUtils::get_property_or_default("ROCKETMQ_ENV_UTILS_STRING", "default_value"),
192                "test_value"
193            );
194        });
195    }
196
197    #[test]
198    fn returns_default_when_env_variable_not_set() {
199        assert_eq!(
200            EnvUtils::get_property_or_default("NON_EXISTENT_KEY", "default_value"),
201            "default_value"
202        );
203    }
204
205    #[test]
206    fn retrieves_env_variable_as_i32() {
207        with_environment("ROCKETMQ_ENV_UTILS_INT", Some("42"), || {
208            assert_eq!(EnvUtils::get_property_as_i32("ROCKETMQ_ENV_UTILS_INT", 0), 42);
209        });
210    }
211
212    #[test]
213    fn returns_default_when_env_variable_as_i32_not_set() {
214        assert_eq!(EnvUtils::get_property_as_i32("NON_EXISTENT_INT_KEY", 10), 10);
215    }
216
217    #[test]
218    fn returns_default_when_env_variable_as_i32_invalid() {
219        with_environment("ROCKETMQ_ENV_UTILS_INVALID_INT", Some("not_a_number"), || {
220            assert_eq!(EnvUtils::get_property_as_i32("ROCKETMQ_ENV_UTILS_INVALID_INT", 5), 5);
221        });
222    }
223
224    #[test]
225    fn returns_default_when_env_variable_as_bool_not_set() {
226        assert!(EnvUtils::get_property_as_bool("NON_EXISTENT_BOOL_KEY", true));
227        assert!(!EnvUtils::get_property_as_bool("NON_EXISTENT_BOOL_KEY", false));
228    }
229
230    #[test]
231    fn returns_default_when_env_variable_as_bool_invalid() {
232        with_environment("ROCKETMQ_ENV_UTILS_INVALID_BOOL", Some("not_a_bool"), || {
233            assert!(EnvUtils::get_property_as_bool("ROCKETMQ_ENV_UTILS_INVALID_BOOL", true));
234            assert!(!EnvUtils::get_property_as_bool(
235                "ROCKETMQ_ENV_UTILS_INVALID_BOOL",
236                false
237            ));
238        });
239    }
240}
rocketmq_model/utils/env_utils.rs

rocketmq_model/utils/
env_utils.rs
