fix: better memory management

config/config/example/derive.rs

@@ -1,7 +1,7 @@
 //! Run with: `cargo run --example derive`
 //! Look at the generated code with: `cargo expand --example derive`
 
-use config::{sources, ConfigBuilder, ConfigError};
+use config::{sources, Config, ConfigBuilder, ConfigError};
 
 type TypeAlias = bool;
 
@@ -17,16 +17,37 @@ struct AppConfig {
 struct LoggingConfig {
     level: String,
     json: bool,
-    #[config(cli(skip), env())]
-    cycle: Option<Box<LoggingConfig>>,
+    manual: Manual,
+    manual2: Manual,
+}
+
+#[derive(Default, Debug, PartialEq, serde::Deserialize)]
+struct Manual {
+    cycle: Box<Option<LoggingConfig>>,
+}
+
+impl config::Config for Manual {
+    fn graph() -> &'static config::KeyGraph {
+        if let Some(tree) = config::KeyGraph::get::<Self>() {
+            return tree;
+        }
+
+        let mut keys = std::collections::BTreeMap::new();
+
+        keys.insert("level".to_string(), config::Key::new(String::graph()));
+        keys.insert("json".to_string(), config::Key::new(bool::graph()));
+
+        config::KeyGraph::store::<Self>(config::KeyGraph::Struct(keys))
+    }
 }
 
 impl Default for LoggingConfig {
     fn default() -> Self {
         Self {
             level: "INFO".to_string(),
             json: false,
-            cycle: None,
+            manual: Manual::default(),
+            manual2: Manual::default(),
         }
     }
 }
@@ -39,21 +60,42 @@ fn main() {
 }
 
 fn parse() -> Result<AppConfig, ConfigError> {
+    let graph = std::thread::spawn(AppConfig::graph).join().unwrap();
+
+    println!("{:#?}", graph);
+
     let mut builder = ConfigBuilder::new();
     builder.add_source(sources::CliSource::new()?);
     builder.add_source(sources::EnvSource::with_prefix("TEST")?);
     builder.add_source(sources::FileSource::json(
         br#"
     {
         "enabled": "on",
-        "optional": null,
-        "map": {
-            "key": {
-                "level": "DEBUG",
-                "json": true,
+        "logging": {
+            "level": "DEBUG",
+            "json": true,
+            "manual": {
+                "cycle": {
+                    "level": "INFO",
+                    "json": false,
+                    "manual": {
+                        "cycle": null
+                    },
+                    "manual2": {
+                        "cycle": null
+                    }
+                }
+            },
+            "manual2": {
                 "cycle": {
-                    "level": "TRACE",
-                    "json": false
+                    "level": "INFO",
+                    "json": false,
+                    "manual": {
+                        "cycle": null
+                    },
+                    "manual2": {
+                        "cycle": null
+                    }
                 }
             }
         }

config/config/src/key.rs

@@ -1,8 +1,8 @@
 use std::{
-    cell::RefCell,
     collections::{btree_map::Entry, BTreeMap},
     fmt::Display,
     ptr::NonNull,
+    sync::Mutex,
 };
 
 #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
@@ -98,56 +98,75 @@ pub enum KeyType {
     CyclicReference,
 }
 
-thread_local! {
-    static MEMO_TREES: RefCell<BTreeMap<std::any::TypeId, KeyGraphRefHelper>> = Default::default();
-}
+static MEMO_GRAPHS: Mutex<BTreeMap<std::any::TypeId, KeyGraphRefHelper>> =
+    Mutex::new(BTreeMap::new());
 
 struct KeyGraphRefHelper {
     ptr: NonNull<KeyGraph>,
     ref_key: &'static KeyGraph,
     is_building: bool,
 }
 
+unsafe impl Send for KeyGraphRefHelper {}
+unsafe impl Sync for KeyGraphRefHelper {}
+
 impl Drop for KeyGraphRefHelper {
     fn drop(&mut self) {
+        // Safety: We are the only ones with access to the pointer.
+        // We are dropping the pointer and therefore we can drop the value.
+        // In practice this should never actually be called because the `MEMO_GRAPHS` is a static
+        // and therefore will never be dropped.
         unsafe {
             let _ = Box::from_raw(self.ptr.as_ptr());
-            let _ = Box::from_raw(self.ref_key as *const _ as *mut KeyGraph);
+            let _ = Box::from_raw(self.ref_key as *const KeyGraph as *mut KeyGraph);
         }
     }
 }
 
 impl KeyGraphRefHelper {
     fn new() -> Self {
+        // We create a pointer to some dummy value.
         let ptr = NonNull::new(Box::leak(Box::new(KeyGraph::Unit))).unwrap();
+
         Self {
+            // We then create a pointer to a ref to the dummy value.
+            // This is because the reference key is a cyclic reference.
+            // And we if we are calling the get function from the same path that is building the
+            // tree then we need to return the ref key. Otherwise we would run into an infinite
+            // loop and stack overflow.
+            // Safety: We created a non null pointer above and its static.
             ref_key: Box::leak(Box::new(KeyGraph::Ref(unsafe { ptr.as_ref() }))),
             is_building: true,
             ptr,
         }
     }
 
     fn get(&self) -> &'static KeyGraph {
-        // This is safe because this is always a valid pointer
+        // The reason we return the ref key here is because if the get function is being
+        // called by the same path that is building the tree then we need to return the
+        // ref key because this is a cyclic reference.
         if self.is_building {
             return self.ref_key;
         }
 
+        // If we are not building then we can return the actual pointer, ie: memoization.
+        // Safety: We are already built and therefore the pointer is valid. It is also static.
         unsafe { self.ptr.as_ref() }
     }
 
-    fn set(&mut self, tree: KeyGraph) {
-        if !self.is_building {
-            panic!("KeyTree already built");
-        }
-
-        // This is safe because we are thread_local! and therefore no other thread can have mutable access to this pointer.
-        // We are moving the content of tree to the pointer. Which will then make every reference to this pointer point to the new content.
-        // The old content will never be destructed but thats fine because the new struct is the exact same size and alignment.
+    fn set(&mut self, mut tree: KeyGraph) {
+        // Safety: We have a mutable ref to `KeyGraphRefHelper` and therefore we are the only
+        // ones with access to the pointer.
+        // We swap the old value of the pointer with the new value. We can then drop the old
+        // value of the pointer. (which will be in `tree`)
+        // It should be impossible to call this function twice on the same type.
+        // However even if it is called twice it is still safe because we are swapping the
+        // pointer and therefore the old value will be dropped and no memory leaks will occur.
         unsafe {
-            *self.ptr.as_mut() = tree;
+            std::mem::swap(self.ptr.as_mut(), &mut tree);
         }
 
+        // Since we are no longer building and are built we stop returning the ref key.
         self.is_building = false;
     }
 }
@@ -231,30 +250,25 @@ impl KeyGraph {
     }
 
     pub fn get<T: crate::Config>() -> Option<&'static Self> {
-        MEMO_TREES.with(|memo_tree| {
-            let mut memo_tree = memo_tree.borrow_mut();
-            let id = std::any::TypeId::of::<T>();
-
-            if let Entry::Vacant(e) = memo_tree.entry(id) {
-                e.insert(KeyGraphRefHelper::new());
-                None
-            } else {
-                memo_tree.get(&id).map(|h| h.get())
-            }
-        })
+        let mut memo_graphs = MEMO_GRAPHS.lock().expect("Failed to lock memo tree");
+        let id = std::any::TypeId::of::<T>();
+
+        if let Entry::Vacant(e) = memo_graphs.entry(id) {
+            e.insert(KeyGraphRefHelper::new());
+            None
+        } else {
+            memo_graphs.get(&id).map(|h| h.get())
+        }
     }
 
     pub fn store<T: crate::Config>(tree: Self) -> &'static Self {
-        MEMO_TREES.with(|memo_tree| {
-            let mut memo_tree = memo_tree.borrow_mut();
-
-            let id = std::any::TypeId::of::<T>();
-            let helper = memo_tree.get_mut(&id).unwrap();
+        let mut memo_graphs = MEMO_GRAPHS.lock().expect("Failed to lock memo tree");
 
-            helper.set(tree);
+        let id = std::any::TypeId::of::<T>();
+        let helper = memo_graphs.get_mut(&id).unwrap();
 
-            helper.get()
-        })
+        helper.set(tree);
+        helper.get()
     }
 }
 

config/config/src/lib.rs

@@ -262,8 +262,6 @@ pub trait Source<C: Config> {
     fn get_key(&self, path: &KeyPath) -> Result<Option<Value>>;
 }
 
-/// You don't need to implemented this trait manually.
-/// Please use the provided derive macro to generate the implementation.
 pub trait Config: serde::de::DeserializeOwned + 'static {
     const PKG_NAME: Option<&'static str> = None;
     const ABOUT: Option<&'static str> = None;