template: factor onTemplateRendered out of template re-renderer

Pull this function out so it can be used by the first-pass renderer when we restore a task from a restarted client.
hashicorp · Dec 11, 2020 · 815f01c · 815f01c
1 parent 863c194
commit 815f01c
Showing 1 changed file with 93 additions and 88 deletions.
diff --git a/client/allocrunner/taskrunner/template/template.go b/client/allocrunner/taskrunner/template/template.go
@@ -368,112 +368,117 @@ func (tm *TaskTemplateManager) handleTemplateRerenders(allRenderedTime time.Time
 					SetFailsTask().
 					SetDisplayMessage(fmt.Sprintf("Template failed: %v", err)))
 		case <-tm.runner.TemplateRenderedCh():
-			// A template has been rendered, figure out what to do
-			var handling []string
-			signals := make(map[string]struct{})
-			restart := false
-			var splay time.Duration
+			tm.onTemplateRendered(handledRenders, allRenderedTime)
+		}
+	}
+}
 
-			events := tm.runner.RenderEvents()
-			for id, event := range events {
+func (tm *TaskTemplateManager) onTemplateRendered(handledRenders map[string]time.Time, allRenderedTime time.Time) {
+	// A template has been rendered, figure out what to do
+	var handling []string
+	signals := make(map[string]struct{})
+	restart := false
+	var splay time.Duration
 
-				// First time through
-				if allRenderedTime.After(event.LastDidRender) || allRenderedTime.Equal(event.LastDidRender) {
-					handledRenders[id] = allRenderedTime
-					continue
-				}
+	events := tm.runner.RenderEvents()
+	for id, event := range events {
 
-				// We have already handled this one
-				if htime := handledRenders[id]; htime.After(event.LastDidRender) || htime.Equal(event.LastDidRender) {
-					continue
-				}
+		// First time through
+		if allRenderedTime.After(event.LastDidRender) || allRenderedTime.Equal(event.LastDidRender) {
+			handledRenders[id] = allRenderedTime
+			continue
+		}
 
-				// Lookup the template and determine what to do
-				tmpls, ok := tm.lookup[id]
-				if !ok {
-					tm.config.Lifecycle.Kill(context.Background(),
-						structs.NewTaskEvent(structs.TaskKilling).
-							SetFailsTask().
-							SetDisplayMessage(fmt.Sprintf("Template runner returned unknown template id %q", id)))
-					return
-				}
+		// We have already handled this one
+		if htime := handledRenders[id]; htime.After(event.LastDidRender) || htime.Equal(event.LastDidRender) {
+			continue
+		}
 
-				// Read environment variables from templates
-				envMap, err := loadTemplateEnv(tm.config.Templates, tm.config.TaskDir, tm.config.EnvBuilder.Build())
-				if err != nil {
-					tm.config.Lifecycle.Kill(context.Background(),
-						structs.NewTaskEvent(structs.TaskKilling).
-							SetFailsTask().
-							SetDisplayMessage(fmt.Sprintf("Template failed to read environment variables: %v", err)))
-					return
-				}
-				tm.config.EnvBuilder.SetTemplateEnv(envMap)
-
-				for _, tmpl := range tmpls {
-					switch tmpl.ChangeMode {
-					case structs.TemplateChangeModeSignal:
-						signals[tmpl.ChangeSignal] = struct{}{}
-					case structs.TemplateChangeModeRestart:
-						restart = true
-					case structs.TemplateChangeModeNoop:
-						continue
-					}
+		// Lookup the template and determine what to do
+		tmpls, ok := tm.lookup[id]
+		if !ok {
+			tm.config.Lifecycle.Kill(context.Background(),
+				structs.NewTaskEvent(structs.TaskKilling).
+					SetFailsTask().
+					SetDisplayMessage(fmt.Sprintf("Template runner returned unknown template id %q", id)))
+			return
+		}
 
-					if tmpl.Splay > splay {
-						splay = tmpl.Splay
-					}
-				}
+		// Read environment variables from templates
+		envMap, err := loadTemplateEnv(tm.config.Templates, tm.config.TaskDir, tm.config.EnvBuilder.Build())
+		if err != nil {
+			tm.config.Lifecycle.Kill(context.Background(),
+				structs.NewTaskEvent(structs.TaskKilling).
+					SetFailsTask().
+					SetDisplayMessage(fmt.Sprintf("Template failed to read environment variables: %v", err)))
+			return
+		}
+		tm.config.EnvBuilder.SetTemplateEnv(envMap)
+
+		for _, tmpl := range tmpls {
+			switch tmpl.ChangeMode {
+			case structs.TemplateChangeModeSignal:
+				signals[tmpl.ChangeSignal] = struct{}{}
+			case structs.TemplateChangeModeRestart:
+				restart = true
+			case structs.TemplateChangeModeNoop:
+				continue
+			}
 
-				handling = append(handling, id)
+			if tmpl.Splay > splay {
+				splay = tmpl.Splay
 			}
+		}
 
-			if restart || len(signals) != 0 {
-				if splay != 0 {
-					ns := splay.Nanoseconds()
-					offset := rand.Int63n(ns)
-					t := time.Duration(offset)
+		handling = append(handling, id)
+	}
 
-					select {
-					case <-time.After(t):
-					case <-tm.shutdownCh:
-						return
-					}
-				}
+	if restart || len(signals) != 0 {
+		if splay != 0 {
+			ns := splay.Nanoseconds()
+			offset := rand.Int63n(ns)
+			t := time.Duration(offset)
 
-				// Update handle time
-				for _, id := range handling {
-					handledRenders[id] = events[id].LastDidRender
-				}
+			select {
+			case <-time.After(t):
+			case <-tm.shutdownCh:
+				return
+			}
+		}
 
-				if restart {
-					tm.config.Lifecycle.Restart(context.Background(),
-						structs.NewTaskEvent(structs.TaskRestartSignal).
-							SetDisplayMessage("Template with change_mode restart re-rendered"), false)
-				} else if len(signals) != 0 {
-					var mErr multierror.Error
-					for signal := range signals {
-						s := tm.signals[signal]
-						event := structs.NewTaskEvent(structs.TaskSignaling).SetTaskSignal(s).SetDisplayMessage("Template re-rendered")
-						if err := tm.config.Lifecycle.Signal(event, signal); err != nil {
-							multierror.Append(&mErr, err)
-						}
-					}
+		// Update handle time
+		for _, id := range handling {
+			handledRenders[id] = events[id].LastDidRender
+		}
 
-					if err := mErr.ErrorOrNil(); err != nil {
-						flat := make([]os.Signal, 0, len(signals))
-						for signal := range signals {
-							flat = append(flat, tm.signals[signal])
-						}
+		if restart {
+			tm.config.Lifecycle.Restart(context.Background(),
+				structs.NewTaskEvent(structs.TaskRestartSignal).
+					SetDisplayMessage("Template with change_mode restart re-rendered"), false)
+		} else if len(signals) != 0 {
+			var mErr multierror.Error
+			for signal := range signals {
+				s := tm.signals[signal]
+				event := structs.NewTaskEvent(structs.TaskSignaling).SetTaskSignal(s).SetDisplayMessage("Template re-rendered")
+				if err := tm.config.Lifecycle.Signal(event, signal); err != nil {
+					multierror.Append(&mErr, err)
+				}
+			}
 
-						tm.config.Lifecycle.Kill(context.Background(),
-							structs.NewTaskEvent(structs.TaskKilling).
-								SetFailsTask().
-								SetDisplayMessage(fmt.Sprintf("Template failed to send signals %v: %v", flat, err)))
-					}
+			if err := mErr.ErrorOrNil(); err != nil {
+				flat := make([]os.Signal, 0, len(signals))
+				for signal := range signals {
+					flat = append(flat, tm.signals[signal])
 				}
+
+				tm.config.Lifecycle.Kill(context.Background(),
+					structs.NewTaskEvent(structs.TaskKilling).
+						SetFailsTask().
+						SetDisplayMessage(fmt.Sprintf("Template failed to send signals %v: %v", flat, err)))
 			}
 		}
 	}
+
 }
 
 // allTemplatesNoop returns whether all the managed templates have change mode noop.