Merge pull request #3 from klauspost/fix-out-of-bounds-cancel

Fix index out of bounds when cancelling.

http_test.go

@@ -245,6 +245,7 @@ func TestWrapHandlerFuncOrder(t *testing.T) {
 	var wait = make(chan bool)
 	var waiting = make(chan bool)
 	fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer close(finished)
 		close(waiting)
 		<-wait
 	})
@@ -255,7 +256,6 @@ func TestWrapHandlerFuncOrder(t *testing.T) {
 		res := httptest.NewRecorder()
 		req, _ := http.NewRequest("", "", bytes.NewBufferString(""))
 		wrapped(res, req)
-		close(finished)
 	}()
 
 	release := time.After(time.Millisecond * 100)

shutdown.go

@@ -147,6 +147,7 @@ func (s *Notifier) Cancel() {
 					b = shutdownQueue[n][i].n
 					if fn.internal.n == b {
 						shutdownQueue[n] = append(shutdownQueue[n][:i], shutdownQueue[n][i+1:]...)
+						break
 					}
 				}
 				// Cancel, so the goroutine exits.
@@ -179,12 +180,10 @@ func (s *Notifier) CancelWait() {
 			if a == b {
 				// Find the matching internal and remove that.
 				for i := range shutdownQueue[n] {
-					if len(shutdownQueue[n]) <= i {
-						continue
-					}
 					b = shutdownQueue[n][i].n
 					if fn.internal.n == b {
 						shutdownQueue[n] = append(shutdownQueue[n][:i], shutdownQueue[n][i+1:]...)
+						break
 					}
 				}
 				// Cancel, so the goroutine exits.

shutdown_test.go

@@ -5,6 +5,7 @@ package shutdown
 import (
 	"bytes"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"runtime"
 	"strings"
@@ -914,6 +915,76 @@ func TestFnSingleCancel(t *testing.T) {
 	}
 }
 
+func TestCancelMulti(t *testing.T) {
+	reset()
+	SetTimeout(time.Second)
+	defer close(startTimer(t))
+	rand.Seed(0xC0CAC01A)
+	for i := 0; i < 1000; i++ {
+		var n Notifier
+		switch rand.Int31n(10) {
+		case 0:
+			n = PreShutdown()
+		case 1:
+			n = First()
+		case 2:
+			n = Second()
+		case 3:
+			n = Third()
+		case 4:
+			n = PreShutdownFn(func() {})
+		case 5:
+			n = FirstFn(func() {})
+		case 6:
+			n = SecondFn(func() {})
+		case 7:
+			n = ThirdFn(func() {})
+		}
+		go func(n Notifier, t time.Duration) {
+			time.Sleep(t)
+			n.Cancel()
+		}(n, time.Millisecond*time.Duration(rand.Intn(100)))
+		time.Sleep(time.Millisecond)
+	}
+	// Start shutdown
+	Shutdown()
+}
+
+func TestCancelWaitMulti(t *testing.T) {
+	reset()
+	SetTimeout(time.Second)
+	defer close(startTimer(t))
+	rand.Seed(0xC0CAC01A)
+	for i := 0; i < 1000; i++ {
+		var n Notifier
+		switch rand.Int31n(10) {
+		case 0:
+			n = PreShutdown()
+		case 1:
+			n = First()
+		case 2:
+			n = Second()
+		case 3:
+			n = Third()
+		case 4:
+			n = PreShutdownFn(func() {})
+		case 5:
+			n = FirstFn(func() {})
+		case 6:
+			n = SecondFn(func() {})
+		case 7:
+			n = ThirdFn(func() {})
+		}
+		go func(n Notifier, t time.Duration) {
+			time.Sleep(t)
+			n.CancelWait()
+		}(n, time.Millisecond*time.Duration(rand.Intn(250)))
+		time.Sleep(time.Millisecond)
+	}
+	// Start shutdown
+	Shutdown()
+}
+
 // Get a notifier and perform our own code when we shutdown
 func ExampleNotifier() {
 	shutdown := First()