client: allow incomplete allocrunners to be removed on restore (#16638)

If an allocrunner is persisted to the client state but the client stops before
task runner can start, we end up with an allocation in the database with
allocrunner state but no taskrunner state. This ends up mimicking an old
pre-0.9.5 state where this state was not recorded and that hits a backwards
compatibility shim. This leaves allocations in the client state that can never
be restored, but won't ever be removed either.

Update the backwards compatibility shim so that we fail the restore for the
allocrunner and remove the allocation from the client state. Taskrunners persist
state during graceful shutdown, so it shouldn't be possible to leak tasks that
have actually started. This lets us "start over" with the allocation, if the
server still wants to place it on the client.

.changelog/16638.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+client: remove incomplete allocation entries from client state database during client restarts
+```

client/client.go

@@ -1206,43 +1206,35 @@ func (c *Client) restoreState() error {
 		return nil
 	}
 
-	//XXX REMOVED! make a note in backward compat / upgrading doc
-	// COMPAT: Remove in 0.7.0
-	// 0.6.0 transitioned from individual state files to a single bolt-db.
-	// The upgrade path is to:
-	// Check if old state exists
-	//   If so, restore from that and delete old state
-	// Restore using state database
-
 	// Restore allocations
 	allocs, allocErrs, err := c.stateDB.GetAllAllocations()
 	if err != nil {
 		return err
 	}
-
 	for allocID, err := range allocErrs {
+		// The boltdb values don't exist or didn't decode correctly. The state
+		// store is corrupt or there was a backwards incompatibility. Either way
+		// log it so we can debug it with operator client-state command.
 		c.logger.Error("error restoring alloc", "error", err, "alloc_id", allocID)
-		//TODO Cleanup
-		// Try to clean up alloc dir
-		// Remove boltdb entries?
-		// Send to server with clientstatus=failed
 	}
 
 	// Load each alloc back
 	for _, alloc := range allocs {
 
-		// COMPAT(0.12): remove once upgrading from 0.9.5 is no longer supported
-		// See hasLocalState for details.  Skipping suspicious allocs
-		// now.  If allocs should be run, they will be started when the client
-		// gets allocs from servers.
+		// If the alloc has no task state, we most likely stopped the client
+		// after the allocrunner was created but before tasks could
+		// start. Remove the client state so that we can start over with this
+		// alloc if the server still wants to place it here.
 		if !c.hasLocalState(alloc) {
-			c.logger.Warn("found an alloc without any local state, skipping restore", "alloc_id", alloc.ID)
+			c.logger.Warn(
+				"found an alloc without any local state, deleting from client state db",
+				"alloc_id", alloc.ID)
+			c.stateDB.DeleteAllocationBucket(alloc.ID, state.WithBatchMode())
 			continue
 		}
 
-		//XXX On Restore we give up on watching previous allocs because
-		//    we need the local AllocRunners initialized first. We could
-		//    add a second loop to initialize just the alloc watcher.
+		// On Restore we give up on watching previous allocs because we need the
+		// local AllocRunners initialized first.
 		prevAllocWatcher := allocwatcher.NoopPrevAlloc{}
 		prevAllocMigrator := allocwatcher.NoopPrevAlloc{}
 
@@ -1305,11 +1297,15 @@ func (c *Client) restoreState() error {
 }
 
 // hasLocalState returns true if we have any other associated state
-// with alloc beyond the task itself
+// with alloc beyond the task itself. We want to detect two possible scenarios:
 //
-// Useful for detecting if a potentially completed alloc got resurrected
-// after AR was destroyed.  In such cases, re-running the alloc lead to
-// unexpected reruns and may lead to process and task exhaustion on node.
+//  1. A potentially completed alloc got resurrected after AR was destroyed. In
+//     such cases, re-running the alloc leads to unexpected reruns and may lead
+//     to process and task exhaustion on node.
+//
+//  2. The client stopped after the allocrunner was persisted to disk but before
+//     tasks started. In this case, we don't restart the AR and instead wait until
+//     we get the desired state from the server.
 //
 // The heuristic used here is an alloc is suspect if we see no other information
 // and no other task/status info is found.
@@ -1321,8 +1317,6 @@ func (c *Client) restoreState() error {
 // See:
 //   - https://github.com/hashicorp/nomad/pull/6207
 //   - https://github.com/hashicorp/nomad/issues/5984
-//
-// COMPAT(0.12): remove once upgrading from 0.9.5 is no longer supported
 func (c *Client) hasLocalState(alloc *structs.Allocation) bool {
 	tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
 	if tg == nil {

client/client_test.go

@@ -749,6 +749,10 @@ func TestClient_SaveRestoreState(t *testing.T) {
 		wait.Gap(time.Millisecond*30),
 	))
 
+	// Create a corrupted allocation that will be removed during restore
+	corruptAlloc := mock.Alloc()
+	c1.stateDB.PutAllocation(corruptAlloc)
+
 	t.Log("shutting down client")
 	must.NoError(t, c1.Shutdown()) // note: this saves the client state DB
 
@@ -864,6 +868,9 @@ func TestClient_SaveRestoreState(t *testing.T) {
 							"alloc %s stopped during shutdown should have updated", a3.ID[:8])
 					}
 
+				case corruptAlloc.ID:
+					return fmt.Errorf("corrupted allocation should not have been restored")
+
 				default:
 					if ar.AllocState().ClientStatus != structs.AllocClientStatusComplete {
 						return fmt.Errorf("expected complete client status, got %v",