system/net: introduce LazyWeak which simplifies parent-child hierarchies

src/net/p2p.rs

@@ -71,11 +71,11 @@ pub struct P2p {
     pub peer_discovery_running: Mutex<bool>,
 
     /// Reference to configured [`ManualSession`]
-    session_manual: Mutex<Option<Arc<ManualSession>>>,
+    session_manual: ManualSessionPtr,
     /// Reference to configured [`InboundSession`]
-    session_inbound: Mutex<Option<Arc<InboundSession>>>,
+    session_inbound: InboundSessionPtr,
     /// Reference to configured [`OutboundSession`]
-    session_outbound: Mutex<Option<Arc<OutboundSession>>>,
+    session_outbound: OutboundSessionPtr,
 
     /// Enable network debugging
     pub dnet_enabled: Mutex<bool>,
@@ -105,19 +105,17 @@ impl P2p {
             settings,
             peer_discovery_running: Mutex::new(false),
 
-            session_manual: Mutex::new(None),
-            session_inbound: Mutex::new(None),
-            session_outbound: Mutex::new(None),
+            session_manual: ManualSession::new(),
+            session_inbound: InboundSession::new(),
+            session_outbound: OutboundSession::new(),
 
             dnet_enabled: Mutex::new(false),
             dnet_subscriber: Subscriber::new(),
         });
 
-        let parent = Arc::downgrade(&self_);
-
-        *self_.session_manual.lock().await = Some(ManualSession::new(parent.clone()));
-        *self_.session_inbound.lock().await = Some(InboundSession::new(parent.clone()));
-        *self_.session_outbound.lock().await = Some(OutboundSession::new(parent).await);
+        self_.session_manual.p2p.init(self_.clone());
+        self_.session_inbound.p2p.init(self_.clone());
+        self_.session_outbound.p2p.init(self_.clone());
 
         register_default_protocols(self_.clone()).await;
 
@@ -274,17 +272,17 @@ impl P2p {
 
     /// Get pointer to manual session
     pub async fn session_manual(&self) -> ManualSessionPtr {
-        self.session_manual.lock().await.as_ref().unwrap().clone()
+        self.session_manual.clone()
     }
 
     /// Get pointer to inbound session
     pub async fn session_inbound(&self) -> InboundSessionPtr {
-        self.session_inbound.lock().await.as_ref().unwrap().clone()
+        self.session_inbound.clone()
     }
 
     /// Get pointer to outbound session
     pub async fn session_outbound(&self) -> OutboundSessionPtr {
-        self.session_outbound.lock().await.as_ref().unwrap().clone()
+        self.session_outbound.clone()
     }
 
     /// Enable network debugging

src/net/session/inbound_session.rs

@@ -23,7 +23,7 @@
 //! an acceptor pointer, and a stoppable task pointer. Using a weak pointer
 //! to P2P allows us to avoid circular dependencies.
 
-use std::sync::{Arc, Weak};
+use std::sync::Arc;
 
 use async_trait::async_trait;
 use log::{debug, error, info};
@@ -39,7 +39,7 @@ use super::{
     Session, SessionBitFlag, SESSION_INBOUND,
 };
 use crate::{
-    system::{StoppableTask, StoppableTaskPtr},
+    system::{LazyWeak, StoppableTask, StoppableTaskPtr},
     Error, Result,
 };
 
@@ -47,15 +47,19 @@ pub type InboundSessionPtr = Arc<InboundSession>;
 
 /// Defines inbound connections session
 pub struct InboundSession {
-    p2p: Weak<P2p>,
+    pub(in crate::net) p2p: LazyWeak<P2p>,
     acceptors: Mutex<Vec<AcceptorPtr>>,
     accept_tasks: Mutex<Vec<StoppableTaskPtr>>,
 }
 
 impl InboundSession {
     /// Create a new inbound session
-    pub fn new(p2p: Weak<P2p>) -> InboundSessionPtr {
-        Arc::new(Self { p2p, acceptors: Mutex::new(vec![]), accept_tasks: Mutex::new(vec![]) })
+    pub fn new() -> InboundSessionPtr {
+        Arc::new(Self {
+            p2p: LazyWeak::new(),
+            acceptors: Mutex::new(Vec::new()),
+            accept_tasks: Mutex::new(Vec::new()),
+        })
     }
 
     /// Starts the inbound session. Begins by accepting connections and fails
@@ -174,7 +178,7 @@ impl InboundSession {
 #[async_trait]
 impl Session for InboundSession {
     fn p2p(&self) -> P2pPtr {
-        self.p2p.upgrade().unwrap()
+        self.p2p.upgrade()
     }
 
     fn type_id(&self) -> SessionBitFlag {

src/net/session/manual_session.rs

@@ -29,7 +29,7 @@
 //! and insures that no other part of the program uses the slots at the
 //! same time.
 
-use std::sync::{Arc, Weak};
+use std::sync::Arc;
 
 use async_trait::async_trait;
 use log::{info, warn};
@@ -45,7 +45,7 @@ use super::{
     Session, SessionBitFlag, SESSION_MANUAL,
 };
 use crate::{
-    system::{sleep, StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr},
+    system::{sleep, LazyWeak, StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr},
     Error, Result,
 };
 
@@ -53,7 +53,7 @@ pub type ManualSessionPtr = Arc<ManualSession>;
 
 /// Defines manual connections session.
 pub struct ManualSession {
-    p2p: Weak<P2p>,
+    pub(in crate::net) p2p: LazyWeak<P2p>,
     connect_slots: Mutex<Vec<StoppableTaskPtr>>,
     /// Subscriber used to signal channels processing
     channel_subscriber: SubscriberPtr<Result<ChannelPtr>>,
@@ -63,9 +63,9 @@ pub struct ManualSession {
 
 impl ManualSession {
     /// Create a new manual session.
-    pub fn new(p2p: Weak<P2p>) -> ManualSessionPtr {
+    pub fn new() -> ManualSessionPtr {
         Arc::new(Self {
-            p2p,
+            p2p: LazyWeak::new(),
             connect_slots: Mutex::new(Vec::new()),
             channel_subscriber: Subscriber::new(),
             notify: Mutex::new(false),
@@ -204,7 +204,7 @@ impl ManualSession {
 #[async_trait]
 impl Session for ManualSession {
     fn p2p(&self) -> P2pPtr {
-        self.p2p.upgrade().unwrap()
+        self.p2p.upgrade()
     }
 
     fn type_id(&self) -> SessionBitFlag {

src/net/session/outbound_session.rs

@@ -45,38 +45,18 @@ use super::{
     Session, SessionBitFlag, SESSION_OUTBOUND,
 };
 use crate::{
-    system::{sleep, CondVar, StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr},
+    system::{
+        sleep, CondVar, LazyWeak, StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr,
+    },
     Error, Result,
 };
 
-use std::sync::OnceLock;
-
-pub struct LazyWeak<Parent>(OnceLock<Weak<Parent>>);
-
-impl<Parent> LazyWeak<Parent> {
-    fn new() -> Self {
-        Self(OnceLock::new())
-    }
-
-    pub fn init(&self, parent: Arc<Parent>) {
-        assert!(self.0.get().is_none());
-        let parent = Arc::downgrade(&parent);
-        self.0.set(parent).unwrap();
-        assert!(self.0.get().is_some());
-    }
-
-    pub fn upgrade(&self) -> Arc<Parent> {
-        assert!(self.0.get().is_some());
-        self.0.get().unwrap().upgrade().unwrap()
-    }
-}
-
 pub type OutboundSessionPtr = Arc<OutboundSession>;
 
 /// Defines outbound connections session.
 pub struct OutboundSession {
     /// Weak pointer to parent p2p object
-    p2p: Weak<P2p>,
+    pub(in crate::net) p2p: LazyWeak<P2p>,
     /// Subscriber used to signal channels processing
     channel_subscriber: SubscriberPtr<Result<ChannelPtr>>,
 
@@ -88,9 +68,9 @@ pub struct OutboundSession {
 
 impl OutboundSession {
     /// Create a new outbound session.
-    pub(crate) async fn new(p2p: Weak<P2p>) -> OutboundSessionPtr {
+    pub(crate) fn new() -> OutboundSessionPtr {
         let self_ = Arc::new(Self {
-            p2p,
+            p2p: LazyWeak::new(),
             channel_subscriber: Subscriber::new(),
             slots: Mutex::new(Vec::new()),
             peer_discovery: PeerDiscovery::new(),
@@ -142,7 +122,7 @@ impl OutboundSession {
 #[async_trait]
 impl Session for OutboundSession {
     fn p2p(&self) -> P2pPtr {
-        self.p2p.upgrade().unwrap()
+        self.p2p.upgrade()
     }
 
     fn type_id(&self) -> SessionBitFlag {

src/system/lazy_weak.rs

@@ -0,0 +1,77 @@
+use std::sync::{Arc, OnceLock, Weak};
+
+/// Sometimes you need a parent-child relationship which results in code like:
+/// ```rust
+/// struct Parent {
+///     child: Mutex<Option<Arc<Child>>>
+/// }
+/// impl Parent {
+///     fn new() -> Arc<Self> {
+///         let self_ = Arc::new(Self {
+///             child: Mutex::new(None)
+///         };
+///         let parent = Arc::downgrade(&self_);
+///         *self_.child.lock().await = Some(Child::new(parent));
+///         self_
+///     }
+/// }
+/// struct Child {
+///     parent: Weak<Parent>
+/// }
+/// impl Child {
+///     fn new(parent: Weak<Parent>) -> Self {
+///         Self { parent }
+///     }
+///     fn upgrade(&self) -> Arc<Parent> {
+///         self.parent.upgrade().unwrap()
+///     }
+/// }
+/// ```
+/// This class simplifies the above code by allowing us instead to do:
+/// ```rust
+/// struct Parent {
+///     child: Arc<Child>
+/// }
+/// impl Parent {
+///     fn new() -> Arc<Self> {
+///         let self_ = Arc::new(Self {
+///             child: Child::new()
+///         };
+///         self_.child.parent.init(self_.clone());
+///         self_
+///     }
+/// }
+/// struct Child {
+///     parent: LazyWeak<Parent>
+/// }
+/// impl Child {
+///     fn new() -> Self {
+///         Self { parent: LazyWeak::new() }
+///     }
+///     fn upgrade(&self) -> Arc<Parent> {
+///         self.parent.upgrade()
+///     }
+/// }
+/// ```
+pub struct LazyWeak<Parent>(OnceLock<Weak<Parent>>);
+
+impl<Parent> LazyWeak<Parent> {
+    /// Create an empty `LazyWeak`, which must immediately be followed by `weak.init()`.
+    pub fn new() -> Self {
+        Self(OnceLock::new())
+    }
+
+    /// Must be called within the same scope as `new()`.
+    pub fn init(&self, parent: Arc<Parent>) {
+        assert!(self.0.get().is_none());
+        let parent = Arc::downgrade(&parent);
+        self.0.set(parent).unwrap();
+        assert!(self.0.get().is_some());
+    }
+
+    /// Access the `Arc<Parent>` pointer
+    pub fn upgrade(&self) -> Arc<Parent> {
+        assert!(self.0.get().is_some());
+        self.0.get().unwrap().upgrade().unwrap()
+    }
+}

src/system/mod.rs

@@ -24,6 +24,10 @@ use smol::{Executor, Timer};
 pub mod condvar;
 pub use condvar::CondVar;
 
+/// Convenient late initialization of `Weak<Foo>`
+pub mod lazy_weak;
+pub use lazy_weak::LazyWeak;
+
 /// Implementation of async background task spawning which are stoppable
 /// using channel signalling.
 pub mod stoppable_task;