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  }