@rivetkit/engine 1.0.0

package/packages/universaldb/src/driver/mod.rs

@@ -0,0 +1,87 @@
+use std::{any::Any, future::Future, pin::Pin, sync::Arc};
+
+use anyhow::{Result, bail};
+
+use crate::{
+	key_selector::KeySelector,
+	options::{ConflictRangeType, DatabaseOption, MutationType},
+	range_option::RangeOption,
+	transaction::{RetryableTransaction, Transaction},
+	utils::IsolationLevel,
+	value::{Slice, Value, Values},
+};
+
+mod postgres;
+pub mod rocksdb;
+
+pub use postgres::PostgresDatabaseDriver;
+pub use rocksdb::RocksDbDatabaseDriver;
+
+pub type BoxFut<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
+pub type Erased = Box<dyn Any + Send>;
+
+pub type DatabaseDriverHandle = Arc<dyn DatabaseDriver>;
+
+pub trait DatabaseDriver: Send + Sync {
+	fn create_trx(&self) -> Result<Transaction>;
+	fn run<'a>(
+		&'a self,
+		closure: Box<dyn Fn(RetryableTransaction) -> BoxFut<'a, Result<Erased>> + Send + Sync + 'a>,
+	) -> BoxFut<'a, Result<Erased>>;
+	fn set_option(&self, opt: DatabaseOption) -> Result<()>;
+}
+
+pub trait TransactionDriver: Send + Sync {
+	fn atomic_op(&self, key: &[u8], param: &[u8], op_type: MutationType);
+
+	// Read operations
+	fn get<'a>(
+		&'a self,
+		key: &[u8],
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Option<Slice>>> + Send + 'a>>;
+	fn get_key<'a>(
+		&'a self,
+		selector: &KeySelector<'a>,
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Slice>> + Send + 'a>>;
+	fn get_range<'a>(
+		&'a self,
+		opt: &RangeOption<'a>,
+		iteration: usize,
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Values>> + Send + 'a>>;
+	fn get_ranges_keyvalues<'a>(
+		&'a self,
+		opt: RangeOption<'a>,
+		isolation_level: IsolationLevel,
+	) -> crate::value::Stream<'a, Value>;
+
+	// Write operations
+	fn set(&self, key: &[u8], value: &[u8]);
+	fn clear(&self, key: &[u8]);
+	fn clear_range(&self, begin: &[u8], end: &[u8]);
+
+	// Transaction management
+	fn commit(self: Box<Self>) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>;
+	fn reset(&mut self);
+	fn cancel(&self);
+	fn add_conflict_range(
+		&self,
+		begin: &[u8],
+		end: &[u8],
+		conflict_type: ConflictRangeType,
+	) -> Result<()>;
+	fn get_estimated_range_size_bytes<'a>(
+		&'a self,
+		begin: &'a [u8],
+		end: &'a [u8],
+	) -> Pin<Box<dyn Future<Output = Result<i64>> + Send + 'a>>;
+
+	// Helper for committing without consuming self (for database drivers that need it)
+	fn commit_ref(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
+		Box::pin(async move {
+			bail!("`commit_ref` unimplemented");
+		})
+	}
+}

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

@@ -0,0 +1,244 @@
+use std::{
+	sync::{
+		Arc,
+		atomic::{AtomicI32, Ordering},
+	},
+	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::{
+	RetryableTransaction, Transaction,
+	driver::{BoxFut, DatabaseDriver, Erased},
+	error::DatabaseError,
+	options::DatabaseOption,
+	utils::{MaybeCommitted, calculate_tx_retry_backoff},
+};
+
+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: AtomicI32,
+	gc_handle: JoinHandle<()>,
+}
+
+impl PostgresDatabaseDriver {
+	pub async fn new(connection_string: String) -> Result<Self> {
+		tracing::debug!(connection_string = ?connection_string, "Creating PostgresDatabaseDriver");
+
+		// Create deadpool config from connection string
+		let mut config = Config::new();
+		config.url = Some(connection_string);
+		config.pool = Some(PoolConfig {
+			max_size: 64,
+			..Default::default()
+		});
+		config.manager = Some(ManagerConfig {
+			recycling_method: RecyclingMethod::Fast,
+		});
+
+		tracing::debug!("Creating Postgres pool");
+		// Create the pool
+		let pool = config
+			.create_pool(Some(Runtime::Tokio1), NoTls)
+			.context("failed to create postgres connection pool")?;
+
+		tracing::debug!("Getting Postgres connection from pool");
+		// Get a connection from the pool to create the table
+		let conn = pool
+			.get()
+			.await
+			.context("failed to get connection from postgres pool")?;
+
+		// Enable btree gist
+		conn.execute("CREATE EXTENSION IF NOT EXISTS btree_gist", &[])
+			.await
+			.context("failed to create btree_gist extension")?;
+
+		conn.execute(
+			"CREATE UNLOGGED SEQUENCE IF NOT EXISTS global_version_seq START WITH 1 INCREMENT BY 1 MINVALUE 1",
+			&[],
+		)
+		.await
+		.context("failed to create global version sequence")?;
+
+		// Create the KV table if it doesn't exist
+		conn.execute(
+			"CREATE TABLE IF NOT EXISTS kv (
+				key BYTEA PRIMARY KEY,
+				value BYTEA NOT NULL
+			)",
+			&[],
+		)
+		.await
+		.context("failed to create kv table")?;
+
+		// Create range_type type if it doesn't exist
+		conn.execute(
+			"DO $$ BEGIN
+				CREATE TYPE range_type AS ENUM ('read', 'write');
+			EXCEPTION
+				WHEN duplicate_object THEN null;
+			END $$",
+			&[],
+		)
+		.await
+		.context("failed to create range_type enum")?;
+
+		// Create bytearange type if it doesn't exist
+		conn.execute(
+			"DO $$ BEGIN
+				CREATE TYPE bytearange AS RANGE (
+					SUBTYPE = bytea,
+					SUBTYPE_OPCLASS = bytea_ops
+				);
+			EXCEPTION
+				WHEN duplicate_object THEN null;
+			END $$",
+			&[],
+		)
+		.await
+		.context("failed to create bytearange type")?;
+
+		// Create the conflict ranges table for non-snapshot reads
+		// This enforces consistent reads for ranges by preventing overlapping conflict ranges
+		conn.execute(
+			"CREATE UNLOGGED TABLE IF NOT EXISTS conflict_ranges (
+				range_data BYTEARANGE NOT NULL,
+				conflict_type range_type NOT NULL,
+				start_version BIGINT NOT NULL,
+				commit_version BIGINT NOT NULL,
+				ts timestamp NOT NULL DEFAULT now(),
+
+				EXCLUDE USING gist (
+					-- Conflict if byte range overlaps...
+					range_data WITH &&,
+					-- And f conflict types are different...
+					conflict_type WITH <>,
+					-- And f the txn versions overlap...
+					int8range(start_version, commit_version, '[]') WITH &&,
+					-- But not if the start_version is the same (from the same txn)
+					start_version WITH <>
+				)
+			)",
+			&[],
+		)
+		.await
+		.context("failed to create conflict_ranges table")?;
+
+		// 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: AtomicI32::new(100),
+			gc_handle,
+		})
+	}
+}
+
+impl DatabaseDriver for PostgresDatabaseDriver {
+	fn create_trx(&self) -> Result<Transaction> {
+		// Pass the connection pool to the transaction driver
+		Ok(Transaction::new(Arc::new(PostgresTransactionDriver::new(
+			self.pool.clone(),
+		))))
+	}
+
+	fn run<'a>(
+		&'a self,
+		closure: Box<dyn Fn(RetryableTransaction) -> BoxFut<'a, Result<Erased>> + Send + Sync + 'a>,
+	) -> BoxFut<'a, Result<Erased>> {
+		Box::pin(async move {
+			let mut maybe_committed = MaybeCommitted(false);
+			let max_retries = self.max_retries.load(Ordering::SeqCst);
+
+			for attempt in 0..max_retries {
+				let tx = self.create_trx()?;
+				let mut retryable = RetryableTransaction::new(tx);
+				retryable.maybe_committed = maybe_committed;
+
+				// Execute transaction
+				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(_) => anyhow::Error::from(DatabaseError::TransactionTooOld),
+					};
+
+				let chain = error
+					.chain()
+					.find_map(|x| x.downcast_ref::<DatabaseError>());
+
+				if let Some(db_error) = chain {
+					// Handle retry or return error
+					if db_error.is_retryable() {
+						if db_error.is_maybe_committed() {
+							maybe_committed = MaybeCommitted(true);
+						}
+
+						let backoff_ms = calculate_tx_retry_backoff(attempt as usize);
+						tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await;
+						continue;
+					}
+				}
+
+				return Err(error);
+			}
+
+			Err(DatabaseError::MaxRetriesReached.into())
+		})
+	}
+
+	fn set_option(&self, opt: DatabaseOption) -> Result<()> {
+		match opt {
+			DatabaseOption::TransactionRetryLimit(limit) => {
+				self.max_retries.store(limit, Ordering::SeqCst);
+				Ok(())
+			}
+		}
+	}
+}
+
+impl Drop for PostgresDatabaseDriver {
+	fn drop(&mut self) {
+		self.gc_handle.abort();
+	}
+}

package/packages/universaldb/src/driver/postgres/mod.rs

@@ -0,0 +1,5 @@
+mod database;
+mod transaction;
+mod transaction_task;
+
+pub use database::PostgresDatabaseDriver;

package/packages/universaldb/src/driver/postgres/transaction.rs

@@ -0,0 +1,337 @@
+use std::{
+	future::Future,
+	pin::Pin,
+	sync::{
+		Arc,
+		atomic::{AtomicBool, Ordering},
+	},
+};
+
+use anyhow::{Context, Result};
+use deadpool_postgres::Pool;
+use tokio::sync::{OnceCell, mpsc, oneshot};
+
+use crate::{
+	driver::TransactionDriver,
+	key_selector::KeySelector,
+	options::{ConflictRangeType, MutationType},
+	range_option::RangeOption,
+	tx_ops::TransactionOperations,
+	utils::IsolationLevel,
+	value::{Slice, Value, Values},
+};
+
+use super::transaction_task::{TransactionCommand, TransactionTask};
+
+pub struct PostgresTransactionDriver {
+	pool: Arc<Pool>,
+	operations: TransactionOperations,
+	committed: AtomicBool,
+	tx_sender: OnceCell<mpsc::Sender<TransactionCommand>>,
+}
+
+impl PostgresTransactionDriver {
+	pub fn new(pool: Arc<Pool>) -> Self {
+		PostgresTransactionDriver {
+			pool,
+			operations: TransactionOperations::default(),
+			committed: AtomicBool::new(false),
+			tx_sender: OnceCell::new(),
+		}
+	}
+
+	/// Get or create the transaction task
+	async fn ensure_transaction(&self) -> Result<&mpsc::Sender<TransactionCommand>> {
+		self.tx_sender
+			.get_or_try_init(|| async {
+				let (sender, receiver) = mpsc::channel(100);
+
+				// Spawn the transaction task with serializable isolation
+				let task = TransactionTask::new(self.pool.as_ref().clone(), receiver);
+				tokio::spawn(task.run());
+
+				anyhow::Ok(sender)
+			})
+			.await
+			.context("failed to initialize postgres transaction task")
+	}
+}
+
+impl TransactionDriver for PostgresTransactionDriver {
+	fn atomic_op(&self, key: &[u8], param: &[u8], op_type: MutationType) {
+		self.operations.atomic_op(key, param, op_type);
+	}
+
+	fn get<'a>(
+		&'a self,
+		key: &[u8],
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Option<Slice>>> + Send + 'a>> {
+		let key = key.to_vec();
+
+		Box::pin(async move {
+			self.operations
+				.get_with_callback(&key, isolation_level, || async {
+					let tx_sender = self.ensure_transaction().await?;
+
+					// Send query command
+					let (response_tx, response_rx) = oneshot::channel();
+					tx_sender
+						.send(TransactionCommand::Get {
+							key: key.clone(),
+							response: response_tx,
+						})
+						.await
+						.context("failed to send postgres transaction command")?;
+
+					// Wait for response
+					response_rx
+						.await
+						.context("failed to receive postgres response")?
+				})
+				.await
+		})
+	}
+
+	fn get_key<'a>(
+		&'a self,
+		selector: &KeySelector<'a>,
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Slice>> + Send + 'a>> {
+		let selector = selector.clone();
+
+		Box::pin(async move {
+			let key = selector.key().to_vec();
+			let offset = selector.offset();
+			let or_equal = selector.or_equal();
+
+			self.operations
+				.get_key(&selector, isolation_level, || async {
+					let tx_sender = self.ensure_transaction().await?;
+
+					// Send query command
+					let (response_tx, response_rx) = oneshot::channel();
+					tx_sender
+						.send(TransactionCommand::GetKey {
+							key: key.clone(),
+							or_equal,
+							offset,
+							response: response_tx,
+						})
+						.await
+						.context("failed to send postgres transaction command")?;
+
+					// Wait for response
+					let result_key = response_rx
+						.await
+						.context("failed to receive postgres key selector response")??;
+
+					// Return the key if found, or empty vector if not
+					Ok(result_key.unwrap_or_else(Slice::new))
+				})
+				.await
+		})
+	}
+
+	fn get_range<'a>(
+		&'a self,
+		opt: &RangeOption<'a>,
+		_iteration: usize,
+		isolation_level: IsolationLevel,
+	) -> Pin<Box<dyn Future<Output = Result<Values>> + Send + 'a>> {
+		let opt = opt.clone();
+
+		Box::pin(async move {
+			let begin = opt.begin.key().to_vec();
+			let begin_or_equal = opt.begin.or_equal();
+			let begin_offset = opt.begin.offset();
+			let end = opt.end.key().to_vec();
+			let end_or_equal = opt.end.or_equal();
+			let end_offset = opt.end.offset();
+			let limit = opt.limit;
+			let reverse = opt.reverse;
+
+			self.operations
+				.get_range(&opt, isolation_level, || async {
+					let tx_sender = self.ensure_transaction().await?;
+
+					// Send query command
+					let (response_tx, response_rx) = oneshot::channel();
+					tx_sender
+						.send(TransactionCommand::GetRange {
+							begin: begin.clone(),
+							begin_or_equal,
+							begin_offset,
+							end: end.clone(),
+							end_or_equal,
+							end_offset,
+							limit,
+							reverse,
+							response: response_tx,
+						})
+						.await
+						.context("failed to send postgres transaction command")?;
+
+					// Wait for response
+					response_rx
+						.await
+						.context("failed to receive postgres range response")?
+				})
+				.await
+		})
+	}
+
+	fn get_ranges_keyvalues<'a>(
+		&'a self,
+		opt: RangeOption<'a>,
+		isolation_level: IsolationLevel,
+	) -> crate::value::Stream<'a, Value> {
+		use futures_util::{StreamExt, stream};
+
+		// Convert the range result into a stream
+		let fut = async move {
+			match self.get_range(&opt, 1, isolation_level).await {
+				Ok(values) => values
+					.into_iter()
+					.map(|kv| Ok(Value::from_keyvalue(kv)))
+					.collect::<Vec<_>>(),
+				Err(e) => vec![Err(e)],
+			}
+		};
+
+		Box::pin(stream::once(fut).flat_map(stream::iter))
+	}
+
+	fn set(&self, key: &[u8], value: &[u8]) {
+		self.operations.set(key, value);
+	}
+
+	fn clear(&self, key: &[u8]) {
+		self.operations.clear(key);
+	}
+
+	fn clear_range(&self, begin: &[u8], end: &[u8]) {
+		self.operations.clear_range(begin, end);
+	}
+
+	fn commit(self: Box<Self>) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
+		Box::pin(async move {
+			if self.committed.load(Ordering::SeqCst) {
+				return Ok(());
+			}
+			self.committed.store(true, Ordering::SeqCst);
+
+			let (operations, conflict_ranges) = self.operations.consume();
+
+			let tx_sender = self.ensure_transaction().await?;
+
+			// Send commit command
+			let (response_tx, response_rx) = oneshot::channel();
+			tx_sender
+				.send(TransactionCommand::Commit {
+					operations,
+					conflict_ranges,
+					response: response_tx,
+				})
+				.await
+				.context("failed to send postgres transaction command")?;
+
+			// Wait for commit response
+			response_rx
+				.await
+				.context("failed to receive postgres commit response")??;
+
+			Ok(())
+		})
+	}
+
+	fn reset(&mut self) {
+		self.operations.clear_all();
+		self.committed.store(false, Ordering::SeqCst);
+
+		// Replace tx sender to get a new txn version
+		self.tx_sender = OnceCell::new();
+	}
+
+	fn cancel(&self) {
+		self.operations.clear_all();
+		self.committed.store(true, Ordering::SeqCst); // Prevent future commits
+
+		// Transaction will be rolled back when dropped
+	}
+
+	fn add_conflict_range(
+		&self,
+		begin: &[u8],
+		end: &[u8],
+		conflict_type: ConflictRangeType,
+	) -> Result<()> {
+		self.operations
+			.add_conflict_range(begin, end, conflict_type);
+
+		Ok(())
+	}
+
+	fn get_estimated_range_size_bytes<'a>(
+		&'a self,
+		begin: &'a [u8],
+		end: &'a [u8],
+	) -> Pin<Box<dyn Future<Output = Result<i64>> + Send + 'a>> {
+		let begin = begin.to_vec();
+		let end = end.to_vec();
+
+		Box::pin(async move {
+			let tx_sender = self.ensure_transaction().await?;
+
+			// Send query command
+			let (response_tx, response_rx) = oneshot::channel();
+			tx_sender
+				.send(TransactionCommand::GetEstimatedRangeSize {
+					begin,
+					end,
+					response: response_tx,
+				})
+				.await
+				.context("failed to send postgres command")?;
+
+			// Wait for response
+			let size = response_rx
+				.await
+				.context("failed to receive postgres size response")??;
+
+			Ok(size)
+		})
+	}
+
+	fn commit_ref(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
+		Box::pin(async move {
+			if self.committed.load(Ordering::SeqCst) {
+				return Ok(());
+			}
+			self.committed.store(true, Ordering::SeqCst);
+
+			let (operations, conflict_ranges) = self.operations.consume();
+
+			// We have operations but no transaction - create one just for commit
+			let tx_sender = self.ensure_transaction().await?;
+
+			// Send commit command
+			let (response_tx, response_rx) = oneshot::channel();
+			tx_sender
+				.send(TransactionCommand::Commit {
+					operations,
+					conflict_ranges,
+					response: response_tx,
+				})
+				.await
+				.context("failed to send postgres transaction command")?;
+
+			// Wait for commit response
+			response_rx
+				.await
+				.context("failed to receive postgres commit response")??;
+
+			Ok(())
+		})
+	}
+}