From d808534da64f71aedd0cae90f0c66222106fe723 Mon Sep 17 00:00:00 2001
From: TorchedSammy <38820196+TorchedSammy@users.noreply.github.com>
Date: Sat, 21 May 2022 21:58:58 -0400
Subject: [PATCH] perf: use waitgroup instead of loop to wait on timers

---
 main.go         | 6 +-----
 timer.go        | 2 ++
 timerhandler.go | 6 ++++++
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index 32484ac..9421bfc 100644
--- a/main.go
+++ b/main.go
@@ -301,11 +301,7 @@ func exit(code int) {
 	// wait for all timers to finish before exiting.
 	// only do that when not interactive
 	if !interactive {
-		for {
-			if timers.running == 0 {
-				os.Exit(code)
-			}
-		}
+		timers.wait()
 	}
 
 	os.Exit(code)
diff --git a/timer.go b/timer.go
index af1a31e..a502087 100644
--- a/timer.go
+++ b/timer.go
@@ -35,6 +35,7 @@ func (t *timer) start() error {
 
 	t.running = true
 	t.th.running++
+	t.th.wg.Add(1)
 	t.ticker = time.NewTicker(t.dur)
 
 	go func() {
@@ -68,6 +69,7 @@ func (t *timer) stop() error {
 	t.channel <- struct{}{}
 	t.running = false
 	t.th.running--
+	t.th.wg.Done()
 	
 	return nil
 }
diff --git a/timerhandler.go b/timerhandler.go
index 1a18edb..2f2e878 100644
--- a/timerhandler.go
+++ b/timerhandler.go
@@ -12,6 +12,7 @@ import (
 var timers *timerHandler
 type timerHandler struct {
 	mu *sync.RWMutex
+	wg *sync.WaitGroup
 	timers map[int]*timer
 	latestID int
 	running int
@@ -22,9 +23,14 @@ func newTimerHandler() *timerHandler {
 		timers: make(map[int]*timer),
 		latestID: 0,
 		mu: &sync.RWMutex{},
+		wg: &sync.WaitGroup{},
 	}
 }
 
+func (th *timerHandler) wait() {
+	th.wg.Wait()
+}
+
 func (th *timerHandler) create(typ timerType, dur time.Duration, fun *rt.Closure) *timer {
 	th.mu.Lock()
 	defer th.mu.Unlock()