allocRunner: add test for allocation wipe

Author: rodrigol-chan

client/allocrunner/alloc_runner_unix_test.go

@@ -300,6 +300,100 @@ func TestAllocRunner_PreStartFailuresLeadToFailed(t *testing.T) {
 	}
 }
 
+// TestAllocRunner_Restore_VolatileAllocDir asserts that restoring a running
+// system alloc rebuilds task directories after client.alloc_dir was wiped,
+// while state in data_dir persists.
+func TestAllocRunner_Restore_VolatileAllocDir(t *testing.T) {
+	ci.Parallel(t)
+
+	alloc := mock.Alloc()
+	alloc.Job.Type = structs.JobTypeSystem
+	taskName := alloc.Job.TaskGroups[0].Tasks[0].Name
+	alloc.Job.TaskGroups[0].Tasks[0].Driver = "mock_driver"
+	alloc.Job.TaskGroups[0].Tasks[0].Config = map[string]interface{}{
+		"run_for": "1h",
+	}
+
+	countTaskDirBuildEvents := func(ar *allocRunner) int {
+		state := ar.AllocState()
+		ts, ok := state.TaskStates[taskName]
+		if !ok || ts == nil {
+			return 0
+		}
+
+		count := 0
+		for _, event := range ts.Events {
+			if event.Type == structs.TaskSetup && event.Message == structs.TaskBuildingTaskDir {
+				count++
+			}
+		}
+
+		return count
+	}
+
+	conf, cleanup := testAllocRunnerConfig(t, alloc.Copy())
+	defer cleanup()
+
+	// Persist task hook state across runner restarts to simulate data_dir.
+	conf.StateDB = state.NewMemDB(conf.Logger)
+
+	arIface, err := NewAllocRunner(conf)
+	must.NoError(t, err)
+	ar := arIface.(*allocRunner)
+
+	go ar.Run()
+	defer destroy(ar)
+
+	testutil.WaitForResult(func() (bool, error) {
+		if countTaskDirBuildEvents(ar) == 0 {
+			return false, fmt.Errorf("expected task dir build hook to run on initial run")
+		}
+		return true, nil
+	}, func(err error) {
+		require.NoError(t, err)
+	})
+
+	initialTaskDirBuildEvents := countTaskDirBuildEvents(ar)
+	require.Greater(t, initialTaskDirBuildEvents, 0)
+	allocDirPath := ar.GetAllocDir().AllocDirPath()
+
+	// Stop the runner and simulate host reboot where client.alloc_dir is wiped.
+	ar.Shutdown()
+	select {
+	case <-ar.ShutdownCh():
+	case <-time.After(30 * time.Second):
+		require.Fail(t, "AR took too long to exit")
+	}
+
+	require.NoError(t, os.RemoveAll(allocDirPath))
+
+	conf2, cleanup2 := testAllocRunnerConfig(t, alloc.Copy())
+	defer cleanup2()
+
+	// Reuse persistent state and the same alloc storage path after restart.
+	conf2.StateDB = conf.StateDB
+	conf2.ClientConfig.AllocDir = conf.ClientConfig.AllocDir
+	conf2.ClientConfig.AllocMountsDir = conf.ClientConfig.AllocMountsDir
+
+	ar2Iface, err := NewAllocRunner(conf2)
+	must.NoError(t, err)
+	ar2 := ar2Iface.(*allocRunner)
+
+	must.NoError(t, ar2.Restore())
+
+	go ar2.Run()
+	defer destroy(ar2)
+
+	testutil.WaitForResult(func() (bool, error) {
+		if countTaskDirBuildEvents(ar2) <= initialTaskDirBuildEvents {
+			return false, fmt.Errorf("expected task dir build hook to rerun on restore run")
+		}
+		return true, nil
+	}, func(err error) {
+		require.NoError(t, err)
+	})
+}
+
 type allocFailingPrestartHook struct{}
 
 func (*allocFailingPrestartHook) Name() string { return "failing_prestart" }