From eea820ecf13011b640a195ffc80151f46db8fce2 Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Wed, 11 Jul 2018 15:42:35 -0700
Subject: [PATCH] tr: pass context to hooks

---
 client/allocrunnerv2/interfaces/task_lifecycle.go    | 12 +++++++-----
 client/allocrunnerv2/taskrunner/task_runner_hooks.go |  8 +++++---
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/client/allocrunnerv2/interfaces/task_lifecycle.go b/client/allocrunnerv2/interfaces/task_lifecycle.go
index f8ef85664..5876dd236 100644
--- a/client/allocrunnerv2/interfaces/task_lifecycle.go
+++ b/client/allocrunnerv2/interfaces/task_lifecycle.go
@@ -1,6 +1,8 @@
 package interfaces
 
 import (
+	"context"
+
 	"github.com/hashicorp/nomad/client/driver/env"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
@@ -49,7 +51,7 @@ type TaskPrerunResponse struct {
 
 type TaskPrerunHook interface {
 	TaskHook
-	Prerun(*TaskPrerunRequest, *TaskPrerunResponse) error
+	Prerun(context.Context, *TaskPrerunRequest, *TaskPrerunResponse) error
 }
 
 // XXX If we want consul style hooks, need to have something that runs after the
@@ -62,7 +64,7 @@ type TaskPostrunResponse struct{}
 type TaskPostrunHook interface {
 	TaskHook
 	Postrun() error
-	//Postrun(*TaskPostrunRequest, *TaskPostrunResponse) error
+	//Postrun(context.Context, *TaskPostrunRequest, *TaskPostrunResponse) error
 }
 
 type TaskPoststopRequest struct{}
@@ -71,7 +73,7 @@ type TaskPoststopResponse struct{}
 type TaskPoststopHook interface {
 	TaskHook
 	Postrun() error
-	//Postrun(*TaskPostrunRequest, *TaskPostrunResponse) error
+	//Postrun(context.Context, *TaskPostrunRequest, *TaskPostrunResponse) error
 }
 
 type TaskDestroyRequest struct{}
@@ -80,7 +82,7 @@ type TaskDestroyResponse struct{}
 type TaskDestroyHook interface {
 	TaskHook
 	Destroy() error
-	//Destroy(*TaskDestroyRequest, *TaskDestroyResponse) error
+	//Destroy(context.Context, *TaskDestroyRequest, *TaskDestroyResponse) error
 }
 
 type TaskUpdateRequest struct {
@@ -92,5 +94,5 @@ type TaskUpdateResponse struct{}
 type TaskUpdateHook interface {
 	TaskHook
 	Update() error
-	//Update(*TaskUpdateRequest, *TaskUpdateResponse) error
+	//Update(context.Context, *TaskUpdateRequest, *TaskUpdateResponse) error
 }
diff --git a/client/allocrunnerv2/taskrunner/task_runner_hooks.go b/client/allocrunnerv2/taskrunner/task_runner_hooks.go
index d3cc170b6..e0f25c5ee 100644
--- a/client/allocrunnerv2/taskrunner/task_runner_hooks.go
+++ b/client/allocrunnerv2/taskrunner/task_runner_hooks.go
@@ -1,6 +1,7 @@
 package taskrunner
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -79,7 +80,7 @@ func (tr *TaskRunner) prerun() error {
 
 		// Run the pre-run hook
 		var resp interfaces.TaskPrerunResponse
-		if err := pre.Prerun(&req, &resp); err != nil {
+		if err := pre.Prerun(tr.ctx, &req, &resp); err != nil {
 			return structs.WrapRecoverable(fmt.Sprintf("pre-run hook %q failed: %v", name, err), err)
 		}
 
@@ -211,7 +212,7 @@ func (h *taskDirHook) Name() string {
 	return "task_dir"
 }
 
-func (h *taskDirHook) Prerun(req *interfaces.TaskPrerunRequest, resp *interfaces.TaskPrerunResponse) error {
+func (h *taskDirHook) Prerun(ctx context.Context, req *interfaces.TaskPrerunRequest, resp *interfaces.TaskPrerunResponse) error {
 	cc := h.runner.clientConfig
 	chroot := cconfig.DefaultChrootEnv
 	if len(cc.ChrootEnv) > 0 {
@@ -256,10 +257,11 @@ func (*artifactHook) Name() string {
 	return "artifacts"
 }
 
-func (h *artifactHook) Prerun(req *interfaces.TaskPrerunRequest, resp *interfaces.TaskPrerunResponse) error {
+func (h *artifactHook) Prerun(ctx context.Context, req *interfaces.TaskPrerunRequest, resp *interfaces.TaskPrerunResponse) error {
 	h.eventEmitter.SetState(structs.TaskStatePending, structs.NewTaskEvent(structs.TaskDownloadingArtifacts))
 
 	for _, artifact := range req.Task.Artifacts {
+		//XXX add ctx to GetArtifact to allow cancelling long downloads
 		if err := getter.GetArtifact(req.TaskEnv, artifact, req.TaskDir); err != nil {
 			wrapped := fmt.Errorf("failed to download artifact %q: %v", artifact.GetterSource, err)
 			h.logger.Debug(wrapped.Error())