From cbe450bacbcd78e0f79c7c561d6899e695cecab6 Mon Sep 17 00:00:00 2001
From: MasterPtato <23087326+MasterPtato@users.noreply.github.com>
Date: Wed, 15 May 2024 00:35:08 +0000
Subject: [PATCH] fix: check for draining before installing/creating dns (#773)

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

## Changes

<!-- If there are frontend changes, please include screenshots. -->
---
 .../worker/src/workers/server_dns_create.rs   | 32 ++++++++++---------
 .../worker/src/workers/server_dns_delete.rs   | 11 ++++---
 .../worker/src/workers/server_install/mod.rs  | 30 +++--------------
 3 files changed, 27 insertions(+), 46 deletions(-)

diff --git a/svc/pkg/cluster/worker/src/workers/server_dns_create.rs b/svc/pkg/cluster/worker/src/workers/server_dns_create.rs
index ceca951ba4..5fef141651 100644
--- a/svc/pkg/cluster/worker/src/workers/server_dns_create.rs
+++ b/svc/pkg/cluster/worker/src/workers/server_dns_create.rs
@@ -11,7 +11,7 @@ use crate::util::CloudflareError;
 struct Server {
 	datacenter_id: Uuid,
 	public_ip: IpAddr,
-	cloud_destroy_ts: Option<i64>,
+	is_destroying_or_draining: bool,
 }
 
 #[worker(name = "cluster-server-dns-create")]
@@ -44,18 +44,6 @@ async fn inner(
 ) -> GlobalResult<()> {
 	let server_id = unwrap_ref!(ctx.server_id).as_uuid();
 
-	let server = sql_fetch_one!(
-		[ctx, Server, @tx tx]
-		"
-		SELECT
-			datacenter_id, public_ip, cloud_destroy_ts
-		FROM db_cluster.servers
-		WHERE server_id = $1
-		",
-		server_id,
-	)
-	.await?;
-
 	// Lock row
 	sql_execute!(
 		[ctx, @tx tx]
@@ -70,8 +58,22 @@ async fn inner(
 	)
 	.await?;
 
-	if server.cloud_destroy_ts.is_some() {
-		tracing::info!("server marked for deletion, not creating dns record");
+	let server = sql_fetch_one!(
+		[ctx, Server, @tx tx]
+		"
+		SELECT
+			datacenter_id,
+			public_ip,
+			(cloud_destroy_ts IS NOT NULL OR drain_ts IS NOT NULL) AS is_destroying_or_draining
+		FROM db_cluster.servers
+		WHERE server_id = $1
+		",
+		server_id,
+	)
+	.await?;
+
+	if server.is_destroying_or_draining {
+		tracing::info!("server marked for deletion/drain, not creating dns record");
 		return Ok(());
 	}
 
diff --git a/svc/pkg/cluster/worker/src/workers/server_dns_delete.rs b/svc/pkg/cluster/worker/src/workers/server_dns_delete.rs
index c8da4f653c..cdfffa917d 100644
--- a/svc/pkg/cluster/worker/src/workers/server_dns_delete.rs
+++ b/svc/pkg/cluster/worker/src/workers/server_dns_delete.rs
@@ -8,7 +8,7 @@ use proto::backend::pkg::*;
 use crate::util::CloudflareError;
 
 #[derive(sqlx::FromRow)]
-struct DnsRecord {
+struct DnsRecordRow {
 	dns_record_id: Option<String>,
 	secondary_dns_record_id: Option<String>,
 }
@@ -44,7 +44,7 @@ async fn inner(
 	let server_id = unwrap_ref!(ctx.server_id).as_uuid();
 
 	let row = sql_fetch_optional!(
-		[ctx, DnsRecord, @tx tx]
+		[ctx, DnsRecordRow, @tx tx]
 		"
 		SELECT dns_record_id, secondary_dns_record_id
 		FROM db_cluster.servers_cloudflare
@@ -57,7 +57,8 @@ async fn inner(
 		util::timestamp::now(),
 	)
 	.await?;
-	let Some(DnsRecord {
+
+	let Some(DnsRecordRow {
 		dns_record_id,
 		secondary_dns_record_id,
 	}) = row
@@ -87,7 +88,7 @@ async fn inner(
 			tracing::warn!(%zone_id, %record_id, "dns record not found");
 		} else {
 			res?;
-			tracing::warn!(%record_id, "deleted dns record");
+			tracing::info!(%record_id, "deleted dns record");
 		}
 	} else {
 		tracing::warn!("server has no primary dns record");
@@ -102,7 +103,7 @@ async fn inner(
 			})
 			.await?;
 
-		tracing::warn!(%record_id, "deleted secondary dns record");
+		tracing::info!(%record_id, "deleted secondary dns record");
 	} else {
 		tracing::warn!("server has no secondary dns record");
 	}
diff --git a/svc/pkg/cluster/worker/src/workers/server_install/mod.rs b/svc/pkg/cluster/worker/src/workers/server_install/mod.rs
index 308cbaea4d..9e7591bfa2 100644
--- a/svc/pkg/cluster/worker/src/workers/server_install/mod.rs
+++ b/svc/pkg/cluster/worker/src/workers/server_install/mod.rs
@@ -21,22 +21,22 @@ async fn worker(ctx: &OperationContext<cluster::msg::server_install::Message>) -
 	}
 
 	if let Some(server_id) = ctx.server_id {
-		let (is_destroying,) = sql_fetch_one!(
+		let (is_destroying_or_draining,) = sql_fetch_one!(
 			[ctx, (bool,)]
 			"
 			SELECT EXISTS(
 				SELECT 1
 				FROM db_cluster.servers
 				WHERE server_id = $1 AND
-				cloud_destroy_ts IS NOT NULL
+				(cloud_destroy_ts IS NOT NULL OR drain_ts IS NOT NULL)
 			)
 			",
 			server_id.as_uuid(),
 		)
 		.await?;
 
-		if is_destroying {
-			tracing::info!("server marked for deletion, not installing");
+		if is_destroying_or_draining {
+			tracing::info!("server marked for deletion/drain, not installing");
 			return Ok(());
 		}
 	}
@@ -133,28 +133,6 @@ async fn worker(ctx: &OperationContext<cluster::msg::server_install::Message>) -
 	})
 	.await??;
 
-	// Check if destroyed again after installing
-	if let Some(server_id) = ctx.server_id {
-		let (is_destroying,) = sql_fetch_one!(
-			[ctx, (bool,)]
-			"
-			SELECT EXISTS(
-				SELECT 1
-				FROM db_cluster.servers
-				WHERE server_id = $1 AND
-				cloud_destroy_ts IS NOT NULL
-			)
-			",
-			server_id.as_uuid(),
-		)
-		.await?;
-
-		if is_destroying {
-			tracing::info!("server marked for deletion, not marking as installed");
-			return Ok(());
-		}
-	}
-
 	let request_id = unwrap_ref!(ctx.request_id).as_uuid();
 	msg!([ctx] cluster::msg::server_install_complete(request_id) {
 		request_id: ctx.request_id,