taskDirHook: ensure hooks rerun if task dir is missing

Author: rodrigol-chan

client/allocrunner/taskrunner/task_dir_hook.go

@@ -5,6 +5,7 @@ package taskrunner
 
 import (
 	"context"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -49,11 +50,18 @@ func (h *taskDirHook) Name() string {
 func (h *taskDirHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
 	fsi := h.runner.driverCapabilities.FSIsolation
 	if v, ok := req.PreviousState[TaskDirHookIsDoneDataKey]; ok && v == "true" {
-		setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
-		resp.State = map[string]string{
-			TaskDirHookIsDoneDataKey: "true",
+		if _, err := os.Stat(h.runner.taskDir.Dir); err == nil {
+			setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
+			resp.State = map[string]string{
+				TaskDirHookIsDoneDataKey: "true",
+			}
+			return nil
 		}
-		return nil
+
+		// Previous state indicated that we ran the hooks. However, the expected
+		// directories are not present. Fallthrough to run the hooks again. This
+		// can happen when e.g. `client.alloc_dir` is in ephemeral storage and
+		// the client reboots.
 	}
 
 	cc := h.runner.clientConfig