Merge pull request buildkite#24 from paulbes/deregister_task

Add flag for deregistering task
sherzberg · May 9, 2019 · e86ae03 · e86ae03
2 parents 33abf90 + 629cdfe
commit e86ae03
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ GLOBAL OPTIONS:
    --env KEY=value, -e KEY=value  An environment variable to add in the form KEY=value or `KEY` (shorthand for `KEY=$KEY` to pass through an env var from the current host). Can be specified multiple times
    --inherit-env, -E              Inherit all of the environment variables from the calling shell
    --count value, -C value        Number of tasks to run (default: 1)
+   --deregister                   Deregister task definition once done
    --help, -h                     show help
    --version, -v                  print the version
 ```
@@ -52,6 +53,7 @@ The following IAM permissions are required:
     - Effect: Allow
       Action:
         - ecs:RegisterTaskDefinition
+        - ecs:DeregisterTaskDefinition
         - ecs:RunTask
         - ecs:DescribeTasks
         - logs:DescribeLogGroups

diff --git a/go.mod b/go.mod
@@ -7,8 +7,7 @@ require (
 	github.com/ghodss/yaml v1.0.0
 	github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af // indirect
 	github.com/kr/pretty v0.1.0 // indirect
-	github.com/mitchellh/gox v0.4.0 // indirect
-	github.com/mitchellh/iochan v1.0.0 // indirect
+	github.com/mitchellh/gox v1.0.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/testify v1.2.2 // indirect
 	github.com/urfave/cli v1.20.0

diff --git a/go.sum b/go.sum
@@ -6,6 +6,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
+github.com/hashicorp/go-version v1.0.0 h1:21MVWPKDphxa7ineQQTrCU5brh7OuVVAzGOCnnCPtE8=
+github.com/hashicorp/go-version v1.0.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
 github.com/jmespath/go-jmespath v0.0.0-20160202185014-0b12d6b521d8/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af h1:pmfjZENx5imkbgOkpRUYLnmbU7UEFbjtDA2hxJ1ichM=
 github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k=
@@ -16,6 +18,8 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/mitchellh/gox v0.4.0 h1:lfGJxY7ToLJQjHHwi0EX6uYBdK78egf954SQl13PQJc=
 github.com/mitchellh/gox v0.4.0/go.mod h1:Sd9lOJ0+aimLBi73mGofS1ycjY8lL3uZM3JPS42BGNg=
+github.com/mitchellh/gox v1.0.1 h1:x0jD3dcHk9a9xPSDN6YEL4xL6Qz0dvNYm8yZqui5chI=
+github.com/mitchellh/gox v1.0.1/go.mod h1:ED6BioOGXMswlXa2zxfh/xdd5QhwYliBFn9V18Ap4z4=
 github.com/mitchellh/iochan v1.0.0 h1:C+X3KsSTLFVBr/tK1eYN/vs4rJcvsiLU338UhYPJWeY=
 github.com/mitchellh/iochan v1.0.0/go.mod h1:JwYml1nuB7xOzsp52dPpHFffvOCDupsG0QubkSMEySY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

diff --git a/main.go b/main.go
@@ -75,6 +75,10 @@ func main() {
 			Value: 1,
 			Usage: "Number of tasks to run",
 		},
+		cli.BoolFlag{
+			Name:  "deregister",
+			Usage: "Deregister task definition once done",
+		},
 	}
 
 	app.Action = func(ctx *cli.Context) error {
@@ -98,6 +102,7 @@ func main() {
 		r.Subnets = ctx.StringSlice("subnet")
 		r.Environment = ctx.StringSlice("env")
 		r.Count = ctx.Int64("count")
+		r.Deregister = ctx.Bool("deregister")
 
 		if ctx.Bool("inherit-env") {
 			for _, env := range os.Environ() {

diff --git a/runner/runner.go b/runner/runner.go
@@ -37,6 +37,7 @@ type Runner struct {
 	Subnets            []string
 	Environment        []string
 	Count              int64
+	Deregister         bool
 }
 
 func New() *Runner {
@@ -86,6 +87,22 @@ func (r *Runner) Run(ctx context.Context) error {
 	taskDefinition := fmt.Sprintf("%s:%d",
 		*resp.TaskDefinition.Family, *resp.TaskDefinition.Revision)
 
+	defer func() {
+		if !r.Deregister {
+			return
+		}
+
+		log.Printf("Deregistering task %s", taskDefinition)
+		_, err := svc.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
+			TaskDefinition: &taskDefinition,
+		})
+		if err != nil {
+			log.Printf("Failed to deregister task %s: %s", taskDefinition, err.Error())
+			return
+		}
+		log.Printf("Successfully deregistered task %s", taskDefinition)
+	}()
+
 	runTaskInput := &ecs.RunTaskInput{
 		TaskDefinition: aws.String(taskDefinition),
 		Cluster:        aws.String(r.Cluster),