f Use versioned write for SqliteStore

Author: tnull

src/io/sqlite_store/mod.rs

@@ -7,10 +7,12 @@
 
 //! Objects related to [`SqliteStore`] live here.
 use std::boxed::Box;
+use std::collections::HashMap;
 use std::fs;
 use std::future::Future;
 use std::path::PathBuf;
 use std::pin::Pin;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 
 use lightning::io;
@@ -41,6 +43,10 @@ const SCHEMA_USER_VERSION: u16 = 2;
 /// [SQLite]: https://sqlite.org
 pub struct SqliteStore {
 	inner: Arc<SqliteStoreInner>,
+
+	// 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_write_version: AtomicU64,
 }
 
 impl SqliteStore {
@@ -54,7 +60,31 @@ impl SqliteStore {
 		data_dir: PathBuf, db_file_name: Option<String>, kv_table_name: Option<String>,
 	) -> io::Result<Self> {
 		let inner = Arc::new(SqliteStoreInner::new(data_dir, db_file_name, kv_table_name)?);
-		Ok(Self { inner })
+		let next_write_version = AtomicU64::new(1);
+		Ok(Self { inner, next_write_version })
+	}
+
+	fn build_locking_key(
+		&self, primary_namespace: &str, secondary_namespace: &str, key: &str,
+	) -> String {
+		if primary_namespace.is_empty() {
+			key.to_owned()
+		} else {
+			format!("{}#{}#{}", primary_namespace, secondary_namespace, key)
+		}
+	}
+
+	fn get_new_version_and_lock_ref(&self, locking_key: String) -> (Arc<Mutex<u64>>, u64) {
+		let version = self.next_write_version.fetch_add(1, Ordering::Relaxed);
+		if version == u64::MAX {
+			panic!("SqliteStore 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(locking_key);
+
+		(inner_lock_ref, version)
 	}
 
 	/// Returns the data directory.
@@ -85,12 +115,22 @@ impl KVStore for SqliteStore {
 	fn write(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
 	) -> Pin<Box<dyn Future<Output = Result<(), io::Error>> + Send>> {
+		let locking_key = self.build_locking_key(primary_namespace, secondary_namespace, key);
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(locking_key.clone());
 		let primary_namespace = primary_namespace.to_string();
 		let secondary_namespace = secondary_namespace.to_string();
 		let key = key.to_string();
 		let inner = Arc::clone(&self.inner);
 		let fut = tokio::task::spawn_blocking(move || {
-			inner.write_internal(&primary_namespace, &secondary_namespace, &key, buf)
+			inner.write_internal(
+				inner_lock_ref,
+				locking_key,
+				version,
+				&primary_namespace,
+				&secondary_namespace,
+				&key,
+				buf,
+			)
 		});
 		Box::pin(async move {
 			fut.await.unwrap_or_else(|e| {
@@ -103,12 +143,22 @@ impl KVStore for SqliteStore {
 	fn remove(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, lazy: bool,
 	) -> Pin<Box<dyn Future<Output = Result<(), io::Error>> + Send>> {
+		let locking_key = self.build_locking_key(primary_namespace, secondary_namespace, key);
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(locking_key.clone());
 		let primary_namespace = primary_namespace.to_string();
 		let secondary_namespace = secondary_namespace.to_string();
 		let key = key.to_string();
 		let inner = Arc::clone(&self.inner);
 		let fut = tokio::task::spawn_blocking(move || {
-			inner.remove_internal(&primary_namespace, &secondary_namespace, &key, lazy)
+			inner.remove_internal(
+				inner_lock_ref,
+				locking_key,
+				version,
+				&primary_namespace,
+				&secondary_namespace,
+				&key,
+				lazy,
+			)
 		});
 		Box::pin(async move {
 			fut.await.unwrap_or_else(|e| {
@@ -146,13 +196,33 @@ impl KVStoreSync for SqliteStore {
 	fn write(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
 	) -> io::Result<()> {
-		self.inner.write_internal(primary_namespace, secondary_namespace, key, buf)
+		let locking_key = self.build_locking_key(primary_namespace, secondary_namespace, key);
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(locking_key.clone());
+		self.inner.write_internal(
+			inner_lock_ref,
+			locking_key,
+			version,
+			primary_namespace,
+			secondary_namespace,
+			key,
+			buf,
+		)
 	}
 
 	fn remove(
 		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, lazy: bool,
 	) -> io::Result<()> {
-		self.inner.remove_internal(primary_namespace, secondary_namespace, key, lazy)
+		let locking_key = self.build_locking_key(primary_namespace, secondary_namespace, key);
+		let (inner_lock_ref, version) = self.get_new_version_and_lock_ref(locking_key.clone());
+		self.inner.remove_internal(
+			inner_lock_ref,
+			locking_key,
+			version,
+			primary_namespace,
+			secondary_namespace,
+			key,
+			lazy,
+		)
 	}
 
 	fn list(&self, primary_namespace: &str, secondary_namespace: &str) -> io::Result<Vec<String>> {
@@ -164,6 +234,7 @@ struct SqliteStoreInner {
 	connection: Arc<Mutex<Connection>>,
 	data_dir: PathBuf,
 	kv_table_name: String,
+	write_version_locks: Mutex<HashMap<String, Arc<Mutex<u64>>>>,
 }
 
 impl SqliteStoreInner {
@@ -237,7 +308,13 @@ impl SqliteStoreInner {
 		})?;
 
 		let connection = Arc::new(Mutex::new(connection));
-		Ok(Self { connection, data_dir, kv_table_name })
+		let write_version_locks = Mutex::new(HashMap::new());
+		Ok(Self { connection, data_dir, kv_table_name, write_version_locks })
+	}
+
+	fn get_inner_lock_ref(&self, locking_key: String) -> Arc<Mutex<u64>> {
+		let mut outer_lock = self.write_version_locks.lock().unwrap();
+		Arc::clone(&outer_lock.entry(locking_key).or_default())
 	}
 
 	fn read_internal(
@@ -289,46 +366,51 @@ impl SqliteStoreInner {
 	}
 
 	fn write_internal(
-		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
+		&self, inner_lock_ref: Arc<Mutex<u64>>, locking_key: String, version: u64,
+		primary_namespace: &str, secondary_namespace: &str, key: &str, buf: Vec<u8>,
 	) -> io::Result<()> {
 		check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "write")?;
 
-		let locked_conn = self.connection.lock().unwrap();
+		self.execute_locked_write(inner_lock_ref, locking_key, version, || {
+			let locked_conn = self.connection.lock().unwrap();
 
-		let sql = format!(
-			"INSERT OR REPLACE INTO {} (primary_namespace, secondary_namespace, key, value) VALUES (:primary_namespace, :secondary_namespace, :key, :value);",
-			self.kv_table_name
-		);
+			let sql = format!(
+				"INSERT OR REPLACE INTO {} (primary_namespace, secondary_namespace, key, value) VALUES (:primary_namespace, :secondary_namespace, :key, :value);",
+				self.kv_table_name
+			);
 
-		let mut stmt = locked_conn.prepare_cached(&sql).map_err(|e| {
-			let msg = format!("Failed to prepare statement: {}", e);
-			io::Error::new(io::ErrorKind::Other, msg)
-		})?;
+			let mut stmt = locked_conn.prepare_cached(&sql).map_err(|e| {
+				let msg = format!("Failed to prepare statement: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?;
 
-		stmt.execute(named_params! {
-			":primary_namespace": primary_namespace,
-			":secondary_namespace": secondary_namespace,
-			":key": key,
-			":value": buf,
-		})
-		.map(|_| ())
-		.map_err(|e| {
-			let msg = format!(
-				"Failed to write to key {}/{}/{}: {}",
-				PrintableString(primary_namespace),
-				PrintableString(secondary_namespace),
-				PrintableString(key),
-				e
-			);
-			io::Error::new(io::ErrorKind::Other, msg)
+			stmt.execute(named_params! {
+				":primary_namespace": primary_namespace,
+				":secondary_namespace": secondary_namespace,
+				":key": key,
+				":value": buf,
+			})
+			.map(|_| ())
+			.map_err(|e| {
+				let msg = format!(
+					"Failed to write to key {}/{}/{}: {}",
+					PrintableString(primary_namespace),
+					PrintableString(secondary_namespace),
+					PrintableString(key),
+					e
+				);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})
 		})
 	}
 
 	fn remove_internal(
-		&self, primary_namespace: &str, secondary_namespace: &str, key: &str, _lazy: bool,
+		&self, inner_lock_ref: Arc<Mutex<u64>>, locking_key: String, version: u64,
+		primary_namespace: &str, secondary_namespace: &str, key: &str, _lazy: bool,
 	) -> io::Result<()> {
 		check_namespace_key_validity(primary_namespace, secondary_namespace, Some(key), "remove")?;
 
+		self.execute_locked_write(inner_lock_ref, locking_key, version, || {
 		let locked_conn = self.connection.lock().unwrap();
 
 		let sql = format!("DELETE FROM {} WHERE primary_namespace=:primary_namespace AND secondary_namespace=:secondary_namespace AND key=:key;", self.kv_table_name);
@@ -354,6 +436,7 @@ impl SqliteStoreInner {
 			io::Error::new(io::ErrorKind::Other, msg)
 		})?;
 		Ok(())
+        })
 	}
 
 	fn list_internal(
@@ -396,6 +479,46 @@ impl SqliteStoreInner {
 
 		Ok(keys)
 	}
+
+	fn execute_locked_write<F: FnOnce() -> Result<(), lightning::io::Error>>(
+		&self, inner_lock_ref: Arc<Mutex<u64>>, locking_key: String, version: u64, callback: F,
+	) -> Result<(), lightning::io::Error> {
+		let res = {
+			let mut last_written_version = inner_lock_ref.lock().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 <= *last_written_version;
+
+			// If the version is not stale, we execute the callback. Otherwise we can and must skip writing.
+			if is_stale_version {
+				Ok(())
+			} else {
+				callback().map(|_| {
+					*last_written_version = version;
+				})
+			}
+		};
+
+		self.clean_locks(&inner_lock_ref, locking_key);
+
+		res
+	}
+
+	fn clean_locks(&self, inner_lock_ref: &Arc<Mutex<u64>>, locking_key: String) {
+		// 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
+		// counted.
+		let mut outer_lock = self.write_version_locks.lock().unwrap();
+
+		let strong_count = Arc::strong_count(&inner_lock_ref);
+		debug_assert!(strong_count >= 2, "Unexpected SqliteStore strong count");
+
+		if strong_count == 2 {
+			outer_lock.remove(&locking_key);
+		}
+	}
 }
 
 #[cfg(test)]