diff --git a/client/allocrunner/alloc_runner.go b/client/allocrunner/alloc_runner.go index 526cb3e3bb37..99acc303a236 100644 --- a/client/allocrunner/alloc_runner.go +++ b/client/allocrunner/alloc_runner.go @@ -1138,6 +1138,9 @@ func (ar *allocRunner) Restart(ctx context.Context, event *structs.TaskEvent, fa var err *multierror.Error var errMutex sync.Mutex + // run alloc task restart hooks + ar.taskRestartHooks() + go func() { var wg sync.WaitGroup defer close(waitCh) @@ -1145,6 +1148,7 @@ func (ar *allocRunner) Restart(ctx context.Context, event *structs.TaskEvent, fa wg.Add(1) go func(taskName string, r agentconsul.WorkloadRestarter) { defer wg.Done() + e := r.Restart(ctx, event, failure) if e != nil { errMutex.Lock() @@ -1170,6 +1174,9 @@ func (ar *allocRunner) Restart(ctx context.Context, event *structs.TaskEvent, fa func (ar *allocRunner) RestartAll(taskEvent *structs.TaskEvent) error { var err *multierror.Error + // run alloc task restart hooks + ar.taskRestartHooks() + for tn := range ar.tasks { rerr := ar.RestartTask(tn, taskEvent.Copy()) if rerr != nil { diff --git a/client/allocrunner/alloc_runner_hooks.go b/client/allocrunner/alloc_runner_hooks.go index df797e44c196..2ad59febfad9 100644 --- a/client/allocrunner/alloc_runner_hooks.go +++ b/client/allocrunner/alloc_runner_hooks.go @@ -360,3 +360,28 @@ func (ar *allocRunner) shutdownHooks() { } } } + +func (ar *allocRunner) taskRestartHooks() { + for _, hook := range ar.runnerHooks { + re, ok := hook.(interfaces.RunnerTaskRestartHook) + if !ok { + continue + } + + name := re.Name() + var start time.Time + if ar.logger.IsTrace() { + start = time.Now() + ar.logger.Trace("running alloc task restart hook", + "name", name, "start", start) + } + + re.PreTaskRestart() + + if ar.logger.IsTrace() { + end := time.Now() + ar.logger.Trace("finished alloc task restart hook", + "name", name, "end", end, "duration", end.Sub(start)) + } + } +} diff --git a/client/allocrunner/groupservice_hook.go b/client/allocrunner/groupservice_hook.go index 616e7f2cf69b..66fd09e29e31 100644 --- a/client/allocrunner/groupservice_hook.go +++ b/client/allocrunner/groupservice_hook.go @@ -145,6 +145,14 @@ func (h *groupServiceHook) Update(req *interfaces.RunnerUpdateRequest) error { return h.consulClient.UpdateWorkload(oldWorkloadServices, newWorkloadServices) } +func (h *groupServiceHook) PreTaskRestart() error { + // TODO: refactor these hooks so that we can call both of them with a + // single mutex to lock both, so that we don't end up having a narrow race + // condition for Update calls + h.PreKill() + return h.Prerun() +} + func (h *groupServiceHook) PreKill() { h.mu.Lock() defer h.mu.Unlock() diff --git a/client/allocrunner/interfaces/runner_lifecycle.go b/client/allocrunner/interfaces/runner_lifecycle.go index 33713b2c1001..7855deaa3f4d 100644 --- a/client/allocrunner/interfaces/runner_lifecycle.go +++ b/client/allocrunner/interfaces/runner_lifecycle.go @@ -55,6 +55,14 @@ type RunnerUpdateRequest struct { Alloc *structs.Allocation } +// RunnerTaskRestartHooks are executed just before the allocation runner is +// going to restart all tasks. +type RunnerTaskRestartHook interface { + RunnerHook + + PreTaskRestart() error +} + // ShutdownHook may be implemented by AllocRunner or TaskRunner hooks and will // be called when the agent process is being shutdown gracefully. type ShutdownHook interface {