diff --git a/api/tasks.go b/api/tasks.go index b9b79af54..8d3920529 100644 --- a/api/tasks.go +++ b/api/tasks.go @@ -632,6 +632,7 @@ type DispatchPayloadConfig struct { const ( TaskLifecycleHookPrestart = "prestart" + TaskLifecycleHookPoststart = "poststart" ) type TaskLifecycle struct { diff --git a/client/allocrunner/task_hook_coordinator.go b/client/allocrunner/task_hook_coordinator.go index 8f6d1c10b..2fd94714e 100644 --- a/client/allocrunner/task_hook_coordinator.go +++ b/client/allocrunner/task_hook_coordinator.go @@ -18,23 +18,31 @@ type taskHookCoordinator struct { mainTaskCtx context.Context mainTaskCtxCancel func() + poststartTaskCtx context.Context + poststartTaskCtxCancel func() + prestartSidecar map[string]struct{} prestartEphemeral map[string]struct{} + + mainTasksPending map[string]struct{} } func newTaskHookCoordinator(logger hclog.Logger, tasks []*structs.Task) *taskHookCoordinator { closedCh := make(chan struct{}) close(closedCh) - mainTaskCtx, cancelFn := context.WithCancel(context.Background()) + mainTaskCtx, mainCancelFn := context.WithCancel(context.Background()) + poststartTaskCtx, poststartCancelFn := context.WithCancel(context.Background()) c := &taskHookCoordinator{ logger: logger, closedCh: closedCh, mainTaskCtx: mainTaskCtx, - mainTaskCtxCancel: cancelFn, + mainTaskCtxCancel: mainCancelFn, prestartSidecar: map[string]struct{}{}, prestartEphemeral: map[string]struct{}{}, + poststartTaskCtx: poststartTaskCtx, + poststartTaskCtxCancel: poststartCancelFn, } c.setTasks(tasks) return c @@ -55,7 +63,6 @@ func (c *taskHookCoordinator) setTasks(tasks []*structs.Task) { } else { c.prestartEphemeral[task.Name] = struct{}{} } - default: c.logger.Error("invalid lifecycle hook", "hook", task.Lifecycle.Hook) } @@ -70,11 +77,25 @@ func (c *taskHookCoordinator) hasPrestartTasks() bool { return len(c.prestartSidecar)+len(c.prestartEphemeral) > 0 } +func (c *taskHookCoordinator) hasPendingMainTasks() bool { + return len(c.mainTasksPending) > 0 +} + func (c *taskHookCoordinator) startConditionForTask(task *structs.Task) <-chan struct{} { if task.Lifecycle != nil && task.Lifecycle.Hook == structs.TaskLifecycleHookPrestart { return c.closedCh } + switch task.Lifecycle.Hook { + case structs.TaskLifecycleHookPrestart: + // Prestart tasks start without checking status of other tasks + return c.closedCh + case structs.TaskLifecycleHookPoststart: + return c.poststartTaskCtx.Done() + default: + return c.mainTaskCtx.Done() + } + return c.mainTaskCtx.Done() } @@ -104,10 +125,23 @@ func (c *taskHookCoordinator) taskStateUpdated(states map[string]*structs.TaskSt delete(c.prestartEphemeral, task) } + for task := range c.mainTasksPending { + st := states[task] + if st == nil || st.StartedAt.IsZero() { + continue + } + + delete(c.mainTasksPending, task) + } + // everything well if !c.hasPrestartTasks() { c.mainTaskCtxCancel() } + + if !c.hasPendingMainTasks() { + c.poststartTaskCtxCancel() + } } // hasNonSidecarTasks returns false if all the passed tasks are sidecar tasks diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go index 4996b3638..437af243c 100644 --- a/nomad/structs/structs.go +++ b/nomad/structs/structs.go @@ -4906,6 +4906,7 @@ func (d *DispatchPayloadConfig) Validate() error { const ( TaskLifecycleHookPrestart = "prestart" + TaskLifecycleHookPoststart = "poststart" ) type TaskLifecycleConfig struct { @@ -4929,6 +4930,7 @@ func (d *TaskLifecycleConfig) Validate() error { switch d.Hook { case TaskLifecycleHookPrestart: + case TaskLifecycleHookPoststart: case "": return fmt.Errorf("no lifecycle hook provided") default: