github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/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/nomad/structs"
     8  )
     9  
    10  // TaskHookCoordinator helps coordinate when main start tasks can launch
    11  // namely after all Prestart Tasks have run, and after all BlockUntilCompleted have completed
    12  type taskHookCoordinator struct {
    13  	logger hclog.Logger
    14  
    15  	closedCh chan struct{}
    16  
    17  	mainTaskCtx       context.Context
    18  	mainTaskCtxCancel func()
    19  
    20  	prestartSidecar   map[string]struct{}
    21  	prestartEphemeral map[string]struct{}
    22  }
    23  
    24  func newTaskHookCoordinator(logger hclog.Logger, tasks []*structs.Task) *taskHookCoordinator {
    25  	closedCh := make(chan struct{})
    26  	close(closedCh)
    27  
    28  	mainTaskCtx, cancelFn := context.WithCancel(context.Background())
    29  
    30  	c := &taskHookCoordinator{
    31  		logger:            logger,
    32  		closedCh:          closedCh,
    33  		mainTaskCtx:       mainTaskCtx,
    34  		mainTaskCtxCancel: cancelFn,
    35  		prestartSidecar:   map[string]struct{}{},
    36  		prestartEphemeral: map[string]struct{}{},
    37  	}
    38  	c.setTasks(tasks)
    39  	return c
    40  }
    41  
    42  func (c *taskHookCoordinator) setTasks(tasks []*structs.Task) {
    43  	for _, task := range tasks {
    44  
    45  		if task.Lifecycle == nil {
    46  			// move nothing
    47  			continue
    48  		}
    49  
    50  		switch task.Lifecycle.Hook {
    51  		case structs.TaskLifecycleHookPrestart:
    52  			if task.Lifecycle.Sidecar {
    53  				c.prestartSidecar[task.Name] = struct{}{}
    54  			} else {
    55  				c.prestartEphemeral[task.Name] = struct{}{}
    56  			}
    57  
    58  		default:
    59  			c.logger.Error("invalid lifecycle hook", "hook", task.Lifecycle.Hook)
    60  		}
    61  	}
    62  
    63  	if !c.hasPrestartTasks() {
    64  		c.mainTaskCtxCancel()
    65  	}
    66  }
    67  
    68  func (c *taskHookCoordinator) hasPrestartTasks() bool {
    69  	return len(c.prestartSidecar)+len(c.prestartEphemeral) > 0
    70  }
    71  
    72  func (c *taskHookCoordinator) startConditionForTask(task *structs.Task) <-chan struct{} {
    73  	if task.Lifecycle != nil && task.Lifecycle.Hook == structs.TaskLifecycleHookPrestart {
    74  		return c.closedCh
    75  	}
    76  
    77  	return c.mainTaskCtx.Done()
    78  
    79  }
    80  
    81  // This is not thread safe! This must only be called from one thread per alloc runner.
    82  func (c *taskHookCoordinator) taskStateUpdated(states map[string]*structs.TaskState) {
    83  	if c.mainTaskCtx.Err() != nil {
    84  		// nothing to do here
    85  		return
    86  	}
    87  
    88  	for task := range c.prestartSidecar {
    89  		st := states[task]
    90  		if st == nil || st.StartedAt.IsZero() {
    91  			continue
    92  		}
    93  
    94  		delete(c.prestartSidecar, task)
    95  	}
    96  
    97  	for task := range c.prestartEphemeral {
    98  		st := states[task]
    99  		if st == nil || !st.Successful() {
   100  			continue
   101  		}
   102  
   103  		delete(c.prestartEphemeral, task)
   104  	}
   105  
   106  	// everything well
   107  	if !c.hasPrestartTasks() {
   108  		c.mainTaskCtxCancel()
   109  	}
   110  }