Is there a way to run a closure on RwLock drop? - rust

Question

I have a program that hides mutable state behind a RwLock. What I'd like to do is that when it's borrowed mutably (RW_LOCK.write()), on drop it should do something (namely, try to write to file, clean up the data behind the rwlock, etc.)
For example:
let DATA: RwLock<Data> = RwLock::new(Data { content: Default::default() } );
fn do_something() {
let mut state = DATA.write().unwrap();
state.change(5);
// ...
// Here, just before `state` goes out of scope (where it gets dropped and `RwLock` will allow
// other threads read/write access to `Data`, I would like for `RwLock` to auto-run `state.cleanup()`.
}
Is there a way to do this, or do I have to reimplement RwLock?

Answer 1

You can do this with a wrapper type:
Playground
use std::ops::{Deref, DerefMut, Drop};
use std::sync::{RwLock, RwLockWriteGuard};
type CleanupClosure<'a> = Fn(&mut RwLockWriteGuard<'a, Data>);
struct Data {
content: String,
}
impl Data {
fn change(&mut self, num: i32) {
println!("Changed to {}", num);
self.content = num.to_string();
}
}
struct RwLockWriteWrapper<'a, F: CleanupClosure<'a>>(RwLockWriteGuard<'a, Data>, F);
impl<'a, F: CleanupClosure<'a>> Deref for RwLockWriteWrapper<'a, F> {
type Target = RwLockWriteGuard<'a, Data>;
fn deref(&self) -> &RwLockWriteGuard<'a, Data> {
&self.0
}
}
impl<'a, F: CleanupClosure<'a>> DerefMut for RwLockWriteWrapper<'a, F> {
fn deref_mut(&mut self) -> &mut RwLockWriteGuard<'a, Data> {
&mut self.0
}
}
impl<'a, F: CleanupClosure<'a>> Drop for RwLockWriteWrapper<'a, F> {
fn drop(&mut self) {
println!("Cleaning up!");
self.1(&mut self.0)
}
}
fn main() {
let data: RwLock<Data> = RwLock::new(Data {
content: "Start".to_owned(),
});
do_something(&data);
do_something(&data);
}
fn do_something(data: &RwLock<Data>) {
// Write your own cleanup logic here
let mut state = RwLockWriteWrapper(data.write().unwrap(), |state| {
state.content = "Cleaned up".to_owned()
});
println!("do_something start: {}", state.content);
state.change(5);
println!("do_something after change: {}", state.content);
} // Automatically run cleanup here
It does require you to remember to wrap the type when you call .write() on it. You can wrap RwLock itself in another type that would return RwLockWriteWrapper to automate that as well.
This does become quite verbose, so I found a crate that impls the deref trait for you.
I am still not sure what closures you mentioned in the title meant.

Answer 2

You can create a wrapper that implements Drop:
struct CleanOnDrop<D, F>
where
F: FnOnce(&mut D),
{
data: D,
cleanup: Option<F>,
}
impl<D, F> CleanOnDrop<D, F>
where
F: FnOnce(&mut D),
{
pub fn new(data: D, cleanup: F) -> Self {
Self { data, cleanup: Some(cleanup) }
}
}
impl<D, F> Drop for CleanOnDrop<D, F>
where
F: FnOnce(&mut D)
{
fn drop(&mut self) {
if let Some(mut cleanup) = self.cleanup.take() {
cleanup(&mut self.data);
}
}
}
For convenience, you may want to implement Deref and DerefMut too, so that you can call methods on it directly:
use std::ops::{Deref, DerefMut};
impl<D, F> Deref for CleanOnDrop<D, F>
where
F: FnOnce(&mut D),
{
type Target = D;
fn deref(&self) -> &D {
&self.data
}
}
impl<D, F> DerefMut for CleanOnDrop<D, F>
where
F: FnOnce(&mut D),
{
fn deref_mut(&mut self) -> &mut D {
&mut self.data
}
}
Use the wrapper like this:
let data = RwLock::new(CleanOnDrop::new(
Data {
content: Default::default(),
},
|state| {
state.cleanup();
},
));