lifecycle: add blocking channel to enable restart of prestart tasks

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/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")
 }