fix(workflows): fix race condition of workflow waking before commit (#2748)

<!-- Please make sure there is an issue that this PR is correlated to. -->

## Changes

<!-- If there are frontend changes, please include screenshots. -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

* **New Features**
  * Enhanced CLI output to display workflow and signal timestamps with millisecond precision.
  * CLI now shows workflow and signal tags and workflow IDs in a formatted, readable manner.

* **Improvements**
  * Grafana traces dashboard updated: simplified queries, improved trace table with error indicators, reorganized panels, and enhanced trace details.
  * Logs and error panels in the dashboard were restructured for better visibility and usability.
  * Worker shutdown process improved to accurately reflect remaining workflows at start.
  * Workflow processing logic updated to ensure timely worker wake-up and avoid race conditions.

* **Bug Fixes**
  * Reduced the threshold for logging long signal receive lag, enabling earlier detection of delays.

* **Chores**
  * Added a new error variant for SQL connection failures to improve error clarity.
  * Updated test structure: converted async test to sync with explicit runtime and removed obsolete test files.
  * Improved internal logging messages for metrics loop shutdown.
  * Minor comment and typo corrections for clarity.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: MasterPtato

docker/dev-full/grafana/dashboards/traces.json

@@ -48,6 +48,8 @@ exporters:
       initial_interval: 5s
       max_interval: 30s
       max_elapsed_time: 300s
+  # debug:
+  #   verbosity: detailed
 
 service:
   pipelines:

docker/dev-full/otel-collector/config.yaml

@@ -83,7 +83,7 @@ impl<'a> ListenCtx<'a> {
 			.with_label_values(&[self.ctx.name(), &signal.signal_name])
 			.observe(recv_lag);
 
-		if recv_lag > 15.0 {
+		if recv_lag > 3.0 {
 			// We print an error here so the trace of this workflow does not get dropped
 			tracing::error!(
 				?recv_lag,

packages/common/chirp-workflow/core/src/ctx/listen.rs

@@ -56,7 +56,7 @@ impl DatabaseCrdbNats {
 			Ok(conn)
 		} else {
 			// Create a new connection
-			self.pool.acquire().await.map_err(WorkflowError::Sqlx)
+			self.pool.acquire().await.map_err(WorkflowError::ConnSqlx)
 		}
 	}
 

packages/common/chirp-workflow/core/src/db/crdb_nats/mod.rs

@@ -1548,7 +1548,8 @@ impl Database for DatabaseFdbSqliteNats {
 		// Evict databases before releasing lease
 		self.evict_wf_sqlite(workflow_id).await?;
 
-		self.pools
+		let wrote_to_wake_idx = self
+			.pools
 			.fdb()?
 			.run(|tx, _mc| {
 				async move {
@@ -1569,9 +1570,11 @@ impl Database for DatabaseFdbSqliteNats {
 						SERIALIZABLE,
 					);
 
-					let (_, tag_keys, wake_deadline_entry) = tokio::try_join!(
+					let (wrote_to_wake_idx, tag_keys, wake_deadline_entry) = 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
@@ -1599,9 +1602,11 @@ impl Database for DatabaseFdbSqliteNats {
 
 								// Clear secondary index
 								tx.clear(entry.key());
+
+								wrote_to_wake_idx = true;
 							}
 
-							Result::<_, fdb::FdbBindingError>::Ok(())
+							Result::<_, fdb::FdbBindingError>::Ok(wrote_to_wake_idx)
 						},
 						// Read tags
 						async {
@@ -1695,13 +1700,16 @@ impl Database for DatabaseFdbSqliteNats {
 						keys::workflow::WorkerInstanceIdKey::new(workflow_id);
 					tx.clear(&self.subspace.pack(&worker_instance_id_key));
 
-					Ok(())
+					Ok(wrote_to_wake_idx)
 				}
 			})
 			.custom_instrument(tracing::info_span!("complete_workflows_tx"))
 			.await?;
 
-		self.wake_worker();
+		// Wake worker again in case some other workflow was waiting for this one to complete
+		if wrote_to_wake_idx {
+			self.wake_worker();
+		}
 
 		let dt = start_instant.elapsed().as_secs_f64();
 		metrics::COMPLETE_WORKFLOW_DURATION
@@ -1849,14 +1857,24 @@ impl Database for DatabaseFdbSqliteNats {
 			.custom_instrument(tracing::info_span!("commit_workflow_tx"))
 			.await?;
 
-		// Wake worker again if the deadline is before the next tick
-		if let Some(deadline_ts) = wake_deadline_ts {
-			if deadline_ts
-				<= rivet_util::timestamp::now() + self.worker_poll_interval().as_millis() as i64
-			{
-				self.wake_worker();
-			}
-		}
+		// Always wake the worker immediately again. This is an IMPORTANT implementation detail to prevent
+		// race conditions with workflow sleep. Imagine the scenario:
+		//
+		// 1. workflow is between user code and commit
+		// 2. worker reads wake condition for said workflow but cannot run it because it is already leased
+		// 3. workflow commits
+		//
+		// This will result in the workflow sleeping instead of immediately running again.
+		//
+		// Adding this wake_worker call ensures that if the workflow has a valid wake condition before commit
+		// then it will immediately wake up again.
+		//
+		// This is simpler than having this commit_workflow fn read wake conditions because:
+		// - the wake conditions are not indexed by wf id
+		// - would involve informing the worker to restart the workflow in memory instead of the usual
+		//   workflow lifecycle
+		// - the worker is already designed to pull wake conditions frequently
+		self.wake_worker();
 
 		let dt = start_instant.elapsed().as_secs_f64();
 		metrics::COMMIT_WORKFLOW_DURATION
@@ -2266,7 +2284,7 @@ impl Database for DatabaseFdbSqliteNats {
 					let wake_signal_key =
 						keys::workflow::WakeSignalKey::new(workflow_id, signal_name.to_string());
 
-					// If the workflow is currently listening for this signal, wake it
+					// If the workflow currently has a wake signal key for this signal, wake it
 					if tx
 						.get(&self.subspace.pack(&wake_signal_key), SERIALIZABLE)
 						.await?

packages/common/chirp-workflow/core/src/db/fdb_sqlite_nats/mod.rs

@@ -144,6 +144,9 @@ pub enum WorkflowError {
 	#[error("sql error: {0}")]
 	Sqlx(#[from] sqlx::Error),
 
+	#[error("failed to create sql connection: {0}")]
+	ConnSqlx(sqlx::Error),
+
 	#[error("fdb error: {0}")]
 	Fdb(#[from] fdb::FdbBindingError),
 

packages/common/chirp-workflow/core/src/error.rs

@@ -112,6 +112,10 @@ impl Worker {
 
 	#[tracing::instrument(skip_all)]
 	async fn shutdown(mut self, mut sigterm: Signal) {
+		// For an accurate count of how many remaining workflows there are, retain before checking length
+		self.running_workflows
+			.retain(|_, wf| !wf.handle.is_finished());
+
 		// Shutdown sequence
 		tracing::info!(
 			duration=?SHUTDOWN_DURATION,

packages/common/chirp-workflow/core/src/worker.rs

@@ -1,16 +1,33 @@
 use std::time::Duration;
 
 use chirp_workflow::prelude::*;
-// use serde_json::json;
-// use uuid::Uuid;
 
-mod common;
-use common::*;
+#[test]
+fn fdb_sqlite_nats_driver() {
+	// SAFETY: No other threads exist yet
+	unsafe {
+		std::env::set_var("RUST_LOG", "DEBUG");
+		std::env::set_var("RUST_LOG_ANSI_COLOR", "1");
+		std::env::set_var("RIVET_OTEL_ENABLED", "1");
+		std::env::set_var("RIVET_OTEL_SAMPLER_RATIO", "1");
+		std::env::set_var("RIVET_SERVICE_NAME", "test");
+		std::env::set_var("RIVET_OTEL_ENDPOINT", "http://127.0.0.1:4317");
+	}
 
-#[tokio::test(flavor = "multi_thread")]
-async fn fdb_sqlite_nats_driver() {
-	setup_tracing();
+	// Build runtime
+	chirp_workflow::prelude::__rivet_runtime::run(
+		chirp_workflow::prelude::tracing::Instrument::instrument(
+			async move {
+				tracing::info!("test starting");
+				fdb_sqlite_nats_driver_inner().await;
+				tracing::info!("test finished");
+			},
+			chirp_workflow::prelude::tracing::info_span!("fdb_sqlite_nats_driver"),
+		),
+	);
+}
 
+async fn fdb_sqlite_nats_driver_inner() {
 	let ctx = chirp_workflow::prelude::TestCtx::from_env::<db::DatabaseFdbSqliteNats>(
 		"fdb_sqlite_nats_driver",
 		true,
@@ -62,8 +79,6 @@ async fn fdb_sqlite_nats_driver() {
 
 mod def {
 	use chirp_workflow::prelude::*;
-	use futures_util::FutureExt;
-	use sqlx::Acquire;
 
 	#[derive(Debug, Serialize, Deserialize)]
 	pub struct Input {}