prevent panic when closing channels

aws/cloudformation.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -158,7 +159,9 @@ type StackWaiter struct {
 	cloudFormation *CloudFormation
 	errc           chan error
 	done           chan struct{}
+	loopDone       chan struct{}
 	events         chan types.StackEvent
+	closer         sync.Once
 }
 
 func newStackWaiter(stackName string, cf *CloudFormation) *StackWaiter {
@@ -167,6 +170,7 @@ func newStackWaiter(stackName string, cf *CloudFormation) *StackWaiter {
 		cloudFormation: cf,
 		errc:           make(chan error),
 		done:           make(chan struct{}),
+		loopDone:       make(chan struct{}),
 		events:         make(chan types.StackEvent),
 	}
 	go sw.pollEvents()
@@ -178,50 +182,35 @@ func (w *StackWaiter) Wait() error {
 }
 
 func (w *StackWaiter) close(err error) {
-	if w.isDone() {
-		return
-	}
-	close(w.done)
-	close(w.events)
-	w.errc <- err
+	w.closer.Do(func() {
+		close(w.done)
+		<-w.loopDone
+		w.errc <- err
+	})
 }
 
 func (w *StackWaiter) Events() <-chan types.StackEvent {
 	return w.events
 }
 
-func (w *StackWaiter) event(e types.StackEvent) {
-	if w.isDone() {
-		return
-	}
-	w.events <- e
-}
-
 func (w *StackWaiter) pollEvents() {
-	ts := time.Now()
 	ticker := time.NewTicker(time.Second)
+	ts := time.Now()
 	for {
 		select {
 		case <-ticker.C:
 			es, _ := w.cloudFormation.stackEvents(w.stackName, &ts)
 			for _, e := range es {
-				w.event(e)
+				w.events <- e
 			}
 			if len(es) > 0 {
 				ts = *es[0].Timestamp
 			}
 		case <-w.done:
 			ticker.Stop()
+			close(w.events)
+			close(w.loopDone)
 			return
 		}
 	}
 }
-
-func (w *StackWaiter) isDone() bool {
-	select {
-	case <-w.done:
-		return true
-	default:
-		return false
-	}
-}

cli/ui/progress/elements.go

@@ -3,6 +3,7 @@ package progress
 import (
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -17,6 +18,7 @@ type Dots struct {
 	currentCnt int
 	updateCh   chan struct{}
 	done       chan struct{}
+	closer     sync.Once
 }
 
 func NewDots() *Dots {
@@ -44,10 +46,9 @@ func (d *Dots) loop() {
 }
 
 func (d *Dots) Stop() {
-	if d.isDone() {
-		return
-	}
-	close(d.done)
+	d.closer.Do(func() {
+		close(d.done)
+	})
 }
 
 func (d *Dots) UpdateChan() <-chan struct{} {

cli/ui/progress/progress.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"reflect"
 	"strings"
+	"sync"
 
 	"github.com/fatih/color"
 )
@@ -16,6 +17,7 @@ type Progress struct {
 	loopDone   chan struct{}
 	printFunc  func(format string, v ...interface{})
 	isTerminal bool
+	closer     sync.Once
 }
 
 func New(prefix string, printFunc func(format string, v ...interface{}), elements ...Element) *Progress {
@@ -42,10 +44,9 @@ func (p *Progress) Run() {
 }
 
 func (p *Progress) Stop() {
-	if p.isDone() {
-		return
-	}
-	close(p.done)
+	p.closer.Do(func() {
+		close(p.done)
+	})
 	<-p.loopDone
 }
 

cli/ui/progress/terraform.go

@@ -10,14 +10,12 @@ type Terraform struct {
 	parser   *terraform.Parser
 	progress *Progress
 	counter  *Counter
-	done     chan struct{}
 }
 
 func NewTerraform() *Terraform {
 	parser := terraform.NewLogParser()
 	p := &Terraform{
 		parser: parser,
-		done:   make(chan struct{}),
 	}
 	return p
 }
@@ -42,7 +40,6 @@ func (p *Terraform) checkState(oldState terraform.ParserState) {
 		fmt.Println()
 	}
 	if newState == terraform.StateDone {
-		p.close()
 		return
 	}
 	p.initProgress()
@@ -66,19 +63,3 @@ func (p *Terraform) updateCounter() {
 	}
 	p.counter.SetCount(p.parser.CurrentResourceCount())
 }
-
-func (p *Terraform) close() {
-	if p.isDone() {
-		return
-	}
-	close(p.done)
-}
-
-func (p *Terraform) isDone() bool {
-	select {
-	case <-p.done:
-		return true
-	default:
-		return false
-	}
-}