From 53c3aebbec4ecbc58a24f45579f7c1220d7b2da1 Mon Sep 17 00:00:00 2001
From: MasterPtato <23087326+MasterPtato@users.noreply.github.com>
Date: Tue, 6 Aug 2024 16:53:39 +0000
Subject: [PATCH] fix(workflows): Throw errors for duplicate workflows (#1011)

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

## Changes

<!-- If there are frontend changes, please include screenshots. -->
---
 lib/chirp-workflow/core/src/error.rs          |  3 +++
 lib/chirp-workflow/core/src/registry.rs       | 20 +++++++++++---
 svc/pkg/cluster/src/lib.rs                    | 26 +++++++++----------
 .../cluster/src/workflows/datacenter/mod.rs   |  7 +++--
 svc/pkg/linode/src/lib.rs                     | 10 +++----
 5 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/lib/chirp-workflow/core/src/error.rs b/lib/chirp-workflow/core/src/error.rs
index 0cb33a10a8..77462f9ee1 100644
--- a/lib/chirp-workflow/core/src/error.rs
+++ b/lib/chirp-workflow/core/src/error.rs
@@ -128,6 +128,9 @@ pub enum WorkflowError {
 
 	#[error("operation timed out")]
 	OperationTimeout,
+
+	#[error("duplicate registered workflow: {0}")]
+	DuplicateRegisteredWorkflow(String),
 }
 
 impl WorkflowError {
diff --git a/lib/chirp-workflow/core/src/registry.rs b/lib/chirp-workflow/core/src/registry.rs
index 70a1872399..395c8418d6 100644
--- a/lib/chirp-workflow/core/src/registry.rs
+++ b/lib/chirp-workflow/core/src/registry.rs
@@ -27,13 +27,25 @@ impl Registry {
 		Arc::new(self)
 	}
 
-	pub fn merge(mut self, registry: Registry) -> Registry {
+	pub fn merge(mut self, registry: Registry) -> WorkflowResult<Registry> {
+		// Check for duplicates
+		for (workflow_name, _) in &registry.workflows {
+			if self.workflows.contains_key(workflow_name.as_str()) {
+				return Err(WorkflowError::DuplicateRegisteredWorkflow(workflow_name.clone()));
+			}
+		}
+
 		self.workflows.extend(registry.workflows.into_iter());
 
-		self
+		Ok(self)
 	}
 
-	pub fn register_workflow<W: Workflow>(&mut self) {
+	pub fn register_workflow<W: Workflow>(&mut self) -> WorkflowResult<()> {
+		// Check for duplicates
+		if self.workflows.contains_key(W::NAME) {
+			return Err(WorkflowError::DuplicateRegisteredWorkflow(W::NAME.to_string()));
+		}
+
 		self.workflows.insert(
 			W::NAME.to_string(),
 			Arc::new(RegistryWorkflow {
@@ -72,6 +84,8 @@ impl Registry {
 				},
 			}),
 		);
+
+		Ok(())
 	}
 
 	pub fn get_workflow(&self, name: &str) -> WorkflowResult<&Arc<RegistryWorkflow>> {
diff --git a/svc/pkg/cluster/src/lib.rs b/svc/pkg/cluster/src/lib.rs
index 5a3fbb5a63..8cbdb35837 100644
--- a/svc/pkg/cluster/src/lib.rs
+++ b/svc/pkg/cluster/src/lib.rs
@@ -5,21 +5,21 @@ pub mod types;
 pub mod util;
 pub mod workflows;
 
-pub fn registry() -> Registry {
+pub fn registry() -> WorkflowResult<Registry> {
 	use workflows::*;
 
 	let mut registry = Registry::new();
-	registry.register_workflow::<cluster::Workflow>();
-	registry.register_workflow::<datacenter::Workflow>();
-	registry.register_workflow::<datacenter::scale::Workflow>();
-	registry.register_workflow::<datacenter::tls_issue::Workflow>();
-	registry.register_workflow::<server::Workflow>();
-	registry.register_workflow::<server::install::Workflow>();
-	registry.register_workflow::<server::dns_create::Workflow>();
-	registry.register_workflow::<server::dns_delete::Workflow>();
-	registry.register_workflow::<server::drain::Workflow>();
-	registry.register_workflow::<server::undrain::Workflow>();
-	registry.register_workflow::<prebake::Workflow>();
+	registry.register_workflow::<cluster::Workflow>()?;
+	registry.register_workflow::<datacenter::Workflow>()?;
+	registry.register_workflow::<datacenter::scale::Workflow>()?;
+	registry.register_workflow::<datacenter::tls_issue::Workflow>()?;
+	registry.register_workflow::<server::Workflow>()?;
+	registry.register_workflow::<server::install::Workflow>()?;
+	registry.register_workflow::<server::dns_create::Workflow>()?;
+	registry.register_workflow::<server::dns_delete::Workflow>()?;
+	registry.register_workflow::<server::drain::Workflow>()?;
+	registry.register_workflow::<server::undrain::Workflow>()?;
+	registry.register_workflow::<prebake::Workflow>()?;
 
-	registry
+	Ok(registry)
 }
diff --git a/svc/pkg/cluster/src/workflows/datacenter/mod.rs b/svc/pkg/cluster/src/workflows/datacenter/mod.rs
index 3657c19fc5..96270f06f8 100644
--- a/svc/pkg/cluster/src/workflows/datacenter/mod.rs
+++ b/svc/pkg/cluster/src/workflows/datacenter/mod.rs
@@ -61,8 +61,11 @@ pub(crate) async fn cluster_datacenter(ctx: &mut WorkflowCtx, input: &Input) ->
 	)
 	.await?;
 
-	// Scale
-	ctx.signal(ctx.workflow_id(), Scale {}).await?;
+	// Scale initially
+	ctx.workflow(scale::Input {
+		datacenter_id: input.datacenter_id,
+	})
+	.await?;
 
 	ctx.repeat(|ctx| {
 		let datacenter_id = input.datacenter_id;
diff --git a/svc/pkg/linode/src/lib.rs b/svc/pkg/linode/src/lib.rs
index ccd841f763..c340e38d74 100644
--- a/svc/pkg/linode/src/lib.rs
+++ b/svc/pkg/linode/src/lib.rs
@@ -5,13 +5,13 @@ pub mod types;
 pub mod util;
 pub mod workflows;
 
-pub fn registry() -> Registry {
+pub fn registry() -> WorkflowResult<Registry> {
 	use workflows::*;
 
 	let mut registry = Registry::new();
-	registry.register_workflow::<image::Workflow>();
-	registry.register_workflow::<server::Workflow>();
-	registry.register_workflow::<server::cleanup::Workflow>();
+	registry.register_workflow::<image::Workflow>()?;
+	registry.register_workflow::<server::Workflow>()?;
+	registry.register_workflow::<server::cleanup::Workflow>()?;
 
-	registry
+	Ok(registry)
 }