github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/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  	// vaultNamespace is the current Vault namespace
    51  	vaultNamespace string
    52  
    53  	// taskDir is the task directory
    54  	taskDir string
    55  }
    56  
    57  func newTemplateHook(config *templateHookConfig) *templateHook {
    58  	h := &templateHook{
    59  		config: config,
    60  	}
    61  	h.logger = config.logger.Named(h.Name())
    62  	return h
    63  }
    64  
    65  func (*templateHook) Name() string {
    66  	return "template"
    67  }
    68  
    69  func (h *templateHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    70  	h.managerLock.Lock()
    71  	defer h.managerLock.Unlock()
    72  
    73  	// If we have already run prerun before exit early.
    74  	if h.templateManager != nil {
    75  		return nil
    76  	}
    77  
    78  	// Store the current Vault token and the task directory
    79  	h.taskDir = req.TaskDir.Dir
    80  	h.vaultToken = req.VaultToken
    81  
    82  	// Set vault namespace if specified
    83  	if req.Task.Vault != nil {
    84  		h.vaultNamespace = req.Task.Vault.Namespace
    85  	}
    86  
    87  	unblockCh, err := h.newManager()
    88  	if err != nil {
    89  		return err
    90  	}
    91  
    92  	// Wait for the template to render
    93  	select {
    94  	case <-ctx.Done():
    95  	case <-unblockCh:
    96  	}
    97  
    98  	return nil
    99  }
   100  
   101  func (h *templateHook) newManager() (unblock chan struct{}, err error) {
   102  	unblock = make(chan struct{})
   103  	m, err := template.NewTaskTemplateManager(&template.TaskTemplateManagerConfig{
   104  		UnblockCh:            unblock,
   105  		Lifecycle:            h.config.lifecycle,
   106  		Events:               h.config.events,
   107  		Templates:            h.config.templates,
   108  		ClientConfig:         h.config.clientConfig,
   109  		VaultToken:           h.vaultToken,
   110  		VaultNamespace:       h.vaultNamespace,
   111  		TaskDir:              h.taskDir,
   112  		EnvBuilder:           h.config.envBuilder,
   113  		MaxTemplateEventRate: template.DefaultMaxTemplateEventRate,
   114  	})
   115  	if err != nil {
   116  		h.logger.Error("failed to create template manager", "error", err)
   117  		return nil, err
   118  	}
   119  
   120  	h.templateManager = m
   121  	return unblock, nil
   122  }
   123  
   124  func (h *templateHook) Stop(ctx context.Context, req *interfaces.TaskStopRequest, resp *interfaces.TaskStopResponse) error {
   125  	h.managerLock.Lock()
   126  	defer h.managerLock.Unlock()
   127  
   128  	// Shutdown any created template
   129  	if h.templateManager != nil {
   130  		h.templateManager.Stop()
   131  	}
   132  
   133  	return nil
   134  }
   135  
   136  // Handle new Vault token
   137  func (h *templateHook) Update(ctx context.Context, req *interfaces.TaskUpdateRequest, resp *interfaces.TaskUpdateResponse) error {
   138  	h.managerLock.Lock()
   139  	defer h.managerLock.Unlock()
   140  
   141  	// Nothing to do
   142  	if h.templateManager == nil {
   143  		return nil
   144  	}
   145  
   146  	// Check if the Vault token has changed
   147  	if req.VaultToken == h.vaultToken {
   148  		return nil
   149  	} else {
   150  		h.vaultToken = req.VaultToken
   151  	}
   152  
   153  	// Shutdown the old template
   154  	h.templateManager.Stop()
   155  	h.templateManager = nil
   156  
   157  	// Create the new template
   158  	if _, err := h.newManager(); err != nil {
   159  		err := fmt.Errorf("failed to build template manager: %v", err)
   160  		h.logger.Error("failed to build template manager", "error", err)
   161  		h.config.lifecycle.Kill(context.Background(),
   162  			structs.NewTaskEvent(structs.TaskKilling).
   163  				SetFailsTask().
   164  				SetDisplayMessage(fmt.Sprintf("Template update %v", err)))
   165  	}
   166  
   167  	return nil
   168  }