address review: use Poststart hook instead of polling goroutine

Replace the runFirstRenderScripts polling goroutine with a
RunFirstRenderScripts public method that is called from the
template hook Poststart handler. The task runner framework
guarantees Poststart fires once the task reaches running state,
so the polling loop is unnecessary.

Also adds TestTaskTemplateManager_FirstRenderScript covering the
new code path.

Author: Vedant Madane

client/allocrunner/taskrunner/template/template.go

@@ -68,6 +68,10 @@ type TaskTemplateManager struct {
 	// shutdownCh is used to signal and started goroutine to shutdown
 	shutdownCh chan struct{}
 
+	// firstRenderScripts holds change_scripts that should fire after the
+	// initial template render once the task reaches the running state.
+	firstRenderScripts []*structs.ChangeScript
+
 	// shutdown marks whether the manager has been shutdown
 	shutdown     bool
 	shutdownLock sync.Mutex
@@ -276,16 +280,13 @@ func (tm *TaskTemplateManager) Run() {
 	// Unblock the task
 	close(tm.config.UnblockCh)
 
-	// Collect any change_script templates that should fire on first render.
-	// These scripts are launched in a goroutine because they need the task
-	// to reach the running state before Exec can succeed.
-	firstRenderScripts := tm.collectFirstRenderScripts()
-	if len(firstRenderScripts) > 0 {
-		go tm.runFirstRenderScripts(firstRenderScripts)
-	}
+	// Collect change_script templates that should fire on first render.
+	// These are stored and executed later when the task reaches the running
+	// state, triggered via RunFirstRenderScripts from the Poststart hook.
+	tm.firstRenderScripts = tm.collectFirstRenderScripts()
 
 	// If all our templates are change mode no-op, then we can exit here
-	if tm.allTemplatesNoop() && len(firstRenderScripts) == 0 {
+	if tm.allTemplatesNoop() && len(tm.firstRenderScripts) == 0 {
 		return
 	}
 
@@ -668,30 +669,22 @@ func (tm *TaskTemplateManager) collectFirstRenderScripts() []*structs.ChangeScri
 	return scripts
 }
 
-// runFirstRenderScripts waits for the task to reach the running state and then
-// executes the given change_scripts. It respects the shutdown channel so it
-// will not block forever if the task is killed before reaching running.
-func (tm *TaskTemplateManager) runFirstRenderScripts(scripts []*structs.ChangeScript) {
-	const pollInterval = 500 * time.Millisecond
-	ticker := time.NewTicker(pollInterval)
-	defer ticker.Stop()
+// RunFirstRenderScripts executes the change_scripts collected during the
+// first template render. It is called by the template hook's Poststart
+// method once the task is running and Exec is available.
+func (tm *TaskTemplateManager) RunFirstRenderScripts() {
+	scripts := tm.firstRenderScripts
+	if len(scripts) == 0 {
+		return
+	}
+	tm.firstRenderScripts = nil
 
-	for {
-		select {
-		case <-tm.shutdownCh:
-			return
-		case <-ticker.C:
-			if tm.config.Lifecycle.IsRunning() {
-				var wg sync.WaitGroup
-				for _, script := range scripts {
-					wg.Add(1)
-					go tm.processScript(script, &wg)
-				}
-				wg.Wait()
-				return
-			}
-		}
+	var wg sync.WaitGroup
+	for _, script := range scripts {
+		wg.Add(1)
+		go tm.processScript(script, &wg)
 	}
+	wg.Wait()
 }
 
 // templateRunner returns a consul-template runner for the given templates and a

client/allocrunner/taskrunner/template/template_test.go

@@ -1347,6 +1347,59 @@ OUTER:
 	}
 }
 
+// TestTaskTemplateManager_FirstRenderScript verifies that a template with
+// change_mode "script" and RunOnFirstRender collects the script so it can
+// be executed once the task reaches the running state via RunFirstRenderScripts.
+func TestTaskTemplateManager_FirstRenderScript(t *testing.T) {
+	ci.Parallel(t)
+	clienttestutil.RequireConsul(t)
+
+	key := "first_render_key"
+	t1 := &structs.Template{
+		EmbeddedTmpl: `FOO={{key "first_render_key"}}` + "\n",
+		DestPath:     "first_render.env",
+		ChangeMode:   structs.TemplateChangeModeScript,
+		ChangeScript: &structs.ChangeScript{
+			Command:          "/bin/foo",
+			Args:             []string{},
+			Timeout:          5 * time.Second,
+			FailOnError:      false,
+			RunOnFirstRender: true,
+		},
+		Envvars: true,
+	}
+
+	harness := newTestHarness(t, []*structs.Template{t1}, true, false)
+	harness.mockHooks.SetupExecTest(0, nil)
+	harness.start(t)
+	defer harness.stop()
+
+	// Write key so the template renders
+	harness.consul.SetKV(t, key, []byte("hello"))
+
+	// Wait for unblock (first render complete)
+	select {
+	case <-harness.mockHooks.UnblockCh:
+	case <-time.After(time.Duration(5*testutil.TestMultiplier()) * time.Second):
+		t.Fatal("Task unblock should have been called")
+	}
+
+	// Simulate the Poststart hook by calling RunFirstRenderScripts directly
+	harness.mockHooks.HasHandle = true
+	harness.manager.RunFirstRenderScripts()
+
+	// Verify script execution event was emitted
+	timeout := time.After(time.Duration(5*testutil.TestMultiplier()) * time.Second)
+	select {
+	case ev := <-harness.mockHooks.EmitEventCh:
+		if !strings.Contains(ev.DisplayMessage, t1.ChangeScript.Command) {
+			t.Fatalf("expected script event, got: %s", ev.DisplayMessage)
+		}
+	case <-timeout:
+		t.Fatal("should have received a script execution event")
+	}
+}
+
 // TestTaskTemplateManager_ScriptExecutionFailTask tests whether we fail the
 // task upon script execution failure if that's how it's configured.
 func TestTaskTemplateManager_ScriptExecutionFailTask(t *testing.T) {

client/allocrunner/taskrunner/template_hook.go

@@ -226,6 +226,16 @@ func (h *templateHook) newManager(tmpls []*structs.Template) (manager *template.
 	return m, unblock, nil
 }
 
+func (h *templateHook) Poststart(_ context.Context, _ *interfaces.TaskPoststartRequest, _ *interfaces.TaskPoststartResponse) error {
+	h.managerLock.Lock()
+	defer h.managerLock.Unlock()
+
+	if h.templateManager != nil {
+		h.templateManager.RunFirstRenderScripts()
+	}
+	return nil
+}
+
 func (h *templateHook) Stop(_ context.Context, req *interfaces.TaskStopRequest, resp *interfaces.TaskStopResponse) error {
 	h.managerLock.Lock()
 	defer h.managerLock.Unlock()