From e0dead151e3c40054de509c1fdd9411285b962e0 Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Thu, 27 Jun 2024 15:57:43 +0200
Subject: [PATCH] feat: Store.weak_ref()

(cherry picked from commit 2fdcc5df62a6cea790bea9b867e1b6d044d4a28f)
---
 rust/nix-store/src/store.rs | 40 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/rust/nix-store/src/store.rs b/rust/nix-store/src/store.rs
index d6cca3d..8ca2031 100644
--- a/rust/nix-store/src/store.rs
+++ b/rust/nix-store/src/store.rs
@@ -7,7 +7,7 @@ use nix_util::{check_call, result_string_init};
 use std::ffi::{c_char, CString};
 use std::ptr::null_mut;
 use std::ptr::NonNull;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 /* TODO make Nix itself thread safe */
 lazy_static! {
@@ -33,6 +33,20 @@ impl Drop for StoreRef {
     }
 }
 
+/// A [Weak] reference to a store.
+pub struct StoreWeak {
+    inner: Weak<StoreRef>,
+}
+impl StoreWeak {
+    /// Upgrade the weak reference to a proper [Store].
+    pub fn upgrade(&self) -> Option<Store> {
+        self.inner.upgrade().map(|inner| Store {
+            inner,
+            context: Context::new(),
+        })
+    }
+}
+
 pub struct Store {
     inner: Arc<StoreRef>,
     /* An error context to reuse. This way we don't have to allocate them for each store operation. */
@@ -109,6 +123,12 @@ impl Store {
         }?;
         r
     }
+
+    pub fn weak_ref(&self) -> StoreWeak {
+        StoreWeak {
+            inner: Arc::downgrade(&self.inner),
+        }
+    }
 }
 
 impl Clone for Store {
@@ -156,4 +176,22 @@ mod tests {
         let uri = store.get_uri().unwrap();
         assert_eq!(uri, "https://cache.nixos.org");
     }
+
+    #[test]
+    fn weak_ref() {
+        let mut store = Store::open("auto", HashMap::new()).unwrap();
+        let uri = store.get_uri().unwrap();
+        let weak = store.weak_ref();
+        let mut store2 = weak.upgrade().unwrap();
+        assert_eq!(store2.get_uri().unwrap(), uri);
+    }
+    #[test]
+    fn weak_ref_gone() {
+        let weak = {
+            let store = Store::open("auto", HashMap::new()).unwrap();
+            store.weak_ref()
+        };
+        assert!(weak.upgrade().is_none());
+        assert!(weak.inner.upgrade().is_none());
+    }
 }