From 9b40b77395f6f704dda0710cd8ffa1079a449630 Mon Sep 17 00:00:00 2001
From: Lyn Nagara <lyn.nagara@gmail.com>
Date: Thu, 22 Feb 2024 14:12:50 -0800
Subject: [PATCH] feat(rust): Add ability to use runtime config from rust
 consumer (#5539)

---
 rust_snuba/src/runtime_config.rs | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/rust_snuba/src/runtime_config.rs b/rust_snuba/src/runtime_config.rs
index 5a314d0cdd..1efc51853a 100644
--- a/rust_snuba/src/runtime_config.rs
+++ b/rust_snuba/src/runtime_config.rs
@@ -1,8 +1,26 @@
 use anyhow::Error;
+use parking_lot::RwLock;
 use pyo3::prelude::{PyModule, Python};
+use std::collections::BTreeMap;
+use std::time::Duration;
 
+use rust_arroyo::utils::timing::Deadline;
+
+#[allow(dead_code)]
+static CONFIG: RwLock<BTreeMap<String, (Option<String>, Deadline)>> = RwLock::new(BTreeMap::new());
+
+/// Runtime config is cached for 10 seconds
 #[allow(dead_code)]
 pub fn get_str_config(key: &str) -> Result<Option<String>, Error> {
+    let deadline = Deadline::new(Duration::from_secs(10));
+
+    if let Some(value) = CONFIG.read().get(key) {
+        let (config, deadline) = value;
+        if !deadline.has_elapsed() {
+            return Ok(config.clone());
+        }
+    }
+
     Python::with_gil(|py| {
         let snuba_state = PyModule::import(py, "snuba.state")?;
         let config = snuba_state
@@ -10,7 +28,10 @@ pub fn get_str_config(key: &str) -> Result<Option<String>, Error> {
             .call1((key,))?
             .extract::<Option<String>>()?;
 
-        Ok(config)
+        CONFIG
+            .write()
+            .insert(key.to_string(), (config.clone(), deadline));
+        Ok(CONFIG.read().get(key).unwrap().0.clone())
     })
 }