@rivetkit/engine 1.0.0

package/packages/universaldb/tests/integration_gas.rs

@@ -0,0 +1,264 @@
+use futures_util::TryStreamExt;
+use rivet_test_deps_docker::TestDatabase;
+use std::sync::Arc;
+use universaldb::{
+	Database, RangeOption, options::StreamingMode, tuple::Subspace, utils::IsolationLevel::*,
+};
+use uuid::Uuid;
+
+#[tokio::test]
+async fn test_postgres_gasoline() {
+	let _ = tracing_subscriber::fmt::try_init();
+
+	let (db_config, docker_config) = TestDatabase::Postgres
+		.config(Uuid::new_v4(), 1)
+		.await
+		.unwrap();
+	let mut docker_config = docker_config.unwrap();
+	docker_config.start().await.unwrap();
+
+	tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
+
+	let rivet_config::config::Database::Postgres(postgres_config) = db_config else {
+		unreachable!();
+	};
+
+	// Get the connection string from the secret
+	let connection_string = postgres_config.url.read().clone();
+
+	let driver = universaldb::driver::PostgresDatabaseDriver::new(connection_string)
+		.await
+		.unwrap();
+	let db = Database::new(Arc::new(driver));
+
+	test_gasoline_operations(&db).await;
+}
+
+#[tokio::test]
+async fn test_rocksdb_gasoline() {
+	// Create a unique temporary directory for this test
+	let temp_dir = tempfile::tempdir().unwrap();
+	let db_path = temp_dir.path();
+
+	let driver = universaldb::driver::RocksDbDatabaseDriver::new(db_path.to_path_buf())
+		.await
+		.unwrap();
+	let db = Database::new(Arc::new(driver));
+
+	test_gasoline_operations(&db).await;
+}
+
+pub async fn test_gasoline_operations(db: &Database) {
+	// First, test a simple write and read to isolate the issue
+	println!("Running simple write/read test...");
+
+	// Simple test: write a single key and read it back
+	db.run(|tx| async move {
+		tx.set(b"simple_test_key", b"simple_test_value");
+		Ok(())
+	})
+	.await
+	.unwrap();
+
+	let value = db
+		.run(|tx| async move {
+			let val = tx.get(b"simple_test_key", Serializable).await?;
+			println!(
+				"Simple test read result: {:?}",
+				val.as_ref()
+					.map(|v| std::str::from_utf8(v).unwrap_or("<binary>"))
+			);
+			Ok(val)
+		})
+		.await
+		.unwrap();
+
+	assert!(
+		value.is_some(),
+		"Simple write/read test failed - value not found!"
+	);
+	println!("Simple write/read test passed");
+
+	// Define the gasoline subspace
+	const RIVET: usize = 0;
+	const GASOLINE: usize = 1;
+	const KV: usize = 2;
+	let workflow_subspace = Subspace::from(&(RIVET, GASOLINE, KV));
+
+	// Generate a workflow ID
+	let workflow_id = Uuid::new_v4();
+
+	// Test 1: Write workflow data like gasoline does
+	db.run(|tx| {
+		let workflow_subspace = workflow_subspace.clone();
+		async move {
+			// Write create timestamp (similar to CreateTsKey)
+			let create_ts_key = workflow_subspace.pack(&("workflow", "create_ts", workflow_id));
+			tx.set(&create_ts_key, b"12345678");
+
+			// Write chunked input data (similar to InputKey)
+			let input_data = r#"{"test": "data"}"#;
+			let input_bytes = input_data.as_bytes();
+
+			// Chunk the data (mimicking FormalChunkedKey behavior)
+			const CHUNK_SIZE: usize = 1024;
+			let chunks: Vec<_> = input_bytes.chunks(CHUNK_SIZE).collect();
+
+			for (i, chunk) in chunks.iter().enumerate() {
+				let chunk_key =
+					workflow_subspace.pack(&("workflow", "input", workflow_id, i as u32));
+				tx.set(&chunk_key, chunk);
+			}
+
+			// Write state data
+			let state_data = r#"{}"#;
+			let state_key = workflow_subspace.pack(&("workflow", "state", workflow_id, 0u32));
+			tx.set(&state_key, state_data.as_bytes());
+
+			// Write has_wake_condition
+			let wake_condition_key =
+				workflow_subspace.pack(&("workflow", "has_wake_condition", workflow_id));
+			tx.set(&wake_condition_key, b"false");
+
+			Ok(())
+		}
+	})
+	.await
+	.unwrap();
+
+	println!("Test 1 completed - data written");
+
+	// Test 2: Read workflow data back like gasoline does
+	let (input_found, state_found, wake_found) = db
+		.run(|tx| {
+			let workflow_subspace = workflow_subspace.clone();
+			async move {
+				// Read input chunks using range query
+				let input_key_base = workflow_subspace.pack(&("workflow", "input", workflow_id));
+				let input_subspace = Subspace::from_bytes(input_key_base.clone());
+
+				let input_chunks = tx
+					.get_ranges_keyvalues(
+						RangeOption {
+							mode: StreamingMode::WantAll,
+							..(&input_subspace).into()
+						},
+						Serializable,
+					)
+					.try_collect::<Vec<_>>()
+					.await?;
+
+				// Read state chunks
+				let state_key_base = workflow_subspace.pack(&("workflow", "state", workflow_id));
+				let state_subspace = Subspace::from_bytes(state_key_base.clone());
+
+				let state_chunks = tx
+					.get_ranges_keyvalues(
+						RangeOption {
+							mode: StreamingMode::WantAll,
+							..(&state_subspace).into()
+						},
+						Serializable,
+					)
+					.try_collect::<Vec<_>>()
+					.await?;
+
+				// Read wake condition
+				let wake_condition_key =
+					workflow_subspace.pack(&("workflow", "has_wake_condition", workflow_id));
+				let wake_condition = tx.get(&wake_condition_key, Serializable).await?;
+
+				println!("Input chunks found: {}", input_chunks.len());
+				println!("State chunks found: {}", state_chunks.len());
+				println!("Wake condition found: {}", wake_condition.is_some());
+
+				Ok((
+					!input_chunks.is_empty(),
+					!state_chunks.is_empty(),
+					wake_condition.is_some(),
+				))
+			}
+		})
+		.await
+		.unwrap();
+
+	assert!(input_found, "Should find input chunks");
+	assert!(state_found, "Should find state chunks");
+	assert!(wake_found, "Should find wake condition");
+
+	// Test 3: Test the exact pattern gasoline uses with subspace operations
+	db.run(|tx| {
+		let workflow_subspace = workflow_subspace.clone();
+		async move {
+			// Create a new workflow ID
+			let workflow_id2 = Uuid::new_v4();
+
+			// Write using the exact subspace pattern gasoline uses
+			let input_key_base = workflow_subspace.pack(&("workflow", "input", workflow_id2));
+			let input_subspace = Subspace::from_bytes(input_key_base.clone());
+
+			// Write chunked data to the input subspace
+			let test_input = r#"{"action": "test_workflow"}"#;
+			let chunk_key = input_subspace.pack(&(0u32,));
+			tx.set(&chunk_key, test_input.as_bytes());
+
+			Ok(())
+		}
+	})
+	.await
+	.unwrap();
+
+	// Test 4: Verify the data was written correctly
+	let workflow_id2 = db
+		.run(|tx| {
+			let workflow_subspace = workflow_subspace.clone();
+			async move {
+				// Generate the same workflow_id2 again (for test purposes, we'll store it)
+				// In real usage, we'd pass it between transactions
+				let workflow_id2 = Uuid::new_v4();
+
+				// Write and immediately read back to verify
+				let input_key_base = workflow_subspace.pack(&("workflow", "input", workflow_id2));
+				let input_subspace = Subspace::from_bytes(input_key_base.clone());
+				let chunk_key = input_subspace.pack(&(0u32,));
+
+				tx.set(&chunk_key, b"test_data");
+
+				// Read it back in the same transaction
+				let value = tx.get(&chunk_key, Serializable).await?;
+				assert_eq!(
+					value,
+					Some(b"test_data".to_vec().into()),
+					"Should read back the same data"
+				);
+
+				Ok(workflow_id2)
+			}
+		})
+		.await
+		.unwrap();
+
+	// Test 5: Read in a separate transaction (like gasoline does)
+	db.run(|tx| {
+		let workflow_subspace = workflow_subspace.clone();
+		async move {
+			let input_key_base = workflow_subspace.pack(&("workflow", "input", workflow_id2));
+			let input_subspace = Subspace::from_bytes(input_key_base.clone());
+
+			// Read using range query like gasoline
+			let input_chunks = tx
+				.get_ranges_keyvalues((&input_subspace).into(), Serializable)
+				.try_collect::<Vec<_>>()
+				.await?;
+
+			assert!(
+				!input_chunks.is_empty(),
+				"Should find input chunks in separate transaction"
+			);
+
+			Ok(())
+		}
+	})
+	.await
+	.unwrap();
+}

package/packages/universaldb/tests/rocksdb.rs

@@ -0,0 +1,141 @@
+use std::{path::Path, sync::Arc};
+
+use anyhow::Context;
+use futures_util::StreamExt;
+use rivet_test_deps_docker::TestDatabase;
+use rocksdb::{OptimisticTransactionDB, Options, WriteOptions};
+use universaldb::{
+	Database,
+	utils::{IsolationLevel::*, calculate_tx_retry_backoff},
+};
+use uuid::Uuid;
+
+#[tokio::test]
+async fn rocksdb_native() {
+	let _ = tracing_subscriber::fmt()
+		.with_env_filter("debug")
+		.with_test_writer()
+		.try_init();
+
+	let db_path = Path::new("/tmp/foobar-db");
+	std::fs::create_dir_all(&db_path)
+		.context("failed to create database directory")
+		.unwrap();
+
+	// Configure RocksDB options
+	let mut opts = Options::default();
+	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
+
+	// Open the OptimisticTransactionDB
+	tracing::debug!(path=%db_path.display(), "opening rocksdb");
+	let db: Arc<OptimisticTransactionDB> = Arc::new(
+		OptimisticTransactionDB::open(&opts, db_path)
+			.context("failed to open rocksdb")
+			.unwrap(),
+	);
+
+	let mut handles = Vec::new();
+
+	for _ in 0..64 {
+		let db = db.clone();
+		let write_opts = WriteOptions::default();
+		let txn_opts = rocksdb::OptimisticTransactionOptions::default();
+
+		let handle = tokio::spawn(async move {
+			for attempt in 0..300 {
+				let txn = db.transaction_opt(&write_opts, &txn_opts);
+
+				let key = vec![1, 2, 3];
+				let value = vec![4, 5, 6];
+
+				txn.get(&key).unwrap();
+				txn.put(&key, &value).unwrap();
+
+				tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+				// Execute transaction
+				let Err(error) = txn.commit() else {
+					tracing::info!("success");
+					return;
+				};
+
+				let err_str = error.to_string();
+				if err_str.contains("conflict") || err_str.contains("Resource busy") {
+					tracing::warn!(?error, "conflict");
+					let backoff_ms = calculate_tx_retry_backoff(attempt);
+					tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await;
+					continue;
+				}
+
+				Err::<(), _>(error).unwrap();
+			}
+		});
+
+		handles.push(handle);
+	}
+
+	futures_util::stream::iter(handles)
+		.buffer_unordered(64)
+		.for_each(|result| async {
+			if let Err(err) = result {
+				tracing::error!(?err, "task failed");
+			}
+		})
+		.await;
+}
+
+#[tokio::test]
+async fn rocksdb_udb() {
+	let _ = tracing_subscriber::fmt()
+		.with_env_filter("debug")
+		.with_test_writer()
+		.try_init();
+
+	let test_id = Uuid::new_v4();
+	let (db_config, _docker_config) = TestDatabase::FileSystem.config(test_id, 1).await.unwrap();
+
+	let rivet_config::config::Database::FileSystem(fs_config) = db_config else {
+		unreachable!()
+	};
+
+	let driver = universaldb::driver::RocksDbDatabaseDriver::new(fs_config.path)
+		.await
+		.unwrap();
+	let db = Database::new(Arc::new(driver));
+
+	let mut handles = Vec::new();
+
+	for _ in 0..64 {
+		let db = db.clone();
+
+		let handle = tokio::spawn(async move {
+			db.run(|tx| async move {
+				let key = vec![1, 2, 3];
+				let value = vec![4, 5, 6];
+
+				tx.get(&key, Serializable).await.unwrap();
+				tx.set(&key, &value);
+
+				Ok(())
+			})
+			.await
+			.unwrap();
+
+			tracing::info!("success");
+		});
+
+		handles.push(handle);
+	}
+
+	futures_util::stream::iter(handles)
+		.buffer_unordered(1024)
+		.for_each(|result| async {
+			if let Err(err) = result {
+				tracing::error!(?err, "task failed");
+			}
+		})
+		.await;
+}

package/packages/universaldb/tests/versionstamp.rs

@@ -0,0 +1,184 @@
+use universaldb::{
+	tuple::{Element, Versionstamp, pack_with_versionstamp, unpack},
+	versionstamp::*,
+};
+
+#[test]
+fn test_generate_versionstamp() {
+	let vs1 = generate_versionstamp(100);
+	let vs2 = generate_versionstamp(200);
+
+	assert!(vs1.is_complete());
+	assert!(vs2.is_complete());
+	assert_eq!(vs1.user_version(), 100);
+	assert_eq!(vs2.user_version(), 200);
+
+	assert_ne!(vs1.as_bytes(), vs2.as_bytes());
+}
+
+#[test]
+fn test_substitute_versionstamp_success() {
+	let incomplete = Versionstamp::from([0xff; 12]);
+	let tuple = vec![
+		Element::String("mykey".into()),
+		Element::Versionstamp(incomplete),
+		Element::Int(42),
+	];
+
+	let mut packed = pack_with_versionstamp(&tuple);
+	let versionstamp = generate_versionstamp(100);
+
+	assert!(substitute_versionstamp(&mut packed, versionstamp).is_ok());
+
+	let unpacked: Vec<Element> = unpack(&packed).unwrap();
+	assert_eq!(unpacked.len(), 3);
+
+	match &unpacked[1] {
+		Element::Versionstamp(v) => {
+			assert!(v.is_complete());
+			assert_eq!(v.user_version(), 100);
+		}
+		_ => panic!("Expected versionstamp"),
+	}
+}
+
+#[test]
+fn test_substitute_versionstamp_no_offset() {
+	let mut packed = vec![1, 2, 3];
+	let versionstamp = generate_versionstamp(100);
+
+	let result = substitute_versionstamp(&mut packed, versionstamp);
+	assert!(result.is_err());
+	assert!(result.unwrap_err().contains("too short"));
+}
+
+#[test]
+fn test_substitute_versionstamp_invalid_offset() {
+	let mut packed = vec![1, 2, 3, 4, 5];
+	packed.extend_from_slice(&100u32.to_le_bytes());
+
+	let versionstamp = generate_versionstamp(100);
+
+	let result = substitute_versionstamp(&mut packed, versionstamp);
+	assert!(result.is_err());
+	assert!(result.unwrap_err().contains("Invalid versionstamp offset"));
+}
+
+#[test]
+fn test_substitute_versionstamp_no_marker() {
+	let mut packed = vec![1, 2, 3, 4, 5, 6, 7, 8];
+	packed.extend_from_slice(&2u32.to_le_bytes());
+
+	let versionstamp = generate_versionstamp(100);
+
+	let result = substitute_versionstamp(&mut packed, versionstamp);
+	assert!(result.is_err());
+	assert!(result.unwrap_err().contains("No versionstamp marker"));
+}
+
+#[test]
+fn test_substitute_versionstamp_already_complete() {
+	// Create an incomplete versionstamp first
+	let incomplete = Versionstamp::from([0xff; 12]);
+	let tuple = vec![Element::Versionstamp(incomplete)];
+
+	let mut packed = pack_with_versionstamp(&tuple);
+
+	// First substitution - this should succeed
+	let versionstamp1 = generate_versionstamp(50);
+	assert!(substitute_versionstamp(&mut packed, versionstamp1).is_ok());
+
+	// Now try to substitute again on the already complete versionstamp
+	// We need to manually add the offset back
+	packed.extend_from_slice(&1u32.to_le_bytes());
+
+	let versionstamp2 = generate_versionstamp(100);
+
+	let result = substitute_versionstamp(&mut packed, versionstamp2);
+	// Should succeed but not modify the already complete versionstamp
+	assert!(result.is_ok());
+}
+
+#[test]
+fn test_pack_and_substitute_versionstamp() {
+	let incomplete = Versionstamp::from([0xff; 12]);
+	let tuple = vec![
+		Element::String("mykey".into()),
+		Element::Versionstamp(incomplete),
+		Element::Int(42),
+	];
+
+	let packed = pack_and_substitute_versionstamp(&tuple, 100).unwrap();
+
+	let unpacked: Vec<Element> = unpack(&packed).unwrap();
+	match &unpacked[1] {
+		Element::Versionstamp(v) => {
+			assert!(v.is_complete());
+			assert_eq!(v.user_version(), 100);
+		}
+		_ => panic!("Expected versionstamp"),
+	}
+}
+
+#[test]
+fn test_versionstamp_incomplete_preserves_user_version() {
+	// Test that Versionstamp::incomplete(user_version) properly preserves the user version
+	let user_version: u16 = 12345;
+	let incomplete = Versionstamp::incomplete(user_version);
+
+	// The versionstamp should be incomplete
+	assert!(
+		!incomplete.is_complete(),
+		"Versionstamp should be incomplete"
+	);
+
+	// The user version should be preserved
+	assert_eq!(
+		incomplete.user_version(),
+		user_version,
+		"User version should be preserved"
+	);
+
+	// Verify the bytes structure: first 10 bytes should be 0xff, last 2 bytes should be user_version
+	let bytes = incomplete.as_bytes();
+	assert_eq!(bytes.len(), 12);
+
+	// First 10 bytes should all be 0xff
+	for i in 0..10 {
+		assert_eq!(bytes[i], 0xff, "Byte {} should be 0xff", i);
+	}
+
+	// Last 2 bytes should contain the user version in big-endian
+	let stored_user_version = u16::from_be_bytes([bytes[10], bytes[11]]);
+	assert_eq!(
+		stored_user_version, user_version,
+		"User version bytes should match"
+	);
+
+	// Test with substitute_versionstamp - the user version from the incomplete versionstamp
+	// should be ignored and replaced with the one from the generated versionstamp
+	let tuple = vec![
+		Element::String("test".into()),
+		Element::Versionstamp(incomplete),
+	];
+
+	let mut packed = pack_with_versionstamp(&tuple);
+	let new_user_version: u16 = 54321;
+	let versionstamp = generate_versionstamp(new_user_version);
+
+	assert!(substitute_versionstamp(&mut packed, versionstamp).is_ok());
+
+	let unpacked: Vec<Element> = unpack(&packed).unwrap();
+	match &unpacked[1] {
+		Element::Versionstamp(v) => {
+			assert!(v.is_complete());
+			// The user version should be from the generated versionstamp, not the original incomplete one
+			assert_eq!(
+				v.user_version(),
+				new_user_version,
+				"Substituted versionstamp should have the new user version"
+			);
+		}
+		_ => panic!("Expected versionstamp"),
+	}
+}

package/packages/universalpubsub/Cargo.toml

@@ -0,0 +1,38 @@
+[package]
+name = "universalpubsub"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+async-nats.workspace = true
+async-trait.workspace = true
+base64.workspace = true
+deadpool-postgres.workspace = true
+futures-util.workspace = true
+rivet-error.workspace = true
+rivet-ups-protocol.workspace = true
+rivet-util.workspace = true
+serde_json.workspace = true
+vbare.workspace = true
+serde.workspace = true
+sha2.workspace = true
+tokio-postgres.workspace = true
+tokio.workspace = true
+tokio-util.workspace = true
+tracing.workspace = true
+uuid.workspace = true
+
+[dev-dependencies]
+rivet-config.workspace = true
+rivet-env.workspace = true
+tabled.workspace = true
+tempfile.workspace = true
+tracing-subscriber.workspace = true
+rivet-test-deps-docker.workspace = true
+
+[[bench]]
+name = "simple"
+harness = false