github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/template_hook.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  
     8  	log "github.com/hashicorp/go-hclog"
     9  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    10  	ti "github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
    11  	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/template"
    12  	"github.com/hashicorp/nomad/client/config"
    13  	"github.com/hashicorp/nomad/client/taskenv"
    14  	"github.com/hashicorp/nomad/nomad/structs"
    15  )
    16  
    17  type templateHookConfig struct {
    18  	// logger is used to log
    19  	logger log.Logger
    20  
    21  	// lifecycle is used to interact with the task's lifecycle
    22  	lifecycle ti.TaskLifecycle
    23  
    24  	// events is used to emit events
    25  	events ti.EventEmitter
    26  
    27  	// templates is the set of templates we are managing
    28  	templates []*structs.Template
    29  
    30  	// clientConfig is the Nomad Client configuration
    31  	clientConfig *config.Config
    32  
    33  	// envBuilder is the environment variable builder for the task.
    34  	envBuilder *taskenv.Builder
    35  }
    36  
    37  type templateHook struct {
    38  	config *templateHookConfig
    39  
    40  	// logger is used to log
    41  	logger log.Logger
    42  
    43  	// templateManager is used to manage any consul-templates this task may have
    44  	templateManager *template.TaskTemplateManager
    45  	managerLock     sync.Mutex
    46  
    47  	// vaultToken is the current Vault token
    48  	vaultToken string
    49  
    50  	// taskDir is the task directory
    51  	taskDir string
    52  }
    53  
    54  func newTemplateHook(config *templateHookConfig) *templateHook {
    55  	h := &templateHook{
    56  		config: config,
    57  	}
    58  	h.logger = config.logger.Named(h.Name())
    59  	return h
    60  }
    61  
    62  func (*templateHook) Name() string {
    63  	return "template"
    64  }
    65  
    66  func (h *templateHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    67  	h.managerLock.Lock()
    68  	defer h.managerLock.Unlock()
    69  
    70  	// If we have already run prerun before exit early.
    71  	if h.templateManager != nil {
    72  		return nil
    73  	}
    74  
    75  	// Store the current Vault token and the task directory
    76  	h.taskDir = req.TaskDir.Dir
    77  	h.vaultToken = req.VaultToken
    78  	unblockCh, err := h.newManager()
    79  	if err != nil {
    80  		return err
    81  	}
    82  
    83  	// Wait for the template to render
    84  	select {
    85  	case <-ctx.Done():
    86  	case <-unblockCh:
    87  	}
    88  
    89  	return nil
    90  }
    91  
    92  func (h *templateHook) newManager() (unblock chan struct{}, err error) {
    93  	unblock = make(chan struct{})
    94  	m, err := template.NewTaskTemplateManager(&template.TaskTemplateManagerConfig{
    95  		UnblockCh:            unblock,
    96  		Lifecycle:            h.config.lifecycle,
    97  		Events:               h.config.events,
    98  		Templates:            h.config.templates,
    99  		ClientConfig:         h.config.clientConfig,
   100  		VaultToken:           h.vaultToken,
   101  		TaskDir:              h.taskDir,
   102  		EnvBuilder:           h.config.envBuilder,
   103  		MaxTemplateEventRate: template.DefaultMaxTemplateEventRate,
   104  	})
   105  	if err != nil {
   106  		h.logger.Error("failed to create template manager", "error", err)
   107  		return nil, err
   108  	}
   109  
   110  	h.templateManager = m
   111  	return unblock, nil
   112  }
   113  
   114  func (h *templateHook) Stop(ctx context.Context, req *interfaces.TaskStopRequest, resp *interfaces.TaskStopResponse) error {
   115  	h.managerLock.Lock()
   116  	defer h.managerLock.Unlock()
   117  
   118  	// Shutdown any created template
   119  	if h.templateManager != nil {
   120  		h.templateManager.Stop()
   121  	}
   122  
   123  	return nil
   124  }
   125  
   126  // Handle new Vault token
   127  func (h *templateHook) Update(ctx context.Context, req *interfaces.TaskUpdateRequest, resp *interfaces.TaskUpdateResponse) error {
   128  	h.managerLock.Lock()
   129  	defer h.managerLock.Unlock()
   130  
   131  	// Nothing to do
   132  	if h.templateManager == nil {
   133  		return nil
   134  	}
   135  
   136  	// Check if the Vault token has changed
   137  	if req.VaultToken == h.vaultToken {
   138  		return nil
   139  	} else {
   140  		h.vaultToken = req.VaultToken
   141  	}
   142  
   143  	// Shutdown the old template
   144  	h.templateManager.Stop()
   145  	h.templateManager = nil
   146  
   147  	// Create the new template
   148  	if _, err := h.newManager(); err != nil {
   149  		err := fmt.Errorf("failed to build template manager: %v", err)
   150  		h.logger.Error("failed to build template manager", "error", err)
   151  		h.config.lifecycle.Kill(context.Background(),
   152  			structs.NewTaskEvent(structs.TaskKilling).
   153  				SetFailsTask().
   154  				SetDisplayMessage(fmt.Sprintf("Template update %v", err)))
   155  	}
   156  
   157  	return nil
   158  }