f: revert to atomic counter again

Author: joostjager

lightning-persister/src/fs_store.rs

@@ -8,7 +8,7 @@ use std::collections::HashMap;
 use std::fs;
 use std::io::{Read, Write};
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
 
 #[cfg(feature = "tokio")]
@@ -41,45 +41,55 @@ fn path_to_windows_str<T: AsRef<OsStr>>(path: &T) -> Vec<u16> {
 // a consistent view and error out.
 const LIST_DIR_CONSISTENCY_RETRIES: usize = 10;
 
-#[derive(Default)]
-struct AsyncState {
-	// Version counter to ensure that writes are applied in the correct order. It is assumed that read, list and remove
-	// operations aren't sensitive to the order of execution.
-	latest_version: u64,
-
-	// The last version that was written to disk.
-	latest_written_version: u64,
-}
-
 struct FilesystemStoreInner {
 	data_dir: PathBuf,
 	tmp_file_counter: AtomicUsize,
 
 	// Per path lock that ensures that we don't have concurrent writes to the same file. The lock also encapsulates the
 	// latest written version per key.
-	locks: Mutex<HashMap<PathBuf, Arc<RwLock<AsyncState>>>>,
+	locks: Mutex<HashMap<PathBuf, Arc<RwLock<u64>>>>,
 }
 
 /// A [`KVStore`] and [`KVStoreSync`] implementation that writes to and reads from the file system.
 ///
 /// [`KVStore`]: lightning::util::persist::KVStore
 pub struct FilesystemStore {
 	inner: Arc<FilesystemStoreInner>,
+
+	// Version counter to ensure that writes are applied in the correct order. It is assumed that read and list
+	// operations aren't sensitive to the order of execution.
+	next_version: AtomicU64,
 }
 
 impl FilesystemStore {
 	/// Constructs a new [`FilesystemStore`].
 	pub fn new(data_dir: PathBuf) -> Self {
 		let locks = Mutex::new(HashMap::new());
 		let tmp_file_counter = AtomicUsize::new(0);
-		Self { inner: Arc::new(FilesystemStoreInner { data_dir, tmp_file_counter, locks }) }
+		Self {
+			inner: Arc::new(FilesystemStoreInner { data_dir, tmp_file_counter, locks }),
+			next_version: AtomicU64::new(1),
+		}
 	}
 
 	/// Returns the data directory.
 	pub fn get_data_dir(&self) -> PathBuf {
 		self.inner.data_dir.clone()
 	}
 
+	fn get_new_version_and_lock_ref(&self, dest_file_path: PathBuf) -> (Arc<RwLock<u64>>, u64) {
+		let version = self.next_version.fetch_add(1, Ordering::Relaxed);
+		if version == u64::MAX {
+			panic!("FilesystemStore version counter overflowed");
+		}
+
+		// Get a reference to the inner lock. We do this early so that the arc can double as an in-flight counter for
+		// cleaning up unused locks.
+		let inner_lock_ref = self.inner.get_inner_lock_ref(dest_file_path);
+
+		(inner_lock_ref, version)
+	}
+
 	#[cfg(any(all(feature = "tokio", test), fuzzing))]
 	/// Returns the size of the async state.
 	pub fn state_size(&self) -> usize {
@@ -110,8 +120,8 @@ impl KVStoreSync for FilesystemStore {
 			Some(key),
 			"write",
 		)?;
-		let inner_lock_ref = self.inner.get_inner_lock_ref(path.clone());
-		self.inner.write_version(inner_lock_ref, path, buf, None)
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(path.clone());
+		self.inner.write_version(inner_lock_ref, path, buf, version)
 	}
 
 	fn remove(
@@ -123,8 +133,8 @@ impl KVStoreSync for FilesystemStore {
 			Some(key),
 			"remove",
 		)?;
-		let inner_lock_ref = self.inner.get_inner_lock_ref(path.clone());
-		self.inner.remove_version(inner_lock_ref, path, lazy, None)
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(path.clone());
+		self.inner.remove_version(inner_lock_ref, path, lazy, version)
 	}
 
 	fn list(
@@ -141,7 +151,7 @@ impl KVStoreSync for FilesystemStore {
 }
 
 impl FilesystemStoreInner {
-	fn get_inner_lock_ref(&self, path: PathBuf) -> Arc<RwLock<AsyncState>> {
+	fn get_inner_lock_ref(&self, path: PathBuf) -> Arc<RwLock<u64>> {
 		let mut outer_lock = self.locks.lock().unwrap();
 		Arc::clone(&outer_lock.entry(path).or_default())
 	}
@@ -184,34 +194,10 @@ impl FilesystemStoreInner {
 		Ok(dest_file_path)
 	}
 
-	#[cfg(feature = "tokio")]
-	fn get_new_version_and_state(&self, dest_file_path: PathBuf) -> (u64, Arc<RwLock<AsyncState>>) {
-		let inner_lock_ref: Arc<RwLock<AsyncState>> = self.get_inner_lock_ref(dest_file_path);
-
-		let new_version = {
-			let mut async_state = inner_lock_ref.write().unwrap();
-			Self::get_new_version(&mut async_state)
-		};
-
-		return (new_version, inner_lock_ref);
-	}
-
-	fn get_new_version(async_state: &mut AsyncState) -> u64 {
-		async_state.latest_version += 1;
-		if async_state.latest_version == 0 {
-			panic!("FilesystemStore version counter overflowed");
-		}
-
-		debug_assert!(async_state.latest_version > async_state.latest_written_version);
-
-		async_state.latest_version
-	}
-
 	fn read(&self, dest_file_path: PathBuf) -> lightning::io::Result<Vec<u8>> {
 		let mut buf = Vec::new();
 
-		let inner_lock_ref = self.get_inner_lock_ref(dest_file_path.clone());
-		self.execute_locked_read(inner_lock_ref, dest_file_path.clone(), || {
+		self.execute_locked_read(dest_file_path.clone(), || {
 			let mut f = fs::File::open(dest_file_path.clone())?;
 			f.read_to_end(&mut buf)?;
 			Ok(())
@@ -221,24 +207,20 @@ impl FilesystemStoreInner {
 	}
 
 	fn execute_locked_write<F: FnOnce() -> Result<(), lightning::io::Error>>(
-		&self, inner_lock_ref: Arc<RwLock<AsyncState>>, dest_file_path: PathBuf,
-		version: Option<u64>, callback: F,
+		&self, inner_lock_ref: Arc<RwLock<u64>>, dest_file_path: PathBuf, version: u64, callback: F,
 	) -> Result<(), lightning::io::Error> {
-		let mut async_state = inner_lock_ref.write().unwrap();
-
-		// Sync calls haven't assigned a version yet because it would require another lock acquisition.
-		let version = version.unwrap_or_else(|| Self::get_new_version(&mut async_state));
+		let mut last_written_version = inner_lock_ref.write().unwrap();
 
 		// Check if we already have a newer version written/removed. This is used in async contexts to realize eventual
 		// consistency.
-		let is_stale_version = version <= async_state.latest_written_version;
+		let is_stale_version = version <= *last_written_version;
 
 		// If the version is not stale, we execute the callback. Otherwise we can and must skip writing.
 		let res = if is_stale_version {
 			Ok(())
 		} else {
 			callback().map(|_| {
-				async_state.latest_written_version = version;
+				*last_written_version = version;
 			})
 		};
 
@@ -248,15 +230,16 @@ impl FilesystemStoreInner {
 	}
 
 	fn execute_locked_read<F: FnOnce() -> Result<(), lightning::io::Error>>(
-		&self, inner_lock_ref: Arc<RwLock<AsyncState>>, dest_file_path: PathBuf, callback: F,
+		&self, dest_file_path: PathBuf, callback: F,
 	) -> Result<(), lightning::io::Error> {
-		let _async_state = inner_lock_ref.read().unwrap();
+		let inner_lock_ref = self.get_inner_lock_ref(dest_file_path.clone());
+		let _guard = inner_lock_ref.read().unwrap();
 		let res = callback();
 		self.clean_locks(&inner_lock_ref, dest_file_path);
 		res
 	}
 
-	fn clean_locks(&self, inner_lock_ref: &Arc<RwLock<AsyncState>>, dest_file_path: PathBuf) {
+	fn clean_locks(&self, inner_lock_ref: &Arc<RwLock<u64>>, dest_file_path: PathBuf) {
 		// If there no arcs in use elsewhere, this means that there are no in-flight writes. We can remove the map entry
 		// to prevent leaking memory. The two arcs that are expected are the one in the map and the one held here in
 		// inner_lock_ref. The outer lock is obtained first, to avoid a new arc being cloned after we've already
@@ -270,8 +253,8 @@ impl FilesystemStoreInner {
 	/// Writes a specific version of a key to the filesystem. If a newer version has been written already, this function
 	/// returns early without writing.
 	fn write_version(
-		&self, inner_lock_ref: Arc<RwLock<AsyncState>>, dest_file_path: PathBuf, buf: Vec<u8>,
-		version: Option<u64>,
+		&self, inner_lock_ref: Arc<RwLock<u64>>, dest_file_path: PathBuf, buf: Vec<u8>,
+		version: u64,
 	) -> lightning::io::Result<()> {
 		let parent_directory = dest_file_path.parent().ok_or_else(|| {
 			let msg =
@@ -343,8 +326,7 @@ impl FilesystemStoreInner {
 	}
 
 	fn remove_version(
-		&self, inner_lock_ref: Arc<RwLock<AsyncState>>, dest_file_path: PathBuf, lazy: bool,
-		version: Option<u64>,
+		&self, inner_lock_ref: Arc<RwLock<u64>>, dest_file_path: PathBuf, lazy: bool, version: u64,
 	) -> lightning::io::Result<()> {
 		self.execute_locked_write(inner_lock_ref, dest_file_path.clone(), version, || {
 			if !dest_file_path.is_file() {
@@ -507,11 +489,10 @@ impl KVStore for FilesystemStore {
 			Err(e) => return Box::pin(async move { Err(e) }),
 		};
 
-		// Obtain a version number to retain the call sequence.
-		let (version, inner_lock_ref) = self.inner.get_new_version_and_state(path.clone());
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(path.clone());
 		Box::pin(async move {
 			tokio::task::spawn_blocking(move || {
-				this.write_version(inner_lock_ref, path, buf, Some(version))
+				this.write_version(inner_lock_ref, path, buf, version)
 			})
 			.await
 			.unwrap_or_else(|e| Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, e)))
@@ -532,11 +513,10 @@ impl KVStore for FilesystemStore {
 			Err(e) => return Box::pin(async move { Err(e) }),
 		};
 
-		// Obtain a version number to retain the call sequence.
-		let (version, inner_lock_ref) = self.inner.get_new_version_and_state(path.clone());
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(path.clone());
 		Box::pin(async move {
 			tokio::task::spawn_blocking(move || {
-				this.remove_version(inner_lock_ref, path, lazy, Some(version))
+				this.remove_version(inner_lock_ref, path, lazy, version)
 			})
 			.await
 			.unwrap_or_else(|e| Err(lightning::io::Error::new(lightning::io::ErrorKind::Other, e)))