@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

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

@@ -1,4 +1,4 @@
-use std::{any::Any, future::Future, pin::Pin, sync::Arc};
+use std::{any::Any, future::Future, path::Path, pin::Pin, sync::Arc};
 
 use anyhow::{Result, bail};
 
@@ -11,7 +11,7 @@ use crate::{
 	value::{Slice, Value, Values},
 };
 
-mod postgres;
+pub mod postgres;
 pub mod rocksdb;
 
 pub use postgres::PostgresDatabaseDriver;
@@ -29,6 +29,11 @@ pub trait DatabaseDriver: Send + Sync {
 		closure: Box<dyn Fn(RetryableTransaction) -> BoxFut<'a, Result<Erased>> + Send + Sync + 'a>,
 	) -> BoxFut<'a, Result<Erased>>;
 	fn set_option(&self, opt: DatabaseOption) -> Result<()>;
+
+	/// Create a consistent point-in-time snapshot of the database at the given path.
+	fn checkpoint(&self, _path: &Path) -> Result<()> {
+		bail!("checkpoint not supported by this database driver")
+	}
 }
 
 pub trait TransactionDriver: Send + Sync {
@@ -78,6 +83,11 @@ pub trait TransactionDriver: Send + Sync {
 		end: &'a [u8],
 	) -> Pin<Box<dyn Future<Output = Result<i64>> + Send + 'a>>;
 
+	fn tag(&self, _tag: &str) -> Result<()> {
+		// No-op unless implemented
+		Ok(())
+	}
+
 	// 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 {

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

@@ -1,4 +1,5 @@
 use std::{
+	path::PathBuf,
 	sync::{
 		Arc,
 		atomic::{AtomicI32, Ordering},
@@ -8,48 +9,108 @@ use std::{
 
 use anyhow::{Context, Result};
 use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime};
+use rivet_postgres_util::build_tls_config;
 use tokio::task::JoinHandle;
-use tokio_postgres::NoTls;
+use tokio_postgres_rustls::MakeRustlsConnect;
+use url::Url;
 
 use crate::{
 	RetryableTransaction, Transaction,
 	driver::{BoxFut, DatabaseDriver, Erased},
 	error::DatabaseError,
 	options::DatabaseOption,
+	transaction::TXN_TIMEOUT,
 	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);
 
+#[derive(Clone, Debug)]
+pub struct PostgresConfig {
+	pub connection_string: String,
+	pub ssl_config: Option<PostgresSslConfig>,
+}
+
+#[derive(Clone, Debug)]
+pub struct PostgresSslConfig {
+	pub ssl_root_cert_path: Option<PathBuf>,
+	pub ssl_client_cert_path: Option<PathBuf>,
+	pub ssl_client_key_path: Option<PathBuf>,
+}
+
+impl PostgresConfig {
+	/// Create a new PostgreSQL configuration with sane defaults
+	pub fn new(connection_string: String) -> Self {
+		Self {
+			connection_string,
+			ssl_config: None,
+		}
+	}
+}
+
 pub struct PostgresDatabaseDriver {
-	pool: Arc<Pool>,
+	pool: 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 a new PostgreSQL driver with custom configuration
+	pub async fn new_with_config(config: PostgresConfig) -> Result<Self> {
+		tracing::debug!(
+			connection_string = ?config.connection_string,
+			"creating PostgresDatabaseDriver"
+		);
 
 		// Create deadpool config from connection string
-		let mut config = Config::new();
-		config.url = Some(connection_string);
-		config.pool = Some(PoolConfig {
+		let mut pool_config = Config::new();
+		pool_config.url = Some(config.connection_string.clone());
+		pool_config.pool = Some(PoolConfig {
 			max_size: 64,
 			..Default::default()
 		});
-		config.manager = Some(ManagerConfig {
+		pool_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!("creating Postgres pool");
+
+		let ssl_disabled = if let Ok(url) = Url::parse(&config.connection_string) {
+			url.query_pairs()
+				.any(|(k, v)| k == "sslmode" && v == "disable")
+		} else {
+			false
+		};
+
+		let pool = if ssl_disabled {
+			let tls = tokio_postgres::NoTls;
+
+			pool_config
+				.create_pool(Some(Runtime::Tokio1), tls)
+				.context("failed to create postgres connection pool")?
+		} else {
+			let tls_config = build_tls_config(
+				config
+					.ssl_config
+					.as_ref()
+					.and_then(|c| c.ssl_root_cert_path.as_ref()),
+				config
+					.ssl_config
+					.as_ref()
+					.and_then(|c| c.ssl_client_cert_path.as_ref()),
+				config
+					.ssl_config
+					.as_ref()
+					.and_then(|c| c.ssl_client_key_path.as_ref()),
+			)?;
+			let tls = MakeRustlsConnect::new(tls_config);
+
+			pool_config
+				.create_pool(Some(Runtime::Tokio1), tls)
+				.context("failed to create postgres connection pool")?
+		};
 
 		tracing::debug!("Getting Postgres connection from pool");
 		// Get a connection from the pool to create the table
@@ -63,12 +124,9 @@ impl PostgresDatabaseDriver {
 			.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")?;
+		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(
@@ -121,9 +179,9 @@ impl PostgresDatabaseDriver {
 				EXCLUDE USING gist (
 					-- Conflict if byte range overlaps...
 					range_data WITH &&,
-					-- And f conflict types are different...
+					-- And if conflict types are different...
 					conflict_type WITH <>,
-					-- And f the txn versions overlap...
+					-- And if 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 <>
@@ -142,6 +200,7 @@ impl PostgresDatabaseDriver {
 		.await
 		.context("failed to create index on conflict_ranges ts column")?;
 
+		let pool2 = pool.clone();
 		let gc_handle = tokio::spawn(async move {
 			let mut interval = tokio::time::interval(GC_INTERVAL);
 			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
@@ -149,6 +208,8 @@ impl PostgresDatabaseDriver {
 			loop {
 				interval.tick().await;
 
+				tracing::debug!(status=?pool2.status(), "postgres pool status");
+
 				// 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
@@ -164,7 +225,7 @@ impl PostgresDatabaseDriver {
 		});
 
 		Ok(PostgresDatabaseDriver {
-			pool: Arc::new(pool),
+			pool,
 			max_retries: AtomicI32::new(100),
 			gc_handle,
 		})
@@ -173,10 +234,10 @@ impl PostgresDatabaseDriver {
 
 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(),
-		))))
+		// Pass the connection pool and config to the transaction driver
+		Ok(Transaction::new(Arc::new(
+			PostgresTransactionDriver::with_config(self.pool.clone()),
+		)))
 	}
 
 	fn run<'a>(

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

@@ -2,4 +2,4 @@ mod database;
 mod transaction;
 mod transaction_task;
 
-pub use database::PostgresDatabaseDriver;
+pub use database::{PostgresConfig, PostgresDatabaseDriver, PostgresSslConfig};

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

@@ -1,10 +1,7 @@
 use std::{
 	future::Future,
 	pin::Pin,
-	sync::{
-		Arc,
-		atomic::{AtomicBool, Ordering},
-	},
+	sync::atomic::{AtomicBool, Ordering},
 };
 
 use anyhow::{Context, Result};
@@ -24,14 +21,14 @@ use crate::{
 use super::transaction_task::{TransactionCommand, TransactionTask};
 
 pub struct PostgresTransactionDriver {
-	pool: Arc<Pool>,
+	pool: Pool,
 	operations: TransactionOperations,
 	committed: AtomicBool,
 	tx_sender: OnceCell<mpsc::Sender<TransactionCommand>>,
 }
 
 impl PostgresTransactionDriver {
-	pub fn new(pool: Arc<Pool>) -> Self {
+	pub fn with_config(pool: Pool) -> Self {
 		PostgresTransactionDriver {
 			pool,
 			operations: TransactionOperations::default(),
@@ -47,7 +44,7 @@ impl PostgresTransactionDriver {
 				let (sender, receiver) = mpsc::channel(100);
 
 				// Spawn the transaction task with serializable isolation
-				let task = TransactionTask::new(self.pool.as_ref().clone(), receiver);
+				let task = TransactionTask::new(self.pool.clone(), receiver);
 				tokio::spawn(task.run());
 
 				anyhow::Ok(sender)

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

@@ -78,6 +78,7 @@ impl TransactionTask {
 			}
 		};
 
+		// Start the read transaction
 		let tx = match conn
 			.build_transaction()
 			.isolation_level(IsolationLevel::RepeatableRead)
@@ -92,6 +93,7 @@ impl TransactionTask {
 			}
 		};
 
+		// TODO: Parallelize future
 		let start_version = match tx
 			.query_one("SELECT nextval('global_version_seq')", &[])
 			.await
@@ -154,9 +156,11 @@ impl TransactionTask {
 					conflict_ranges,
 					response,
 				} => {
-					let result = self
-						.handle_commit(tx, start_version, operations, conflict_ranges)
-						.await;
+					let (_, result) = tokio::join!(
+						// Read-only txn, we don't care about the result
+						tx.commit(),
+						self.handle_commit(start_version, operations, conflict_ranges),
+					);
 
 					let _ = response.send(result);
 					// Exit after commit
@@ -327,25 +331,20 @@ impl TransactionTask {
 
 	async fn handle_commit(
 		&mut self,
-		tx: Transaction<'_>,
 		start_version: i64,
 		operations: Vec<Operation>,
 		conflict_ranges: Vec<(Vec<u8>, Vec<u8>, ConflictRangeType)>,
 	) -> Result<()> {
-		// // 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')", &[]),
-		);
+		// Get connection from pool
+		let mut conn = self.pool.get().await?;
 
-		let commit_version = version_res
-			.context("failed to get postgres txn commit_version")?
-			.get::<_, i64>(0);
+		// Start write transaction
+		let tx = conn
+			.build_transaction()
+			.isolation_level(IsolationLevel::ReadCommitted)
+			.start()
+			.await
+			.context("failed to start write txn")?;
 
 		let mut begins = Vec::with_capacity(conflict_ranges.len());
 		let mut ends = Vec::with_capacity(conflict_ranges.len());
@@ -363,28 +362,22 @@ impl TransactionTask {
 		}
 
 		let query = "
+			WITH data AS (
+				SELECT nextval('global_version_seq') AS commit_version
+			)
 			INSERT INTO conflict_ranges (range_data, conflict_type, start_version, commit_version)
 			SELECT
 				bytearange(begin_key, end_key, '[)'),
 				conflict_type::range_type,
 				$4,
-				$5
-			FROM UNNEST($1::bytea[], $2::bytea[], $3::text[]) AS t(begin_key, end_key, conflict_type)";
+				data.commit_version
+			FROM UNNEST($1::bytea[], $2::bytea[], $3::text[]) AS t(begin_key, end_key, conflict_type), data";
 		let stmt = tx.prepare_cached(query).await.map_err(map_postgres_error)?;
 
 		// Insert all conflict ranges at once
-		tx.execute(
-			&stmt,
-			&[
-				&begins,
-				&ends,
-				&conflict_types,
-				&start_version,
-				&commit_version,
-			],
-		)
-		.await
-		.map_err(map_postgres_error)?;
+		tx.execute(&stmt, &[&begins, &ends, &conflict_types, &start_version])
+			.await
+			.map_err(map_postgres_error)?;
 
 		for op in operations {
 			match op {
@@ -392,16 +385,6 @@ impl TransactionTask {
 					// 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)?;
 
@@ -453,16 +436,6 @@ 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)
@@ -515,7 +488,12 @@ impl TransactionTask {
 
 /// Maps PostgreSQL error to DatabaseError
 fn map_postgres_error(err: tokio_postgres::Error) -> anyhow::Error {
-	let error_str = err.to_string();
+	let error_str = if let Some(err) = err.as_db_error() {
+		err.to_string()
+	} else {
+		err.to_string()
+	};
+
 	if error_str.contains("exclusion_violation")
 		|| error_str.contains("violates exclusion constraint")
 	{

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

@@ -1,20 +1,20 @@
 use std::{
-	path::PathBuf,
+	path::{Path, PathBuf},
 	sync::{
 		Arc,
 		atomic::{AtomicI32, Ordering},
 	},
-	time::Duration,
 };
 
 use anyhow::{Context, Result};
-use rocksdb::{OptimisticTransactionDB, Options};
+use rocksdb::{OptimisticTransactionDB, Options, checkpoint::Checkpoint};
 
 use crate::{
 	RetryableTransaction, Transaction,
 	driver::{BoxFut, DatabaseDriver, Erased},
 	error::DatabaseError,
 	options::DatabaseOption,
+	transaction::TXN_TIMEOUT,
 	utils::{MaybeCommitted, calculate_tx_retry_backoff},
 };
 
@@ -22,8 +22,6 @@ use super::{
 	transaction::RocksDbTransactionDriver, transaction_conflict_tracker::TransactionConflictTracker,
 };
 
-const TXN_TIMEOUT: Duration = Duration::from_secs(5);
-
 pub struct RocksDbDatabaseDriver {
 	db: Arc<OptimisticTransactionDB>,
 	max_retries: AtomicI32,
@@ -32,7 +30,7 @@ pub struct RocksDbDatabaseDriver {
 
 impl RocksDbDatabaseDriver {
 	pub async fn new(db_path: PathBuf) -> Result<Self> {
-		tracing::info!(?db_path, "starting file system driver");
+		tracing::info!(db_path=%db_path.display(), "starting file system driver");
 
 		// Create directory if it doesn't exist
 		std::fs::create_dir_all(&db_path).context("failed to create database directory")?;
@@ -42,7 +40,8 @@ impl RocksDbDatabaseDriver {
 		opts.create_if_missing(true);
 		opts.set_max_open_files(10000);
 		opts.set_keep_log_file_num(10);
-		opts.set_max_total_wal_size(64 * 1024 * 1024); // 64MB
+		opts.set_max_total_wal_size(64 * 1024 * 1024); // 64MiB
+		opts.set_write_buffer_size(256 * 1024 * 1024); // 256MiB for conflict detection
 
 		// Open the OptimisticTransactionDB
 		tracing::debug!(path=%db_path.display(), "opening rocksdb");
@@ -120,6 +119,13 @@ impl DatabaseDriver for RocksDbDatabaseDriver {
 			}
 		}
 	}
+
+	fn checkpoint(&self, path: &Path) -> Result<()> {
+		let cp = Checkpoint::new(&*self.db).context("failed to create checkpoint handle")?;
+		cp.create_checkpoint(path)
+			.context("failed to create rocksdb checkpoint")?;
+		Ok(())
+	}
 }
 
 impl Drop for RocksDbDatabaseDriver {

package/packages/universaldb/src/driver/rocksdb/transaction_conflict_tracker.rs

@@ -64,12 +64,12 @@ impl TransactionConflictTracker {
 					for (cr2_start, cr2_end, cr2_type) in &txn2.conflict_ranges {
 						// Check conflict ranges overlap
 						if cr1_start < cr2_end && cr2_start < cr1_end && cr1_type != cr2_type {
-							tracing::info!(
-								?cr1_start,
-								?cr1_end,
+							tracing::debug!(
+								cr1_start=%hex::encode(cr1_start),
+								cr1_end=%hex::encode(cr1_end),
 								?cr1_type,
-								?cr2_start,
-								?cr2_end,
+								cr2_start=%hex::encode(cr2_start),
+								cr2_end=%hex::encode(cr2_end),
 								?cr2_type,
 								txn1_start_version,
 								txn1_commit_version,

package/packages/universaldb/src/driver/rocksdb/transaction_task.rs

@@ -139,7 +139,8 @@ impl TransactionTask {
 
 	fn create_transaction(&self) -> RocksDbTransaction<'_, OptimisticTransactionDB> {
 		let write_opts = WriteOptions::default();
-		let txn_opts = rocksdb::OptimisticTransactionOptions::default();
+		let mut txn_opts = rocksdb::OptimisticTransactionOptions::default();
+		txn_opts.set_snapshot(true);
 		self.db.transaction_opt(&write_opts, &txn_opts)
 	}
 

package/packages/universaldb/src/metrics.rs

@@ -1,27 +1,43 @@
-use rivet_metrics::{
-	MICRO_BUCKETS,
-	otel::{global::*, metrics::*},
-};
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-universaldb");
+	pub static ref KEY_PACK_COUNT: IntCounterVec = register_int_counter_vec_with_registry!(
+		"udb_key_pack_count",
+		"How many times a key has been packed.",
+		&["type"],
+		*REGISTRY
+	).unwrap();
+	pub static ref KEY_UNPACK_COUNT: IntCounterVec = register_int_counter_vec_with_registry!(
+		"udb_key_unpack_count",
+		"How many times a key has been unpacked.",
+		&["type"],
+		*REGISTRY
+	).unwrap();
 
-	/// Has no expected attributes
-	pub static ref PING_DURATION: Histogram<f64> = METER.f64_histogram("rivet_udb_ping_duration")
-		.with_description("Total duration to retrieve a single value from udb.")
-		.with_boundaries(MICRO_BUCKETS.to_vec())
-		.build();
-	/// Has no expected attributes
-	pub static ref MISSED_PING: Gauge<u64> = METER.u64_gauge("rivet_udb_missed_ping")
-		.with_description("1 if udb missed the last ping.")
-		.build();
-
-	/// Expected attributes: "type"
-	pub static ref KEY_PACK_COUNT: Counter<u64> = METER.u64_counter("rivet_udb_key_pack_count")
-		.with_description("How many times a key has been packed.")
-		.build();
-	/// Expected attributes: "type"
-	pub static ref KEY_UNPACK_COUNT: Counter<u64> = METER.u64_counter("rivet_udb_key_unpack_count")
-		.with_description("How many times a key has been unpacked.")
-		.build();
+	pub static ref TRANSACTION_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"udb_transaction_total",
+		"How many transactions have been started.",
+		&["name"],
+		*REGISTRY
+	).unwrap();
+	pub static ref TRANSACTION_PENDING: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"udb_transaction_pending",
+		"How many transactions have been started.",
+		&["name"],
+		*REGISTRY
+	).unwrap();
+	pub static ref TRANSACTION_DURATION: HistogramVec = register_histogram_vec_with_registry!(
+		"udb_transaction_duration",
+		"Duration of a transaction.",
+		&["name"],
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
+	pub static ref TRANSACTION_ATTEMPTS: HistogramVec = register_histogram_vec_with_registry!(
+		"udb_transaction_attempts",
+		"Amount of attempts (1 + retries) taken for a transaction.",
+		&["name"],
+		vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0, 16.0],
+		*REGISTRY
+	).unwrap();
 }

package/packages/universaldb/src/prelude.rs

@@ -1,6 +1,6 @@
 pub use crate::{
 	key_selector::KeySelector,
-	options::StreamingMode,
+	options::{ConflictRangeType, MutationType, StreamingMode},
 	range_option::RangeOption,
 	tuple::{PackError, PackResult, TupleDepth, TuplePack, TupleUnpack, VersionstampOffset},
 	utils::{FormalChunkedKey, FormalKey, IsolationLevel::*, OptSliceExt, SliceExt, keys::*},

package/packages/universaldb/src/transaction.rs

@@ -1,4 +1,4 @@
-use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
+use std::{future::Future, ops::Deref, pin::Pin, sync::Arc, time::Duration};
 
 use anyhow::{Context, Result};
 
@@ -15,6 +15,8 @@ use crate::{
 	value::{Slice, Value, Values},
 };
 
+pub const TXN_TIMEOUT: Duration = Duration::from_secs(5);
+
 #[derive(Clone)]
 pub struct Transaction {
 	pub(crate) driver: Arc<dyn TransactionDriver>,
@@ -111,7 +113,7 @@ impl Transaction {
 
 	pub fn delete_key_subspace<T: TuplePack>(&self, key: &T) {
 		self.informal()
-			.clear_subspace_range(&self.subspace.subspace(&self.subspace.pack(key)));
+			.clear_subspace_range(&self.subspace.subspace(&key));
 	}
 
 	pub fn read_entry<T: FormalKey + for<'de> TupleUnpack<'de>>(
@@ -261,6 +263,11 @@ impl Transaction {
 	) -> Pin<Box<dyn Future<Output = Result<i64>> + Send + 'a>> {
 		self.driver.get_estimated_range_size_bytes(begin, end)
 	}
+
+	/// Adds a tag to the current transaction
+	pub fn tag(&self, tag: &str) -> Result<()> {
+		self.driver.tag(tag)
+	}
 }
 
 pub struct InformalTransaction<'t> {