From 4de87e12e8aead60743738f6092cc4266e5560d0 Mon Sep 17 00:00:00 2001
From: Siddarth Kumar <siddarthkay@gmail.com>
Date: Sat, 11 May 2024 21:58:16 +0530
Subject: [PATCH] Revert "fix: update status of maintenance check"

This reverts commit 02c65d66be49fddbbd94240c6cacb359928f112c.
---
 .changelog/11332.txt |  2 +-
 agent/agent.go       | 24 ++++++------------------
 agent/agent_test.go  | 19 ++++++-------------
 3 files changed, 13 insertions(+), 32 deletions(-)

diff --git a/.changelog/11332.txt b/.changelog/11332.txt
index e3f506bd1f..b8c96368ee 100644
--- a/.changelog/11332.txt
+++ b/.changelog/11332.txt
@@ -1,3 +1,3 @@
 ```release-note:bug
-agent: update maintenance check to passing instead of removing the check
+agent: update maintenance check to passing before removing
 ```
diff --git a/agent/agent.go b/agent/agent.go
index 0eb540167b..28433dc08a 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -4194,7 +4194,7 @@ func (a *Agent) EnableServiceMaintenance(serviceID structs.ServiceID, reason, to
 
 	// Check if maintenance mode is not already enabled
 	checkID := serviceMaintCheckID(serviceID)
-	if a.State.Check(checkID) != nil && a.State.Check(checkID).Status != api.HealthPassing {
+	if a.State.Check(checkID) != nil {
 		return nil
 	}
 
@@ -4203,7 +4203,7 @@ func (a *Agent) EnableServiceMaintenance(serviceID structs.ServiceID, reason, to
 		reason = defaultServiceMaintReason
 	}
 
-	// New Critical Health Check
+	// Create and register the critical health check
 	check := &structs.HealthCheck{
 		Node:           a.config.NodeName,
 		CheckID:        checkID.ID,
@@ -4215,17 +4215,7 @@ func (a *Agent) EnableServiceMaintenance(serviceID structs.ServiceID, reason, to
 		Type:           "maintenance",
 		EnterpriseMeta: checkID.EnterpriseMeta,
 	}
-
-	// If check exists, update status, else create a new check
-	if a.State.Check(checkID) != nil {
-		a.State.UpdateCheck(checkID, api.HealthCritical, "")
-	} else {
-		err := a.AddCheck(check, nil, true, token, ConfigSourceLocal)
-		if err != nil {
-			return nil
-		}
-	}
-
+	a.AddCheck(check, nil, true, token, ConfigSourceLocal)
 	a.logger.Info("Service entered maintenance mode", "service", serviceID.String())
 
 	return nil
@@ -4248,11 +4238,9 @@ func (a *Agent) DisableServiceMaintenance(serviceID structs.ServiceID) error {
 	// Update check to trigger an event for watchers
 	a.State.UpdateCheck(checkID, api.HealthPassing, "")
 	// Make sure state change is propagated
-	err := a.State.SyncFull()
-	if err != nil {
-		return err
-	}
-
+	a.State.SyncChanges()
+	// Deregister the maintenance check
+	a.RemoveCheck(checkID, true)
 	a.logger.Info("Service left maintenance mode", "service", serviceID.String())
 
 	return nil
diff --git a/agent/agent_test.go b/agent/agent_test.go
index 7572d3aa95..316adeb3dd 100644
--- a/agent/agent_test.go
+++ b/agent/agent_test.go
@@ -3377,8 +3377,9 @@ func TestAgent_Service_MaintenanceMode(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
-	// Ensure the check has updated status
-	if found := a.State.Check(checkID); found != nil && found.Status != api.HealthPassing {
+	// Ensure the check was deregistered
+
+	if found := a.State.Check(checkID); found != nil {
 		t.Fatalf("should have deregistered maintenance check")
 	}
 
@@ -3392,10 +3393,7 @@ func TestAgent_Service_MaintenanceMode(t *testing.T) {
 	if check == nil {
 		t.Fatalf("should have registered critical check")
 	}
-	if check.Notes != "broken" {
-		t.Fatalf("bad: %#v", check)
-	}
-	if check.Status != api.HealthCritical {
+	if check.Notes != defaultServiceMaintReason {
 		t.Fatalf("bad: %#v", check)
 	}
 }
@@ -3622,10 +3620,8 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) {
 	// Leave maintenance mode
 	a.DisableNodeMaintenance()
 
-	// Ensure the check indicates passing status
-	if check.Status == api.HealthPassing {
-		t.Fatalf("bad: %#v", check)
-	}
+	// Ensure the check was deregistered
+	requireCheckMissing(t, a, structs.NodeMaint)
 
 	// Enter maintenance mode without passing a reason
 	a.EnableNodeMaintenance("", "")
@@ -3635,9 +3631,6 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) {
 	if check.Notes != defaultNodeMaintReason {
 		t.Fatalf("bad: %#v", check)
 	}
-	if check.Status != api.HealthCritical {
-		t.Fatalf("check status must be passing")
-	}
 }
 
 func TestAgent_checkStateSnapshot(t *testing.T) {