graceful shutdown of goroutines

3 files changed, 72 insertions, 27 deletions

diff --git a/focus.go b/focus.go
index 331d232..5f19e6c 100644
--- a/focus.go
+++ b/focus.go
@@ -19,6 +19,14 @@ func NewFocus(rows []int) Focus {
 	return f
 }
 
+func (f *Focus) Close() {
+	for i := range f.widgets {
+		for j := range f.widgets[i] {
+			close(f.widgets[i][j])
+		}
+	}
+}
+
 func (f *Focus) Left() {
 	f.widgets[f.p.Y][f.p.X] <- false
 	if f.p.X <= 0 {
diff --git a/gui/widget/widget.go b/gui/widget/widget.go
index cbc837e..f9fe658 100644
--- a/gui/widget/widget.go
+++ b/gui/widget/widget.go
@@ -3,6 +3,7 @@ package widget
 import (
 	"cmp"
 	"fmt"
+	"sync"
 
 	"image"
 	"image/color"
@@ -19,11 +20,15 @@ var (
 	WHITE       = color.Gray{255}
 )
 
-func Label(text string, r image.Rectangle, env gui.Env) {
+func Label(text string, r image.Rectangle, env gui.Env, wg *sync.WaitGroup) {
+	defer wg.Done()
+	defer close(env.Draw())
+
 	redraw := func(drw draw.Image) image.Rectangle {
 		drawText([]byte(text), drw, r, BLACK, WHITE)
 		return r
 	}
+
 	env.Draw() <- redraw
 	for event := range env.Events() {
 		switch event := event.(type) {
@@ -33,10 +38,12 @@ func Label(text string, r image.Rectangle, env gui.Env) {
 			}
 		}
 	}
-	close(env.Draw())
 }
 
-func Input(val chan<- uint, r image.Rectangle, focusChan <-chan bool, env gui.Env) {
+func Input(val chan<- uint, r image.Rectangle, focusChan <-chan bool, env gui.Env, wg *sync.WaitGroup) {
+	defer wg.Done()
+	defer close(env.Draw())
+
 	redraw := func(text []byte, focus bool) func(draw.Image) image.Rectangle {
 		return func(drw draw.Image) image.Rectangle {
 			if focus {
@@ -51,12 +58,15 @@ func Input(val chan<- uint, r image.Rectangle, focusChan <-chan bool, env gui.En
 	focus := false
 
 	env.Draw() <- redraw(text, focus)
-
+Loop:
 	for {
 		select {
 		case focus = <-focusChan:
 			env.Draw() <- redraw(text, focus)
-		case event := <-env.Events():
+		case event, ok := <-env.Events():
+			if !ok { // channel closed
+				break Loop
+			}
 			switch event := event.(type) {
 			case win.WiFocus:
 				if event.Focused {
@@ -77,10 +87,12 @@ func Input(val chan<- uint, r image.Rectangle, focusChan <-chan bool, env gui.En
 			}
 		}
 	}
-	close(env.Draw())
 }
 
-func Output(val <-chan uint, r image.Rectangle, env gui.Env) {
+func Output(val <-chan uint, r image.Rectangle, env gui.Env, wg *sync.WaitGroup) {
+	defer wg.Done()
+	defer close(env.Draw())
+
 	redraw := func(n uint) func(draw.Image) image.Rectangle {
 		return func(drw draw.Image) image.Rectangle {
 			drawText([]byte(fmt.Sprint(n)), drw, r, BLACK, WHITE)
@@ -89,8 +101,8 @@ func Output(val <-chan uint, r image.Rectangle, env gui.Env) {
 	}
 
 	var n uint = 0
-	env.Draw() <- redraw(n)
 
+	env.Draw() <- redraw(n)
 Loop:
 	for {
 		select {
@@ -105,7 +117,6 @@ Loop:
 			}
 		}
 	}
-	close(env.Draw())
 }
 
 func isDigit(r rune) bool {
diff --git a/main.go b/main.go
index 7e7f21d..d0efc6c 100644
--- a/main.go
+++ b/main.go
@@ -1,8 +1,11 @@
 package main
 
 import (
+	"fmt"
 	"image"
 	"image/color"
+	"os"
+	"sync"
 
 	"github.com/faiface/mainthread"
 	"volute/gui"
@@ -22,28 +25,40 @@ const (
 )
 
 func run() {
-	w, err := win.New(win.Title("volute"), win.Size(WIDTH, HEIGHT))
-	if err != nil {
-		panic(err)
-	}
-	mux, env := gui.NewMux(w)
-
 	var (
+		wg = new(sync.WaitGroup)
+
+		focus = NewFocus([]int{1, POINTS, POINTS, POINTS, POINTS})
+
 		displacementChan = make(chan uint)
 
 		rpmChan  [POINTS]chan uint
 		veChan   [POINTS]chan uint
 		imapChan [POINTS]chan uint
 		actChan  [POINTS]chan uint
-
-		focus = NewFocus([]int{1, POINTS, POINTS, POINTS, POINTS})
 	)
+	defer wg.Wait()
+	defer focus.Close()
+	defer close(displacementChan)
 	for i := 0; i < POINTS; i++ {
 		rpmChan[i] = make(chan uint)
 		veChan[i] = make(chan uint)
 		imapChan[i] = make(chan uint)
 		actChan[i] = make(chan uint)
+
+		defer close(rpmChan[i])
+		defer close(veChan[i])
+		defer close(imapChan[i])
+		defer close(actChan[i])
+	}
+
+	w, err := win.New(win.Title("volute"), win.Size(WIDTH, HEIGHT))
+	if err != nil {
+		fmt.Println("error creating window:", err)
+		os.Exit(1)
 	}
+	mux, env := gui.NewMux(w)
+	defer close(env.Draw())
 
 	bounds := layout.Grid{
 		Rows:        []int{2, 7, 7, 7, 7},
@@ -57,41 +72,56 @@ func run() {
 		Flip:        false,
 	}.Lay(image.Rect(0, 0, WIDTH, HEIGHT))
 
-	go widget.Label("displacement (cc)", bounds[0], mux.MakeEnv())
+	wg.Add(1)
+	go widget.Label("displacement (cc)", bounds[0], mux.MakeEnv(), wg)
+	wg.Add(1)
 	go widget.Input(
 		displacementChan,
 		bounds[1],
 		focus.widgets[0][0],
 		mux.MakeEnv(),
+		wg,
 	)
-	go widget.Label("speed (rpm)", bounds[2], mux.MakeEnv())
-	go widget.Label("VE (%)", bounds[3+POINTS], mux.MakeEnv())
-	go widget.Label("IMAP (mbar)", bounds[4+2*POINTS], mux.MakeEnv())
-	go widget.Label("ACT (*C)", bounds[5+3*POINTS], mux.MakeEnv())
+	wg.Add(1)
+	go widget.Label("speed (rpm)", bounds[2], mux.MakeEnv(), wg)
+	wg.Add(1)
+	go widget.Label("VE (%)", bounds[3+POINTS], mux.MakeEnv(), wg)
+	wg.Add(1)
+	go widget.Label("IMAP (mbar)", bounds[4+2*POINTS], mux.MakeEnv(), wg)
+	wg.Add(1)
+	go widget.Label("ACT (*C)", bounds[5+3*POINTS], mux.MakeEnv(), wg)
 	for i := 0; i < POINTS; i++ {
+		wg.Add(1)
 		go widget.Input( // speed
 			rpmChan[i],
 			bounds[3+i],
 			focus.widgets[1][i],
 			mux.MakeEnv(),
+			wg,
 		)
+		wg.Add(1)
 		go widget.Input( // VE
 			veChan[i],
 			bounds[4+POINTS+i],
 			focus.widgets[2][i],
 			mux.MakeEnv(),
+			wg,
 		)
+		wg.Add(1)
 		go widget.Input( // IMAP
 			imapChan[i],
 			bounds[5+2*POINTS+i],
 			focus.widgets[3][i],
 			mux.MakeEnv(),
+			wg,
 		)
+		wg.Add(1)
 		go widget.Input( // ACT
 			actChan[i],
 			bounds[6+3*POINTS+i],
 			focus.widgets[4][i],
 			mux.MakeEnv(),
+			wg,
 		)
 	}
 
@@ -126,11 +156,7 @@ Loop:
 			}
 		}
 	}
-	close(env.Draw())
-	close(displacementChan)
-	for i := range rpmChan {
-		close(rpmChan[i])
-	}
+	fmt.Println("Shutting down...")
 }
 
 func split(elements int, space int) []int {