prestart-restart

client/allocrunner/alloc_runner.go

@@ -260,6 +260,7 @@ func (ar *allocRunner) initTaskRunners(tasks []*structs.Task) error {
 			DriverManager:        ar.driverManager,
 			ServersContactedCh:   ar.serversContactedCh,
 			StartConditionMetCtx: ar.taskHookCoordinator.startConditionForTask(task),
+			EndConditionMetCtx:   ar.taskHookCoordinator.endConditionForTask(task),
 		}
 
 		// Create, but do not Run, the task runner

client/allocrunner/task_hook_coordinator.go

@@ -121,6 +121,11 @@ func (c *taskHookCoordinator) startConditionForTask(task *structs.Task) <-chan s
 	}
 }
 
+// Tasks are able to exit the taskrunner Run() loop when poststop tasks are ready to start
+func (c *taskHookCoordinator) endConditionForTask(task *structs.Task) <-chan struct{} {
+	return c.poststopTaskCtx.Done()
+}
+
 // This is not thread safe! This must only be called from one thread per alloc runner.
 func (c *taskHookCoordinator) taskStateUpdated(states map[string]*structs.TaskState) {
 	for task := range c.prestartSidecar {

client/allocrunner/taskrunner/lifecycle.go

@@ -2,6 +2,7 @@ package taskrunner
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/hashicorp/nomad/nomad/structs"
 )
@@ -15,7 +16,8 @@ func (tr *TaskRunner) Restart(ctx context.Context, event *structs.TaskEvent, fai
 	handle := tr.getDriverHandle()
 
 	// Check it is running
-	if handle == nil {
+	if handle == nil && !tr.IsPrestartTask() {
+		fmt.Println("!!! task not running, not restarting")
 		return ErrTaskNotRunning
 	}
 

client/allocrunner/taskrunner/restarts/restarts.go

@@ -164,6 +164,7 @@ func (r *RestartTracker) GetState() (string, time.Duration) {
 
 	// Hot path if a restart was triggered
 	if r.restartTriggered {
+		fmt.Println("!!! restart triggered")
 		r.reason = ""
 		return structs.TaskRestarting, 0
 	}

client/allocrunner/taskrunner/task_runner.go

@@ -215,9 +215,16 @@ type TaskRunner struct {
 	// GetClientAllocs has been called in case of a failed restore.
 	serversContactedCh <-chan struct{}
 
-	// startConditionMetCtx is done when TR should start the task
+	// startConditionMetCtx is done when TaskRunner should start the task
+	// within the allocation lifecycle. This will allow the task to proceed
+	// runtime execution.
 	startConditionMetCtx <-chan struct{}
 
+	// endConditionMetCtx blocks tasks from exiting the taskrunner Run loop
+	// until the whole allocation is ready to finish. This allows prestart ephemeral
+	// tasks to be restarted successfully
+	endConditionMetCtx <-chan struct{}
+
 	// waitOnServers defaults to false but will be set true if a restore
 	// fails and the Run method should wait until serversContactedCh is
 	// closed.
@@ -278,6 +285,10 @@ type Config struct {
 
 	// startConditionMetCtx is done when TR should start the task
 	StartConditionMetCtx <-chan struct{}
+
+	// EndConditionMetCtx is done when TR can let the task exit the run loop
+	// i.e. when the whole allocation is ready to finish
+	EndConditionMetCtx <-chan struct{}
 }
 
 func NewTaskRunner(config *Config) (*TaskRunner, error) {
@@ -332,6 +343,7 @@ func NewTaskRunner(config *Config) (*TaskRunner, error) {
 		maxEvents:            defaultMaxEvents,
 		serversContactedCh:   config.ServersContactedCh,
 		startConditionMetCtx: config.StartConditionMetCtx,
+		endConditionMetCtx:   config.EndConditionMetCtx,
 	}
 
 	// Create the logger based on the allocation ID
@@ -598,6 +610,7 @@ MAIN:
 			tr.logger.Trace("gracefully shutting down during restart delay")
 			return
 		}
+
 	}
 
 	// Ensure handle is cleaned up. Restore could have recovered a task
@@ -623,6 +636,16 @@ MAIN:
 		tr.logger.Error("stop failed", "error", err)
 	}
 
+	// Block until the allocation is ready to finish
+	select {
+	case <-tr.endConditionMetCtx:
+		tr.logger.Debug("lifecycle end condition has been met, proceeding")
+		// yay proceed
+	case <-tr.killCtx.Done():
+	case <-tr.shutdownCtx.Done():
+		return
+	}
+
 	tr.logger.Debug("task run loop exiting")
 }
 

client/allocrunner/taskrunner/task_runner_getters.go

@@ -28,6 +28,11 @@ func (tr *TaskRunner) IsLeader() bool {
 	return tr.taskLeader
 }
 
+// IsPrestartTask returns true if this task is a prestart task in its task group.
+func (tr *TaskRunner) IsPrestartTask() bool {
+	return tr.Task().Lifecycle != nil && tr.Task().Lifecycle.Hook == structs.TaskLifecycleHookPrestart
+}
+
 // IsPoststopTask returns true if this task is a poststop task in its task group.
 func (tr *TaskRunner) IsPoststopTask() bool {
 	return tr.Task().Lifecycle != nil && tr.Task().Lifecycle.Hook == structs.TaskLifecycleHookPoststop