Don't close channels (or range over them) when done.

This is groundwork for #20 and #23
glassechidna · May 12, 2019 · 614efda · 614efda
1 parent fcb3924
commit 614efda
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 28 deletions.
diff --git a/cmd/down.go b/cmd/down.go
@@ -38,10 +38,25 @@ var downCmd = &cobra.Command{
 
 		sess := awsSession(profile, region)
 		sit := stackit.NewStackit(cloudformation.New(sess), sts.New(sess))
-		go sit.Down(context.Background(), stackName, events)
 
-		for tailEvent := range events {
-			printer.PrintTailEvent(tailEvent)
+		ctx := context.Background()
+		printerCtx, printerCancel := context.WithCancel(ctx)
+		defer printerCancel()
+
+		go func() {
+			for {
+				select {
+				case <-printerCtx.Done():
+					return
+				case tailEvent := <-events:
+					printer.PrintTailEvent(tailEvent)
+				}
+			}
+		}()
+
+		err := sit.Down(ctx, stackName, events)
+		if err != nil {
+			panic(err)
 		}
 	},
 }

diff --git a/cmd/execute.go b/cmd/execute.go
@@ -31,9 +31,17 @@ func executeChangeSet(ctx context.Context, region, profile, stackName, changeSet
 	events := make(chan stackit.TailStackEvent)
 	printer := stackit.NewTailPrinterWithOptions(true, true, writer)
 
+	printerCtx, printerCancel := context.WithCancel(ctx)
+	defer printerCancel()
+
 	go func() {
-		for event := range events {
-			printer.PrintTailEvent(event)
+		for {
+			select {
+			case <-printerCtx.Done():
+				return
+			case tailEvent := <-events:
+				printer.PrintTailEvent(tailEvent)
+			}
 		}
 	}()
 

diff --git a/cmd/package.go b/cmd/package.go
@@ -48,9 +48,17 @@ func packageTemplate(ctx context.Context, region, profile, stackName, templatePa
 	events := make(chan stackit.TailStackEvent)
 	printer := stackit.NewTailPrinterWithOptions(true, true, writer)
 
+	printerCtx, printerCancel := context.WithCancel(ctx)
+	defer printerCancel()
+
 	go func() {
-		for event := range events {
-			printer.PrintTailEvent(event)
+		for {
+			select {
+			case <-printerCtx.Done():
+				return
+			case tailEvent := <-events:
+				printer.PrintTailEvent(tailEvent)
+			}
 		}
 	}()
 

diff --git a/cmd/up.go b/cmd/up.go
@@ -69,13 +69,23 @@ var upCmd = &cobra.Command{
 		sess := awsSession(profile, region)
 		sit := stackit.NewStackit(cloudformation.New(sess), sts.New(sess))
 
+		ctx := context.Background()
+
+		printerCtx, printerCancel := context.WithCancel(ctx)
+		defer printerCancel()
+
 		go func() {
-			for tailEvent := range events {
-				printer.PrintTailEvent(tailEvent)
+			for {
+				select {
+				case <-printerCtx.Done():
+					return
+				case tailEvent := <-events:
+					printer.PrintTailEvent(tailEvent)
+				}
 			}
 		}()
 
-		prepared, err := sit.Prepare(context.Background(), parsed, events)
+		prepared, err := sit.Prepare(ctx, parsed, events)
 		if err != nil {
 			panic(err)
 		}
@@ -84,7 +94,7 @@ var upCmd = &cobra.Command{
 			return // no-op change set
 		}
 
-		err = sit.Execute(context.Background(), *prepared.Output.StackId, *prepared.Output.Id, events)
+		err = sit.Execute(ctx, *prepared.Output.StackId, *prepared.Output.Id, events)
 		if err != nil {
 			panic(err)
 		}

diff --git a/pkg/stackit/down.go b/pkg/stackit/down.go
@@ -18,15 +18,13 @@ func (s *Stackit) Down(ctx context.Context, stackName string, events chan<- Tail
 		}
 		_, err = s.api.DeleteStack(input)
 		if err != nil {
-			close(events)
 			return err
 		}
 
 		finalEvent, err := s.PollStackEvents(ctx, *stack.StackId, token, func(event TailStackEvent) {
 			events <- event
 		})
 		if err != nil {
-			close(events)
 			return err
 		}
 
@@ -35,32 +33,28 @@ func (s *Stackit) Down(ctx context.Context, stackName string, events chan<- Tail
 			input.ClientRequestToken = &token
 			input.RetainResources, err = s.resourcesToBeRetainedDuringDelete(*stack.StackId, events)
 			if err != nil {
-				close(events)
 				return errors.Wrap(err, "determining resources to be kept")
 			}
 
 			_, err = s.api.DeleteStack(input)
 			if err != nil {
-				close(events)
 				return errors.Wrap(err, "deleting stack")
 			}
 
 			_, err = s.PollStackEvents(ctx, *stack.StackId, token, func(event TailStackEvent) {
 				events <- event
 			})
 			if err != nil {
-				close(events)
 				return errors.Wrap(err, "deleting stack")
 			}
 		}
 	}
 
-	close(events)
 	return nil
 }
 
 func (s *Stackit) resourcesToBeRetainedDuringDelete(stackName string, events chan<- TailStackEvent) ([]*string, error) {
-	names := []*string{}
+	var names []*string
 
 	err := s.api.ListStackResourcesPages(&cloudformation.ListStackResourcesInput{StackName: &stackName}, func(page *cloudformation.ListStackResourcesOutput, lastPage bool) bool {
 		for _, resource := range page.StackResourceSummaries {
@@ -70,10 +64,6 @@ func (s *Stackit) resourcesToBeRetainedDuringDelete(stackName string, events cha
 		}
 		return !lastPage
 	})
-	if err != nil {
-		close(events)
-		return nil, err
-	}
 
-	return names, nil
+	return names, errors.Wrap(err, "listing stack resources")
 }
diff --git a/pkg/stackit/up.go b/pkg/stackit/up.go
@@ -73,7 +73,6 @@ func (s *Stackit) ensureStackReady(ctx context.Context, stackName string, events
 		token := generateToken()
 		_, err := s.api.DeleteStack(&cloudformation.DeleteStackInput{StackName: &stackId, ClientRequestToken: &token})
 		if err != nil {
-			close(events)
 			return err
 		}
 
@@ -225,14 +224,12 @@ func (s *Stackit) Execute(ctx context.Context, stackId, changeSetId string, even
 	})
 
 	if err != nil {
-		close(events)
 		return errors.Wrap(err, "executing change set")
 	}
 
 	_, err = s.PollStackEvents(ctx, stackId, token, func(event TailStackEvent) {
 		events <- event
 	})
 
-	close(events)
 	return nil
 }
diff --git a/pkg/stackit/up_test.go b/pkg/stackit/up_test.go
@@ -1,6 +1,7 @@
 package stackit
 
 import (
+	"context"
 	"errors"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
@@ -45,7 +46,7 @@ func TestServiceRoleArnCanBeName(t *testing.T) {
 	}
 
 	ch := make(chan TailStackEvent)
-	_, err := s.Prepare(input, ch)
+	_, err := s.Prepare(context.Background(), input, ch)
 	assert.EqualError(t, err, "creating change set: done")
 }
 
@@ -83,6 +84,6 @@ func TestServiceRoleArnDoesntTriggerStsCall(t *testing.T) {
 	}
 
 	ch := make(chan TailStackEvent)
-	_, err := s.Prepare(input, ch)
+	_, err := s.Prepare(context.Background(), input, ch)
 	assert.EqualError(t, err, "creating change set: done")
 }