diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
index 025a18ab7..721a6e016 100644
--- a/client/alloc_runner_test.go
+++ b/client/alloc_runner_test.go
@@ -321,7 +321,8 @@ func TestAllocRunner_SaveRestoreState(t *testing.T) {
 	}
 
 	// Create a new alloc runner
-	ar2 := NewAllocRunner(ar.logger, ar.config, upd.Update,
+	l2 := prefixedTestLogger("----- ar2:  ")
+	ar2 := NewAllocRunner(l2, ar.config, upd.Update,
 		&structs.Allocation{ID: ar.alloc.ID}, ar.vaultClient)
 	err = ar2.RestoreState()
 	if err != nil {
@@ -341,7 +342,7 @@ func TestAllocRunner_SaveRestoreState(t *testing.T) {
 		last := upd.Allocs[upd.Count-1]
 		return last.ClientStatus == structs.AllocClientStatusRunning, nil
 	}, func(err error) {
-		t.Fatalf("err: %v %#v %#v", err, upd.Allocs[0], ar.alloc.TaskStates)
+		t.Fatalf("err: %v %#v %#v", err, upd.Allocs[0], ar2.alloc.TaskStates["web"])
 	})
 
 	// Destroy and wait
diff --git a/client/task_runner.go b/client/task_runner.go
index 02e721102..36c0d7646 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -1212,8 +1212,13 @@ func (r *TaskRunner) buildTaskDir(fsi cstructs.FSIsolation) error {
 	built := r.taskDirBuilt
 	r.persistLock.Unlock()
 
-	r.setState(structs.TaskStatePending, structs.NewTaskEvent(structs.TaskSetup).
-		SetMessage(structs.TaskBuildingTaskDir))
+	// We do not set the state again since this only occurs during restoration
+	// and the task dir is already built. The reason we call Build again is to
+	// ensure that the task dir invariants are still held.
+	if !built {
+		r.setState(structs.TaskStatePending, structs.NewTaskEvent(structs.TaskSetup).
+			SetMessage(structs.TaskBuildingTaskDir))
+	}
 
 	chroot := config.DefaultChrootEnv
 	if len(r.config.ChrootEnv) > 0 {