Fix checks removal when removing service (#5457)

Fix my recently discovered issue described here: #5456

agent/agent.go

@@ -2068,7 +2068,10 @@ func (a *Agent) removeServiceLocked(serviceID string, persist bool) error {
 
 	checks := a.State.Checks()
 	var checkIDs []types.CheckID
-	for id := range checks {
+	for id, check := range checks {
+		if check.ServiceID != serviceID {
+			continue
+		}
 		checkIDs = append(checkIDs, id)
 	}
 

agent/agent_test.go

@@ -605,6 +605,7 @@ func TestAgent_RemoveService(t *testing.T) {
 
 	// Removing a service with multiple checks works
 	{
+		// add a service to remove
 		srv := &structs.NodeService{
 			ID:      "redis",
 			Service: "redis",
@@ -618,6 +619,20 @@ func TestAgent_RemoveService(t *testing.T) {
 			t.Fatalf("err: %v", err)
 		}
 
+		// add another service that wont be affected
+		srv = &structs.NodeService{
+			ID:      "mysql",
+			Service: "mysql",
+			Port:    3306,
+		}
+		chkTypes = []*structs.CheckType{
+			&structs.CheckType{TTL: time.Minute},
+			&structs.CheckType{TTL: 30 * time.Second},
+		}
+		if err := a.AddService(srv, chkTypes, false, "", ConfigSourceLocal); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+
 		// Remove the service
 		if err := a.RemoveService("redis", false); err != nil {
 			t.Fatalf("err: %v", err)
@@ -636,13 +651,33 @@ func TestAgent_RemoveService(t *testing.T) {
 			t.Fatalf("check redis:2 should be removed")
 		}
 
-		// Ensure a TTL is setup
+		// Ensure the redis checks are removed
 		if _, ok := a.checkTTLs["service:redis:1"]; ok {
 			t.Fatalf("check ttl for redis:1 should be removed")
 		}
+		if check := a.State.Check(types.CheckID("service:redis:1")); check != nil {
+			t.Fatalf("check ttl for redis:1 should be removed")
+		}
 		if _, ok := a.checkTTLs["service:redis:2"]; ok {
 			t.Fatalf("check ttl for redis:2 should be removed")
 		}
+		if check := a.State.Check(types.CheckID("service:redis:2")); check != nil {
+			t.Fatalf("check ttl for redis:2 should be removed")
+		}
+
+		// check the mysql service is unnafected
+		if _, ok := a.checkTTLs["service:mysql:1"]; !ok {
+			t.Fatalf("check ttl for mysql:1 should not be removed")
+		}
+		if check := a.State.Check(types.CheckID("service:mysql:1")); check == nil {
+			t.Fatalf("check ttl for mysql:1 should not be removed")
+		}
+		if _, ok := a.checkTTLs["service:mysql:2"]; !ok {
+			t.Fatalf("check ttl for mysql:2 should not be removed")
+		}
+		if check := a.State.Check(types.CheckID("service:mysql:2")); check == nil {
+			t.Fatalf("check ttl for mysql:2 should not be removed")
+		}
 	}
 }