fix(udb): optimize pg driver

Author: MasterPtato

out/openapi.json

@@ -1,7 +1,11 @@
-use std::sync::{Arc, Mutex};
+use std::{
+	sync::{Arc, Mutex},
+	time::Duration,
+};
 
 use anyhow::{Context, Result};
 use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime};
+use tokio::task::JoinHandle;
 use tokio_postgres::NoTls;
 
 use crate::{
@@ -14,9 +18,13 @@ use crate::{
 
 use super::transaction::PostgresTransactionDriver;
 
+const TXN_TIMEOUT: Duration = Duration::from_secs(5);
+const GC_INTERVAL: Duration = Duration::from_secs(5);
+
 pub struct PostgresDatabaseDriver {
 	pool: Arc<Pool>,
 	max_retries: Arc<Mutex<i32>>,
+	gc_handle: JoinHandle<()>,
 }
 
 impl PostgresDatabaseDriver {
@@ -53,7 +61,7 @@ impl PostgresDatabaseDriver {
 			.context("failed to create btree_gist extension")?;
 
 		conn.execute(
-			"CREATE SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1",
+			"CREATE UNLOGGED SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1",
 			&[],
 		)
 		.await
@@ -123,12 +131,39 @@ impl PostgresDatabaseDriver {
 		.await
 		.context("failed to create conflict_ranges table")?;
 
-		// Connection is automatically returned to the pool when dropped
-		drop(conn);
+		// Create index on ts column for efficient garbage collection
+		conn.execute(
+			"CREATE INDEX IF NOT EXISTS idx_conflict_ranges_ts ON conflict_ranges (ts)",
+			&[],
+		)
+		.await
+		.context("failed to create index on conflict_ranges ts column")?;
+
+		let gc_handle = tokio::spawn(async move {
+			let mut interval = tokio::time::interval(GC_INTERVAL);
+			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
+			loop {
+				interval.tick().await;
+
+				// NOTE: Transactions have a max limit of 5 seconds, we delete after 10 seconds for extra padding
+				// Delete old conflict ranges
+				if let Err(err) = conn
+					.execute(
+						"DELETE FROM conflict_ranges where ts < now() - interval '10 seconds'",
+						&[],
+					)
+					.await
+				{
+					tracing::error!(?err, "failed postgres gc task");
+				}
+			}
+		});
 
 		Ok(PostgresDatabaseDriver {
 			pool: Arc::new(pool),
 			max_retries: Arc::new(Mutex::new(100)),
+			gc_handle,
 		})
 	}
 }
@@ -155,13 +190,15 @@ impl DatabaseDriver for PostgresDatabaseDriver {
 				retryable.maybe_committed = maybe_committed;
 
 				// Execute transaction
-				let error = match closure(retryable.clone()).await {
-					Ok(res) => match retryable.inner.driver.commit_ref().await {
-						Ok(_) => return Ok(res),
-						Err(e) => e,
-					},
-					Err(e) => e,
-				};
+				let error =
+					match tokio::time::timeout(TXN_TIMEOUT, closure(retryable.clone())).await {
+						Ok(Ok(res)) => match retryable.inner.driver.commit_ref().await {
+							Ok(_) => return Ok(res),
+							Err(e) => e,
+						},
+						Ok(Err(e)) => e,
+						Err(e) => anyhow::Error::from(DatabaseError::TransactionTooOld),
+					};
 
 				let chain = error
 					.chain()
@@ -196,3 +233,9 @@ impl DatabaseDriver for PostgresDatabaseDriver {
 		}
 	}
 }
+
+impl Drop for PostgresDatabaseDriver {
+	fn drop(&mut self) {
+		self.gc_handle.abort();
+	}
+}

packages/common/universaldb/src/driver/postgres/database.rs

@@ -330,11 +330,20 @@ impl TransactionTask {
 		tx: Transaction<'_>,
 		start_version: i64,
 		operations: Vec<Operation>,
-		conflict_ranges: Vec<(Vec<u8>, Vec<u8>, ConflictRangeType)>,
+		mut conflict_ranges: Vec<(Vec<u8>, Vec<u8>, ConflictRangeType)>,
 	) -> Result<()> {
-		let commit_version = tx
-			.query_one("SELECT nextval('global_version_seq')", &[])
-			.await
+		// // Defer all constraint checks until commit
+		// tx.execute("SET CONSTRAINTS ALL DEFERRED", &[])
+		// 	.await
+		// 	.map_err(map_postgres_error)?;
+
+		let (_, _, version_res) = tokio::join!(
+			tx.execute("SET LOCAL lock_timeout = '0'", &[],),
+			tx.execute("SET LOCAL deadlock_timeout = '10ms'", &[],),
+			tx.query_one("SELECT nextval('global_version_seq')", &[]),
+		);
+
+		let commit_version = version_res
 			.context("failed to get postgres txn commit_version")?
 			.get::<_, i64>(0);
 
@@ -355,7 +364,7 @@ impl TransactionTask {
 
 		let query = "
 			INSERT INTO conflict_ranges (range_data, conflict_type, start_version, commit_version)
-			SELECT 
+			SELECT
 				bytearange(begin_key, end_key, '[)'),
 				conflict_type::range_type,
 				$4,
@@ -377,13 +386,22 @@ impl TransactionTask {
 		.await
 		.map_err(map_postgres_error)?;
 
-		// TODO: Parallelize
 		for op in operations {
 			match op {
 				Operation::Set { key, value } => {
 					// TODO: versionstamps need to be calculated on the sql side, not in rust
 					let value = substitute_versionstamp_if_incomplete(value.clone(), 0);
 
+					// // Poor man's upsert, you cant use ON CONFLICT with deferred constraints
+					// let query = "WITH updated AS (
+					// 		UPDATE kv
+					// 		SET value = $2
+					// 		WHERE key = $1
+					// 		RETURNING 1
+					// 	)
+					// 	INSERT INTO kv (key, value)
+					// 	SELECT $1, $2
+					// 	WHERE NOT EXISTS (SELECT 1 FROM updated)";
 					let query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2";
 					let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?;
 
@@ -435,6 +453,16 @@ impl TransactionTask {
 
 					// Store the result
 					if let Some(new_value) = new_value {
+						// // Poor man's upsert, you cant use ON CONFLICT with deferred constraints
+						// let update_query = "WITH updated AS (
+						// 		UPDATE kv
+						// 		SET value = $2
+						// 		WHERE key = $1
+						// 		RETURNING 1
+						// 	)
+						// 	INSERT INTO kv (key, value)
+						// 	SELECT $1, $2
+						// 	WHERE NOT EXISTS (SELECT 1 FROM updated)";
 						let update_query = "INSERT INTO kv (key, value) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET value = $2";
 						let stmt = tx
 							.prepare_cached(update_query)

packages/common/universaldb/src/driver/postgres/transaction_task.rs

@@ -264,7 +264,7 @@ pub enum MutationType {
 	/// Performs an atomic ``compare and clear`` operation. If the existing value in the database is equal to the given value, then given key is cleared.
 	CompareAndClear,
 }
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
 #[non_exhaustive]
 pub enum ConflictRangeType {
 	/// Used to add a read conflict range

packages/common/universaldb/src/options.rs

@@ -34,6 +34,17 @@ pub enum Operation {
 	},
 }
 
+impl Operation {
+	pub fn sorting_key(&self) -> &[u8] {
+		match self {
+			Operation::Set { key, .. } => key,
+			Operation::Clear { key } => key,
+			Operation::ClearRange { begin, .. } => begin,
+			Operation::AtomicOp { key, .. } => key,
+		}
+	}
+}
+
 #[derive(Debug, Clone)]
 pub enum GetOutput {
 	Value(Vec<u8>),