From 663805cd239c778c5021dc551d840287d710a38d Mon Sep 17 00:00:00 2001
From: Julian Sommer <115516918+juliansommer@users.noreply.github.com>
Date: Fri, 13 Mar 2026 01:53:00 +0000
Subject: [PATCH] fix: race conditions on node failure

---
 hive/daemon/daemon.go                         |  9 +++--
 .../ec2/instance/TerminateInstances.go        | 27 ++++++++++----
 .../ec2/instance/TerminateInstances_test.go   | 35 ++++++-------------
 3 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/hive/daemon/daemon.go b/hive/daemon/daemon.go
index 5e8e8c2d3..d41ddc4c3 100644
--- a/hive/daemon/daemon.go
+++ b/hive/daemon/daemon.go
@@ -797,9 +797,12 @@ func (d *Daemon) restoreInstances() {
 
 		if instance.Status == vm.StateTerminated {
 			if !d.migrateTerminatedToKV(instance) {
-				// KV write failed — still remove from local map. The VM is gone;
-				// keeping it here would make it retry forever on every restart.
-				delete(d.Instances.VMS, instance.ID)
+				// KV write failed — keep in local state so the next restart
+				// retries the migration. Deleting here would create a "void":
+				// the instance disappears from both local state and the
+				// terminated KV, making it invisible to DescribeInstances.
+				slog.Warn("Terminated instance KV migration failed, will retry on next restart",
+					"instance", instance.ID)
 			}
 			continue
 		}
diff --git a/hive/gateway/ec2/instance/TerminateInstances.go b/hive/gateway/ec2/instance/TerminateInstances.go
index 13cfc2871..5c100cdb0 100644
--- a/hive/gateway/ec2/instance/TerminateInstances.go
+++ b/hive/gateway/ec2/instance/TerminateInstances.go
@@ -51,12 +51,26 @@ func TerminateInstances(input *ec2.TerminateInstancesInput, natsConn *nats.Conn,
 			continue
 		}
 
-		// Send NATS request to the specific instance topic with account ID header
+		// Send NATS request to the specific instance topic with account ID header.
+		// Retry briefly on ErrNoResponders — after a cluster restart the
+		// per-instance NATS subscription may not have propagated to all
+		// servers yet.
 		subject := fmt.Sprintf("ec2.cmd.%s", instanceID)
-		reqMsg := nats.NewMsg(subject)
-		reqMsg.Data = jsonData
-		reqMsg.Header.Set(utils.AccountIDHeader, accountID)
-		msg, err := natsConn.RequestMsg(reqMsg, 5*time.Second)
+		var msg *nats.Msg
+		for attempt := range 3 {
+			reqMsg := nats.NewMsg(subject)
+			reqMsg.Data = jsonData
+			reqMsg.Header.Set(utils.AccountIDHeader, accountID)
+			msg, err = natsConn.RequestMsg(reqMsg, 5*time.Second)
+			if err == nil || !errors.Is(err, nats.ErrNoResponders) {
+				break
+			}
+			if attempt < 2 {
+				slog.Debug("TerminateInstances: No responder on per-instance topic, retrying",
+					"instance_id", instanceID, "attempt", attempt+1)
+				time.Sleep(time.Duration(attempt+1) * time.Second)
+			}
+		}
 		if err != nil {
 			// If no daemon owns this instance, try the ec2.terminate topic for stopped instances
 			if errors.Is(err, nats.ErrNoResponders) {
@@ -88,8 +102,7 @@ func TerminateInstances(input *ec2.TerminateInstancesInput, natsConn *nats.Conn,
 			}
 
 			slog.Error("TerminateInstances: Failed to send command", "instance_id", instanceID, "err", err)
-			stateChanges = append(stateChanges, newStateChange(instanceID, 16, "running", 16, "running"))
-			continue
+			return nil, fmt.Errorf("failed to terminate instance %s: %w", instanceID, err)
 		}
 
 		// Check if the daemon returned an error response (e.g. ownership check failure)
diff --git a/hive/gateway/ec2/instance/TerminateInstances_test.go b/hive/gateway/ec2/instance/TerminateInstances_test.go
index 5bd7aa75e..dc60cc71d 100644
--- a/hive/gateway/ec2/instance/TerminateInstances_test.go
+++ b/hive/gateway/ec2/instance/TerminateInstances_test.go
@@ -114,18 +114,11 @@ func TestTerminateInstances_NATSRequestFails(t *testing.T) {
 		InstanceIds: []*string{aws.String(instanceID)},
 	}
 
-	output, err := TerminateInstances(input, nc, "123456789012")
-
-	require.NoError(t, err)
-	require.Len(t, output.TerminatingInstances, 1)
-
-	// On NATS failure, state should reflect "still running"
-	sc := output.TerminatingInstances[0]
-	assert.Equal(t, instanceID, *sc.InstanceId)
-	assert.Equal(t, int64(16), *sc.CurrentState.Code)
-	assert.Equal(t, "running", *sc.CurrentState.Name)
-	assert.Equal(t, int64(16), *sc.PreviousState.Code)
-	assert.Equal(t, "running", *sc.PreviousState.Name)
+	// When no subscriber exists and all fallback paths fail, an error is
+	// returned so the caller knows the terminate did not succeed.
+	_, err := TerminateInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), instanceID)
 }
 
 func TestTerminateInstances_MixedSuccessAndFailure(t *testing.T) {
@@ -142,18 +135,12 @@ func TestTerminateInstances_MixedSuccessAndFailure(t *testing.T) {
 		InstanceIds: []*string{aws.String(goodID), aws.String(badID)},
 	}
 
-	output, err := TerminateInstances(input, nc, "123456789012")
-
-	require.NoError(t, err)
-	require.Len(t, output.TerminatingInstances, 2)
-
-	// First: success → shutting-down
-	assert.Equal(t, goodID, *output.TerminatingInstances[0].InstanceId)
-	assert.Equal(t, "shutting-down", *output.TerminatingInstances[0].CurrentState.Name)
-
-	// Second: failure → still running
-	assert.Equal(t, badID, *output.TerminatingInstances[1].InstanceId)
-	assert.Equal(t, "running", *output.TerminatingInstances[1].CurrentState.Name)
+	// When a running instance cannot be reached (ErrNoResponders exhausted after
+	// retries), TerminateInstances returns an error rather than silently
+	// reporting the instance as "running".
+	_, err := TerminateInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), badID)
 }
 
 func TestTerminateInstances_VerifiesQMPAttributes(t *testing.T) {