fix: add max param to prune signals cmd

Author: MasterPtato

engine/docker/template/grafana-dashboards/gasoline.json

@@ -232,7 +232,7 @@
             "uid": "prometheus"
           },
           "editorMode": "code",
-          "expr": "sum by (workflow_name) (\n    max by(workflow_name, rivet_project, rivet_datacenter) (\n        rivet_gasoline_workflow_total{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\",workflow_name=~\"$workflow_name\"}\n        * on(k8s_pod_name) group_left(rivet_project, rivet_datacenter)\n        (\n            (rivet_gasoline_worker_last_metrics_publish{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\"} == bool\n            on(rivet_project, rivet_datacenter) group_left()\n            max by (rivet_project, rivet_datacenter) (rivet_gasoline_worker_last_metrics_publish{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\"}))\n        )\n    )\n)",
+          "expr": "sum by (workflow_name) (\n    max by(workflow_name, rivet_project, rivet_datacenter) (\n        rivet_gasoline_workflow_sleeping{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\",workflow_name=~\"$workflow_name\"}\n        * on(k8s_pod_name) group_left(rivet_project, rivet_datacenter)\n        (\n            (rivet_gasoline_worker_last_metrics_publish{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\"} == bool\n            on(rivet_project, rivet_datacenter) group_left()\n            max by (rivet_project, rivet_datacenter) (rivet_gasoline_worker_last_metrics_publish{rivet_project=~\"$project\",rivet_datacenter=~\"$datacenter\"}))\n        )\n    )\n)",
           "instant": false,
           "legendFormat": "{{workflow_name}}",
           "range": true,

engine/packages/engine/src/commands/wf/mod.rs

@@ -42,7 +42,7 @@ pub enum SubCommand {
 		#[clap(short = 'd', long)]
 		dry_run: bool,
 		#[clap(short = 'p', long)]
-		parallelization: Option<u128>,
+		parallelization: Option<u16>,
 	},
 	/// Deletes the history for completed workflows that match the name and before filter.
 	PruneHistory {
@@ -53,7 +53,7 @@ pub enum SubCommand {
 		#[clap(short = 'd', long)]
 		dry_run: bool,
 		#[clap(short = 'p', long)]
-		parallelization: Option<u128>,
+		parallelization: Option<u16>,
 	},
 	/// Lists the entire event history of a workflow.
 	History {

engine/packages/engine/src/commands/wf/signal.rs

@@ -36,7 +36,9 @@ pub enum SubCommand {
 		#[clap(short = 'd', long)]
 		dry_run: bool,
 		#[clap(short = 'p', long)]
-		parallelization: Option<u128>,
+		parallelization: Option<u16>,
+		#[clap(short = 'm', long)]
+		max_per_txn: Option<usize>,
 	},
 }
 
@@ -73,13 +75,15 @@ impl SubCommand {
 				before,
 				dry_run,
 				parallelization,
+				max_per_txn,
 			} => {
 				let total = db
 					.prune_acked_signals(
 						&name.iter().map(|x| x.as_str()).collect::<Vec<_>>(),
 						before.timestamp_millis(),
 						dry_run,
 						parallelization.unwrap_or(1),
+						max_per_txn,
 					)
 					.await?;
 

engine/packages/gasoline/src/db/debug.rs

@@ -45,7 +45,7 @@ pub trait DatabaseDebug: Database {
 		names: &[&str],
 		error_like: &[&str],
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
 	) -> Result<usize>;
 
 	/// Used by pruner workflow for automatic pruning.
@@ -57,7 +57,7 @@ pub trait DatabaseDebug: Database {
 		names: &[&str],
 		before_ts: i64,
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
 	) -> Result<usize>;
 
 	/// Used for manual pruning.
@@ -66,7 +66,8 @@ pub trait DatabaseDebug: Database {
 		names: &[&str],
 		before_ts: i64,
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
+		max_per_txn: Option<usize>,
 	) -> Result<usize>;
 }
 

engine/packages/gasoline/src/db/kv/debug.rs

@@ -1221,15 +1221,15 @@ impl DatabaseDebug for DatabaseKv {
 		names: &[&str],
 		error_like: &[&str],
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
 	) -> Result<usize> {
 		ensure!(parallelization > 0);
 		ensure!(parallelization < 1024);
 
-		let chunk_size = u128::MAX / parallelization;
+		let chunk_size = u128::MAX / parallelization as u128;
 		let mut futs = FuturesUnordered::new();
 
-		for i in 0..parallelization {
+		for i in 0..parallelization as u128 {
 			let start = i * chunk_size;
 			futs.push(self.revive_workflows_inner(
 				names,
@@ -1407,15 +1407,15 @@ impl DatabaseDebug for DatabaseKv {
 		names: &[&str],
 		before_ts: i64,
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
 	) -> Result<usize> {
 		ensure!(parallelization > 0);
 		ensure!(parallelization < 1024);
 
-		let chunk_size = u128::MAX / parallelization;
+		let chunk_size = u128::MAX / parallelization as u128;
 		let mut futs = FuturesUnordered::new();
 
-		for i in 0..parallelization {
+		for i in 0..parallelization as u128 {
 			let start = i * chunk_size;
 			futs.push(self.prune_workflow_history_inner(
 				names,
@@ -1442,22 +1442,24 @@ impl DatabaseDebug for DatabaseKv {
 		names: &[&str],
 		before_ts: i64,
 		dry_run: bool,
-		parallelization: u128,
+		parallelization: u16,
+		max_per_txn: Option<usize>,
 	) -> Result<usize> {
 		ensure!(parallelization > 0);
 		ensure!(parallelization < 1024);
 
-		let chunk_size = u128::MAX / parallelization;
+		let chunk_size = u128::MAX / parallelization as u128;
 		let mut futs = FuturesUnordered::new();
 
-		for i in 0..parallelization {
+		for i in 0..parallelization as u128 {
 			let start = i * chunk_size;
 			futs.push(self.prune_signals_inner(
 				names,
 				before_ts,
 				dry_run,
 				start,
 				start + chunk_size,
+				max_per_txn,
 			));
 		}
 
@@ -1813,6 +1815,7 @@ impl DatabaseKv {
 		dry_run: bool,
 		start: u128,
 		end: u128,
+		max_per_txn: Option<usize>,
 	) -> Result<usize> {
 		let mut total = 0;
 		let mut current_signal_id = Some(Id::v1(Uuid::from_u128(start), self.config.dc_label()));
@@ -1858,14 +1861,24 @@ impl DatabaseKv {
 						let mut ts_matches = false;
 
 						let fut = async {
-							while let Some(entry) = stream.try_next().await? {
+							loop {
+								let Some(entry) = stream.try_next().await? else {
+									break;
+								};
+
 								let signal_id = *self.subspace.unpack::<JustId>(entry.key())?;
 
 								if let Some(curr) = current_signal_id {
 									if signal_id != curr {
 										// Save if matches query
 										if name_matches && state_matches && ts_matches {
 											signal_ids.push(curr);
+
+											if let Some(max_per_txn) = max_per_txn {
+												if signal_ids.len() >= max_per_txn {
+													return anyhow::Ok(());
+												}
+											}
 										}
 
 										signals_processed += 1;