Fix race condition (#2)

* Fix race condition with CancelAll

* Update CHANGELOG

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.2] - 2020-01-29
+
+### Fixed
+
+- Fix race condition in CancelAll
+
 ## [0.2.1] - 2020-01-28
 
 ### Changed
@@ -34,3 +40,4 @@ All notable changes to this project will be documented in this file. The format
 ### Added
 
 - initial commit of rununtil library
+

rununtil.go

@@ -95,22 +95,17 @@ type canceller struct {
 
 func (canc *canceller) addChannel(key string, c chan struct{}) {
 	canc.mux.Lock()
+	defer canc.mux.Unlock()
 	canc.signals[key] = c
-	canc.mux.Unlock()
-}
-
-func (canc *canceller) removeChannel(key string) {
-	canc.mux.Lock()
-	delete(canc.signals, key)
-	canc.mux.Unlock()
 }
 
 func (canc *canceller) cancelAll() {
 	canc.mux.Lock()
+	defer canc.mux.Unlock()
 	for key := range canc.signals {
 		close(canc.signals[key])
+		delete(canc.signals, key)
 	}
-	canc.mux.Unlock()
 }
 
 var globalCanceller canceller
@@ -146,7 +141,6 @@ func AwaitKillSignals(signals []os.Signal, runnerFuncs ...RunnerFunc) {
 	finish := make(chan struct{})
 	uuid := uuid.New()
 	globalCanceller.addChannel(uuid.String(), finish)
-	defer globalCanceller.removeChannel(uuid.String())
 
 	for _, runner := range runnerFuncs {
 		shutdown := runner()

rununtil_test.go

@@ -147,9 +147,19 @@ func TestRununtilCancelAll_MultipleTimes(t *testing.T) {
 }
 
 func TestRununtilCancelAll_Threadsafe(t *testing.T) {
+	var hasBeenKilledVec [100]bool
 	for idx := 0; idx < 100; idx++ {
+		cancel := rununtil.Killed(helperMakeMain(&hasBeenKilledVec[idx]))
+		cancel()
 		rununtil.CancelAll()
 	}
+	// yield control back to scheduler so that killing can actually happen
+	time.Sleep(time.Millisecond)
+	for idx, hasBeenKilled := range hasBeenKilledVec {
+		if !hasBeenKilled {
+			t.Fatalf("expected main to have been killed: %d", idx)
+		}
+	}
 }
 
 // Annoyingly this test has to be run by itself to actually fail...