Repository: hashicorp/nomad
File: client/allocrunner/taskrunner/template_hook.go (Go)

Code before:

// Update is used to handle updates to vault and/or nomad tokens.
func (h *templateHook) Update(ctx context.Context, req *interfaces.TaskUpdateRequest, resp *interfaces.TaskUpdateResponse) error {
	h.managerLock.Lock()
	defer h.managerLock.Unlock()

	// no template manager to manage
	if h.templateManager == nil {
		return nil
	}

	// neither vault or nomad token has been updated, nothing to do
	if req.VaultToken == h.vaultToken && req.NomadToken == h.nomadToken {
		return nil
	} else {
		h.vaultToken = req.VaultToken
		h.nomadToken = req.NomadToken
	}

	tmpls := h.templateManager.Templates()

	// shutdown the old template
	h.templateManager.Stop()
	h.templateManager = nil

	err := h.renderTemplates(ctx, nil, tmpls)
	if err != nil {
		err = fmt.Errorf("failed to build template manager: %v", err)
		h.logger.Error("failed to build template manager", "error", err)
		_ = h.config.lifecycle.Kill(context.Background(),
			structs.NewTaskEvent(structs.TaskKilling).
				SetFailsTask().
				SetDisplayMessage(fmt.Sprintf("Template update %v", err)))
	}

	return nil
}

// renderTemplates creates the template managers and waits until each template has rendered, setting the watch
// templateManger on the hook when complete so it can be referenced during token updates.
func (h *templateHook) renderTemplates(ctx context.Context, once []*structs.Template, watch []*structs.Template) error {
	onceMgr, unblockOne, err := h.newManager(ctx, once)
	if err != nil {
		return err
	}

	watchMgr, unblockWatch, err := h.newManager(ctx, watch)
	if err != nil {
		return err
	}


Review comment (suggestion):
The assumption here is that we don't want to re-render templates with `once = true` on a task update.  This may or may not be correct.

Code after:

// Update is used to handle updates to vault and/or nomad tokens.
func (h *templateHook) Update(ctx context.Context, req *interfaces.TaskUpdateRequest, resp *interfaces.TaskUpdateResponse) error {
	h.managerLock.Lock()
	defer h.managerLock.Unlock()

	// no template manager to manage
	if h.templateManager == nil {
		return nil
	}

	// neither vault or nomad token has been updated, nothing to do
	if req.VaultToken == h.vaultToken && req.NomadToken == h.nomadToken {
		return nil
	} else {
		h.vaultToken = req.VaultToken
		h.nomadToken = req.NomadToken
	}

	tmpls := h.templateManager.Templates()

	// shutdown the old template
	h.templateManager.Stop()
	h.templateManager = nil

	err := h.renderTemplates(ctx, nil, tmpls)
	if err != nil {
		err = fmt.Errorf("failed to build template manager: %v", err)
		h.logger.Error("failed to build template manager", "error", err)
		_ = h.config.lifecycle.Kill(context.Background(),
			structs.NewTaskEvent(structs.TaskKilling).
				SetFailsTask().
				SetDisplayMessage(fmt.Sprintf("Template update %v", err)))
	}

	return nil
}

// renderTemplates creates the template managers and waits until each template has rendered, setting the watch
// templateManger on the hook when complete so it can be referenced during token updates.
func (h *templateHook) renderTemplates(ctx context.Context, once []*structs.Template, watch []*structs.Template) error {
	onceMgr, unblockOne, err := h.newManager(once)
	if err != nil {
		return err
	}

	watchMgr, unblockWatch, err := h.newManager(watch)
	if err != nil {
		return err
	}