fix: clear all pending signals after workflow complete

Author: MasterPtato

Cargo.lock

@@ -804,6 +804,10 @@ impl PendingSignalKey {
 	pub fn subspace(workflow_id: Id, signal_name: String) -> PendingSignalSubspaceKey {
 		PendingSignalSubspaceKey::new(workflow_id, signal_name)
 	}
+
+	pub fn workflow_subspace(workflow_id: Id) -> PendingSignalSubspaceKey {
+		PendingSignalSubspaceKey::workflow(workflow_id)
+	}
 }
 
 impl FormalKey for PendingSignalKey {
@@ -854,14 +858,21 @@ impl<'de> TupleUnpack<'de> for PendingSignalKey {
 
 pub struct PendingSignalSubspaceKey {
 	workflow_id: Id,
-	signal_name: String,
+	signal_name: Option<String>,
 }
 
 impl PendingSignalSubspaceKey {
 	pub fn new(workflow_id: Id, signal_name: String) -> Self {
 		PendingSignalSubspaceKey {
 			workflow_id,
-			signal_name,
+			signal_name: Some(signal_name),
+		}
+	}
+
+	pub fn workflow(workflow_id: Id) -> Self {
+		PendingSignalSubspaceKey {
+			workflow_id,
+			signal_name: None,
 		}
 	}
 }
@@ -872,14 +883,16 @@ impl TuplePack for PendingSignalSubspaceKey {
 		w: &mut W,
 		tuple_depth: TupleDepth,
 	) -> std::io::Result<VersionstampOffset> {
-		let t = (
-			WORKFLOW,
-			SIGNAL,
-			self.workflow_id,
-			PENDING,
-			&self.signal_name,
-		);
-		t.pack(w, tuple_depth)
+		let mut offset = VersionstampOffset::None { size: 0 };
+
+		let t = (WORKFLOW, SIGNAL, self.workflow_id, PENDING);
+		offset += t.pack(w, tuple_depth)?;
+
+		if let Some(signal_name) = &self.signal_name {
+			offset += signal_name.pack(w, tuple_depth)?;
+		}
+
+		Ok(offset)
 	}
 }
 

engine/packages/gasoline/src/db/kv/keys/workflow.rs

@@ -1719,12 +1719,15 @@ impl Database for DatabaseKv {
 	) -> WorkflowResult<()> {
 		let start_instant = Instant::now();
 
-		let wrote_to_wake_idx = self
+		let (wrote_to_wake_idx, pending_signal_cleared_count) = self
 			.pools
 			.udb()
 			.map_err(WorkflowError::PoolsGeneric)?
 			.run(|tx| {
 				async move {
+					let start = Instant::now();
+					let tx = tx.with_subspace(self.subspace.clone());
+
 					let sub_workflow_wake_subspace = self
 						.subspace
 						.subspace(&keys::wake::SubWorkflowWakeKey::subspace(workflow_id));
@@ -1742,34 +1745,29 @@ impl Database for DatabaseKv {
 						Serializable,
 					);
 
-					let (wrote_to_wake_idx, tag_keys, wake_deadline_entry) = tokio::try_join!(
+					let (wrote_to_wake_idx, tag_keys, wake_deadline) = tokio::try_join!(
 						// Check for other workflows waiting on this one, wake all
 						async {
 							let mut wrote_to_wake_idx = false;
 
 							while let Some(entry) = stream.try_next().await? {
-								let sub_workflow_wake_key =
-									self.subspace
-										.unpack::<keys::wake::SubWorkflowWakeKey>(&entry.key())?;
-								let workflow_name =
-									sub_workflow_wake_key.deserialize(entry.value())?;
-
-								let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new(
-									workflow_name,
-									sub_workflow_wake_key.workflow_id,
-									keys::wake::WakeCondition::SubWorkflow {
-										sub_workflow_id: workflow_id,
-									},
-								);
+								let (sub_workflow_wake_key, workflow_name) =
+									tx.read_entry::<keys::wake::SubWorkflowWakeKey>(&entry)?;
 
 								// Add wake condition for workflow
-								tx.set(
-									&self.subspace.pack(&wake_condition_key),
-									&wake_condition_key.serialize(())?,
-								);
+								tx.write(
+									&keys::wake::WorkflowWakeConditionKey::new(
+										workflow_name,
+										sub_workflow_wake_key.workflow_id,
+										keys::wake::WakeCondition::SubWorkflow {
+											sub_workflow_id: workflow_id,
+										},
+									),
+									(),
+								)?;
 
 								// Clear secondary index
-								tx.clear(entry.key());
+								tx.delete(&sub_workflow_wake_key);
 
 								wrote_to_wake_idx = true;
 							}
@@ -1785,73 +1783,111 @@ impl Database for DatabaseKv {
 							Serializable,
 						)
 						.map(|res| {
-							self.subspace
-								.unpack::<keys::workflow::TagKey>(res?.key())
+							tx.unpack::<keys::workflow::TagKey>(res?.key())
 								.map_err(anyhow::Error::from)
 						})
 						.try_collect::<Vec<_>>(),
-						tx.get(&self.subspace.pack(&wake_deadline_key), Serializable),
+						tx.read_opt(&wake_deadline_key, Serializable),
 					)?;
 
 					for key in tag_keys {
-						let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::new(
+						tx.delete(&keys::workflow::ByNameAndTagKey::new(
 							workflow_name.to_string(),
 							key.k,
 							key.v,
 							workflow_id,
-						);
-						tx.clear(&self.subspace.pack(&by_name_and_tag_key));
+						));
 					}
 
 					// Clear null key
-					{
-						let by_name_and_tag_key = keys::workflow::ByNameAndTagKey::null(
-							workflow_name.to_string(),
-							workflow_id,
-						);
-						tx.clear(&self.subspace.pack(&by_name_and_tag_key));
-					}
+					tx.delete(&keys::workflow::ByNameAndTagKey::null(
+						workflow_name.to_string(),
+						workflow_id,
+					));
 
 					// Get and clear the pending deadline wake condition, if any. This could be put in the
 					// `pull_workflows` function (where we clear secondary indexes) but we chose to clear it
 					// here and in `commit_workflow` because its not a secondary index so theres no worry of
 					// it inserting more wake conditions. This reduces the load on `pull_workflows`. The
 					// reason this isn't immediately cleared in `pull_workflows` along with the rest of the
 					// wake conditions is because it might be in the future.
-					if let Some(raw) = wake_deadline_entry {
-						let deadline_ts = wake_deadline_key.deserialize(&raw)?;
-
-						let wake_condition_key = keys::wake::WorkflowWakeConditionKey::new(
+					if let Some(deadline_ts) = wake_deadline {
+						tx.delete(&keys::wake::WorkflowWakeConditionKey::new(
 							workflow_name.to_string(),
 							workflow_id,
 							keys::wake::WakeCondition::Deadline { deadline_ts },
-						);
-
-						tx.clear(&self.subspace.pack(&wake_condition_key));
+						));
 					}
 
 					// Clear "has wake condition"
-					let has_wake_condition_key =
-						keys::workflow::HasWakeConditionKey::new(workflow_id);
-					tx.clear(&self.subspace.pack(&has_wake_condition_key));
+					tx.delete(&keys::workflow::HasWakeConditionKey::new(workflow_id));
 
 					// Write output
 					let output_key = keys::workflow::OutputKey::new(workflow_id);
 
 					for (i, chunk) in output_key.split_ref(output)?.into_iter().enumerate() {
 						let chunk_key = output_key.chunk(i);
 
-						tx.set(&self.subspace.pack(&chunk_key), &chunk);
+						tx.set(&tx.pack(&chunk_key), &chunk);
+					}
+
+					// Mark all remaining pending signals as ack'd
+					let now = rivet_util::timestamp::now();
+
+					// This subspace contains all pending signals for the workflow
+					let pending_signal_subspace = self.subspace.subspace(
+						&keys::workflow::PendingSignalKey::workflow_subspace(workflow_id),
+					);
+
+					let mut stream = tx.get_ranges_keyvalues(
+						universaldb::RangeOption {
+							mode: StreamingMode::WantAll,
+							..(&pending_signal_subspace).into()
+						},
+						Serializable,
+					);
+
+					let mut pending_signal_cleared_count = 0;
+
+					loop {
+						if start.elapsed() > EARLY_TXN_TIMEOUT {
+							tracing::warn!(
+								"timed out processing pending signals for complete workflow"
+							);
+							break;
+						}
+
+						let Some(entry) = stream.try_next().await? else {
+							break;
+						};
+
+						let (key, _) = tx.read_entry::<keys::workflow::PendingSignalKey>(&entry)?;
+
+						// Ack signal
+						tx.add_conflict_key(&key, ConflictRangeType::Read)?;
+
+						tx.write(&keys::signal::AckTsKey::new(key.signal_id), now)?;
+
+						update_metric(
+							&tx,
+							Some(keys::metric::GaugeMetric::SignalPending(
+								key.signal_name.to_string(),
+							)),
+							None,
+						);
+
+						// Clear pending signal key
+						tx.delete(&key);
+
+						pending_signal_cleared_count += 1;
 					}
 
 					// Clear lease
-					let lease_key = keys::workflow::LeaseKey::new(workflow_id);
-					tx.clear(&self.subspace.pack(&lease_key));
-					let worker_id_key = keys::workflow::WorkerIdKey::new(workflow_id);
-					tx.clear(&self.subspace.pack(&worker_id_key));
+					tx.delete(&keys::workflow::LeaseKey::new(workflow_id));
+					tx.delete(&keys::workflow::WorkerIdKey::new(workflow_id));
 
 					update_metric(
-						&tx.with_subspace(self.subspace.clone()),
+						&tx,
 						Some(keys::metric::GaugeMetric::WorkflowActive(
 							workflow_name.to_string(),
 						)),
@@ -1860,7 +1896,7 @@ impl Database for DatabaseKv {
 						)),
 					);
 
-					Ok(wrote_to_wake_idx)
+					Ok((wrote_to_wake_idx, pending_signal_cleared_count))
 				}
 			})
 			.custom_instrument(tracing::info_span!("complete_workflows_tx"))
@@ -1874,6 +1910,10 @@ impl Database for DatabaseKv {
 			self.bump(BumpSubSubject::Worker);
 		}
 
+		if pending_signal_cleared_count != 0 {
+			tracing::debug!(count=%pending_signal_cleared_count, "cleared pending signals after workflow completed");
+		}
+
 		let dt = start_instant.elapsed().as_secs_f64();
 		metrics::COMPLETE_WORKFLOW_DURATION
 			.with_label_values(&[workflow_name])

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

@@ -11,7 +11,7 @@ use universaldb::{
 use universalpubsub::PublishOpts;
 use vbare::OwnedVersionedData;
 
-use crate::{keys, metrics, workflows::actor::Allocate};
+use crate::{keys, workflows::actor::Allocate};
 
 /// Batch size of how many events to ack.
 const EVENT_ACK_BATCH_SIZE: i64 = 500;

engine/packages/pegboard/src/workflows/runner.rs

@@ -11,7 +11,7 @@ use universaldb::{
 use universalpubsub::PublishOpts;
 use vbare::OwnedVersionedData;
 
-use crate::{keys, metrics, workflows::actor::Allocate};
+use crate::{keys, workflows::actor::Allocate};
 
 const EARLY_TXN_TIMEOUT: Duration = Duration::from_millis(2500);
 

engine/packages/pegboard/src/workflows/runner2.rs

@@ -389,7 +389,7 @@ impl PubSubDriver for PostgresDriver {
 		let encoded = BASE64.encode(payload);
 		let hashed = self.hash_subject(subject);
 
-		tracing::debug!("attempting to get connection for publish");
+		tracing::trace!("attempting to get connection for publish");
 
 		// Wait for listen connection to be ready first if this channel has subscribers
 		// This ensures that if we're reconnecting, the LISTEN is re-registered before NOTIFY

engine/packages/universalpubsub/src/driver/postgres/mod.rs

@@ -7,10 +7,9 @@ edition.workspace = true
 
 [dependencies]
 anyhow.workspace = true
-gas.workspace = true
 epoxy.workspace = true
-rivet-config.workspace = true
-tracing.workspace = true
-
+gas.workspace = true
 namespace.workspace = true
 pegboard.workspace = true
+rivet-config.workspace = true
+tracing.workspace = true