github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/allocrunner/task_hook_coordinator.go (about) 1 package allocrunner 2 3 import ( 4 "context" 5 6 "github.com/hashicorp/go-hclog" 7 "github.com/hashicorp/nomad/client/allocrunner/taskrunner" 8 "github.com/hashicorp/nomad/nomad/structs" 9 ) 10 11 // TaskHookCoordinator helps coordinate when mainTasks start tasks can launch 12 // namely after all Prestart Tasks have run, and after all BlockUntilCompleted have completed 13 type taskHookCoordinator struct { 14 logger hclog.Logger 15 16 // constant for quickly starting all prestart tasks 17 closedCh chan struct{} 18 19 // Each context is used to gate task runners launching the tasks. A task 20 // runner waits until the context associated its lifecycle context is 21 // done/cancelled. 22 mainTaskCtx context.Context 23 mainTaskCtxCancel func() 24 25 poststartTaskCtx context.Context 26 poststartTaskCtxCancel func() 27 poststopTaskCtx context.Context 28 poststopTaskCtxCancel context.CancelFunc 29 30 prestartSidecar map[string]struct{} 31 prestartEphemeral map[string]struct{} 32 mainTasksRunning map[string]struct{} // poststop: main tasks running -> finished 33 mainTasksPending map[string]struct{} // poststart: main tasks pending -> running 34 } 35 36 func newTaskHookCoordinator(logger hclog.Logger, tasks []*structs.Task) *taskHookCoordinator { 37 closedCh := make(chan struct{}) 38 close(closedCh) 39 40 mainTaskCtx, mainCancelFn := context.WithCancel(context.Background()) 41 poststartTaskCtx, poststartCancelFn := context.WithCancel(context.Background()) 42 poststopTaskCtx, poststopTaskCancelFn := context.WithCancel(context.Background()) 43 44 c := &taskHookCoordinator{ 45 logger: logger, 46 closedCh: closedCh, 47 mainTaskCtx: mainTaskCtx, 48 mainTaskCtxCancel: mainCancelFn, 49 prestartSidecar: map[string]struct{}{}, 50 prestartEphemeral: map[string]struct{}{}, 51 mainTasksRunning: map[string]struct{}{}, 52 mainTasksPending: map[string]struct{}{}, 53 poststartTaskCtx: poststartTaskCtx, 54 poststartTaskCtxCancel: poststartCancelFn, 55 poststopTaskCtx: poststopTaskCtx, 56 poststopTaskCtxCancel: poststopTaskCancelFn, 57 } 58 c.setTasks(tasks) 59 return c 60 } 61 62 func (c *taskHookCoordinator) setTasks(tasks []*structs.Task) { 63 for _, task := range tasks { 64 65 if task.Lifecycle == nil { 66 c.mainTasksPending[task.Name] = struct{}{} 67 c.mainTasksRunning[task.Name] = struct{}{} 68 continue 69 } 70 71 switch task.Lifecycle.Hook { 72 case structs.TaskLifecycleHookPrestart: 73 if task.Lifecycle.Sidecar { 74 c.prestartSidecar[task.Name] = struct{}{} 75 } else { 76 c.prestartEphemeral[task.Name] = struct{}{} 77 } 78 case structs.TaskLifecycleHookPoststart: 79 // Poststart hooks don't need to be tracked. 80 case structs.TaskLifecycleHookPoststop: 81 // Poststop hooks don't need to be tracked. 82 default: 83 c.logger.Error("invalid lifecycle hook", "task", task.Name, "hook", task.Lifecycle.Hook) 84 } 85 } 86 87 if !c.hasPrestartTasks() { 88 c.mainTaskCtxCancel() 89 } 90 } 91 92 func (c *taskHookCoordinator) hasPrestartTasks() bool { 93 return len(c.prestartSidecar)+len(c.prestartEphemeral) > 0 94 } 95 96 func (c *taskHookCoordinator) hasRunningMainTasks() bool { 97 return len(c.mainTasksRunning) > 0 98 } 99 100 func (c *taskHookCoordinator) hasPendingMainTasks() bool { 101 return len(c.mainTasksPending) > 0 102 } 103 104 func (c *taskHookCoordinator) startConditionForTask(task *structs.Task) <-chan struct{} { 105 if task.Lifecycle == nil { 106 return c.mainTaskCtx.Done() 107 } 108 109 switch task.Lifecycle.Hook { 110 case structs.TaskLifecycleHookPrestart: 111 // Prestart tasks start without checking status of other tasks 112 return c.closedCh 113 case structs.TaskLifecycleHookPoststart: 114 return c.poststartTaskCtx.Done() 115 case structs.TaskLifecycleHookPoststop: 116 return c.poststopTaskCtx.Done() 117 default: 118 // it should never have a lifecycle stanza w/o a hook, so report an error but allow the task to start normally 119 c.logger.Error("invalid lifecycle hook", "task", task.Name, "hook", task.Lifecycle.Hook) 120 return c.mainTaskCtx.Done() 121 } 122 } 123 124 // This is not thread safe! This must only be called from one thread per alloc runner. 125 func (c *taskHookCoordinator) taskStateUpdated(states map[string]*structs.TaskState) { 126 for task := range c.prestartSidecar { 127 st := states[task] 128 if st == nil || st.StartedAt.IsZero() { 129 continue 130 } 131 132 delete(c.prestartSidecar, task) 133 } 134 135 for task := range c.prestartEphemeral { 136 st := states[task] 137 if st == nil || !st.Successful() { 138 continue 139 } 140 141 delete(c.prestartEphemeral, task) 142 } 143 144 for task := range c.mainTasksRunning { 145 st := states[task] 146 147 if st == nil || st.State != structs.TaskStateDead { 148 continue 149 } 150 151 delete(c.mainTasksRunning, task) 152 } 153 154 for task := range c.mainTasksPending { 155 st := states[task] 156 if st == nil || st.StartedAt.IsZero() { 157 continue 158 } 159 160 delete(c.mainTasksPending, task) 161 } 162 163 if !c.hasPrestartTasks() { 164 c.mainTaskCtxCancel() 165 } 166 167 if !c.hasPendingMainTasks() { 168 c.poststartTaskCtxCancel() 169 } 170 if !c.hasRunningMainTasks() { 171 c.poststopTaskCtxCancel() 172 } 173 } 174 175 func (c *taskHookCoordinator) StartPoststopTasks() { 176 c.poststopTaskCtxCancel() 177 } 178 179 // hasNonSidecarTasks returns false if all the passed tasks are sidecar tasks 180 func hasNonSidecarTasks(tasks []*taskrunner.TaskRunner) bool { 181 for _, tr := range tasks { 182 lc := tr.Task().Lifecycle 183 if lc == nil || !lc.Sidecar { 184 return true 185 } 186 } 187 188 return false 189 } 190 191 // hasSidecarTasks returns true if all the passed tasks are sidecar tasks 192 func hasSidecarTasks(tasks map[string]*taskrunner.TaskRunner) bool { 193 for _, tr := range tasks { 194 lc := tr.Task().Lifecycle 195 if lc != nil && lc.Sidecar { 196 return true 197 } 198 } 199 200 return false 201 }