From f2b72f59b432758123ffc9b72ff386c5cd8f038b Mon Sep 17 00:00:00 2001 From: Torben Schweren Date: Fri, 4 Oct 2024 19:02:24 +0200 Subject: [PATCH] refactor: use Weak instead of Arc self-references --- src/event/event_repeater.rs | 54 +++++++++++++++++++++++++--------- src/service/service_manager.rs | 39 +++++++++++++++++------- 2 files changed, 68 insertions(+), 25 deletions(-) diff --git a/src/event/event_repeater.rs b/src/event/event_repeater.rs index d585df7..7d3a06a 100644 --- a/src/event/event_repeater.rs +++ b/src/event/event_repeater.rs @@ -1,20 +1,26 @@ -use std::{collections::HashMap, sync::Arc}; +use log::error; +use std::{ + collections::HashMap, + sync::{Arc, Weak}, +}; use thiserror::Error; use tokio::{sync::Mutex, task::JoinHandle}; use uuid::Uuid; +use crate::setlock::{SetLock, SetLockError}; + use super::{Event, Subscription}; #[derive(Debug, Error)] pub enum AttachError { - #[error("Tried to attach event {event_name} to EventRepeater {repeater_name} before it was initialized. Did you not use EventRepeater::new()?")] + #[error("Tried to attach event {event_name} to EventRepeater {repeater_name} while it was uninitialized. Did you not use EventRepeater::new()?")] NotInitialized { event_name: String, repeater_name: String, }, #[error( - "Tried to attach event {event_name} to EventRepeater {repeater_name}, which was already attached." + "Tried to attach event {event_name} to EventRepeater {repeater_name}, which was already attached to it." )] AlreadyAttached { event_name: String, @@ -47,7 +53,7 @@ where T: Send + Sync + 'static, { pub event: Event, - self_arc: Mutex>>, + weak: Mutex>>, subscriptions: Mutex)>>, } @@ -62,18 +68,28 @@ where { let event = Event::new(name); let event_repeater = Self { - self_arc: Mutex::new(None), + weak: Mutex::new(SetLock::new()), event, subscriptions: Mutex::new(HashMap::new()), }; - let self_arc = Arc::new(event_repeater); - let mut lock = self_arc.self_arc.lock().await; - let self_arc_clone = Arc::clone(&self_arc); - *lock = Some(self_arc_clone); - drop(lock); + let arc = Arc::new(event_repeater); + let weak = Arc::downgrade(&arc); + + let result = arc.weak.lock().await.set(weak); + if let Err(err) = result { + match err { + SetLockError::AlreadySet => { + error!("Failed to set EventRepeater {}'s Weak self-reference because it was already set. This should never happen. Shutting down ungracefully to prevent further undefined behavior.", arc.event.name); + unreachable!( + "Unable to set EventRepeater {}'s Weak self-reference because it was already set.", + arc.event.name + ); + } + } + } - self_arc + arc } pub async fn subscription_count(&self) -> usize { @@ -81,8 +97,9 @@ where } pub async fn attach(&self, event: &Event, buffer: usize) -> Result<(), AttachError> { - let self_arc = match self.self_arc.lock().await.as_ref() { - Some(arc) => Arc::clone(arc), + let lock = self.weak.lock().await; + let weak = match lock.get() { + Some(weak) => weak, None => { return Err(AttachError::NotInitialized { event_name: event.name.clone(), @@ -91,6 +108,15 @@ where } }; + // This can't fail because the Arc is guaranteed to be valid as long as &self is valid. + let arc = match weak.upgrade() { + Some(arc) => arc, + None => { + error!("EventRepeater {}'s Weak self-reference could not be upgraded to Arc while attaching event {}. This should never happen. Shutting down ungracefully to prevent further undefined behavior.", self.event.name, event.name); + unreachable!("EventRepeater {}'s Weak self-reference could not be upgraded to Arc while attaching event {}.", self.event.name, event.name); + } + }; + let mut subscriptions = self.subscriptions.lock().await; if subscriptions.contains_key(&event.uuid) { return Err(AttachError::AlreadyAttached { @@ -108,7 +134,7 @@ where let join_handle = tokio::spawn(async move { while let Some(value) = receiver.recv().await { - let _ = self_arc.event.dispatch(value).await; + let _ = arc.event.dispatch(value).await; } }); subscriptions.insert(event.uuid, (subscription, join_handle)); diff --git a/src/service/service_manager.rs b/src/service/service_manager.rs index 6d25e44..78f4ecb 100644 --- a/src/service/service_manager.rs +++ b/src/service/service_manager.rs @@ -6,7 +6,7 @@ use crate::{ event::EventRepeater, service::Watchdog, setlock::{SetLock, SetLockError} }; use log::{error, info, warn}; -use std::{collections::HashMap, fmt::Display, mem, sync::Arc, time::Duration}; +use std::{collections::HashMap, fmt::Display, mem, sync::{Arc, Weak}, time::Duration}; use tokio::{ spawn, sync::{Mutex, MutexGuard}, @@ -54,31 +54,31 @@ pub fn new() -> Self { pub async fn build(self) -> Arc { let service_manager = ServiceManager { - arc: Mutex::new(SetLock::new()), + weak: Mutex::new(SetLock::new()), services: self.services, background_tasks: Mutex::new(HashMap::new()), on_status_change: EventRepeater::new("service_manager_on_status_change").await, }; - let self_arc = Arc::new(service_manager); - let self_arc_clone = Arc::clone(&self_arc); - - let result = self_arc_clone.arc.lock().await.set(Arc::clone(&self_arc_clone)); + let arc = Arc::new(service_manager); + let weak = Arc::downgrade(&arc); + let result = arc.weak.lock().await.set(weak); if let Err(err) = result { match err { SetLockError::AlreadySet => { - unreachable!("Unable to set ServiceManager's self-arc in ServiceManagerBuilder because it was already set. This should never happen. How did you...?"); + error!("Unable to set ServiceManager's Weak self-reference in ServiceManagerBuilder because it was already set. This should never happen. Shutting down ungracefully to prevent further undefined behavior."); + unreachable!("Unable to set ServiceManager's Weak self-reference in ServiceManagerBuilder because it was already set."); } } } - self_arc + arc } } pub struct ServiceManager { - arc: Mutex>>, + weak: Mutex>>, background_tasks: Mutex>>, pub services: Vec>>, @@ -326,10 +326,27 @@ impl ServiceManager { &self, service: &mut MutexGuard<'_, dyn Service>, ) -> Result<(), StartupError> { - let service_manager = Arc::clone(self.arc.lock().await.unwrap()); + let lock = self.weak.lock().await; + let weak = match lock.get() { + Some(weak) => weak, + None => { + error!("ServiceManager's Weak self-reference was None while initializing service {}. This should never happen. Did you not use a ServiceManagerBuilder? Shutting down ungracefully to prevent further undefined behavior.", service.info().name); + unreachable!("ServiceManager's Weak self-reference was None while initializing service {}.", service.info().name); + } + }; + + // This can't fail because the Arc is guaranteed to be valid as long as &self is valid. + let arc = match weak.upgrade() { + Some(arc) => arc, + None => { + error!("ServiceManager's Weak self-reference could not be upgraded to Arc while initializing service {}. This should never happen. Shutting down ungracefully to prevent further undefined behavior.", service.info().name); + unreachable!("ServiceManager's Weak self-reference could not be upgraded to Arc while initializing service {}.", service.info().name); + } + }; + //TODO: Add to config instead of hardcoding duration - let start = service.start(service_manager); + let start = service.start(arc); let timeout_result = timeout(Duration::from_secs(10), start).await; match timeout_result {