From 7380c5a9fa8403fe6eb6354ceacf0cccdc0c61aa Mon Sep 17 00:00:00 2001
From: Prashant Varanasi <prashant@uber.com>
Date: Mon, 11 Nov 2019 13:21:39 -0800
Subject: [PATCH] Fix flaky tests (#36)

Some tests assume that stack.All() will give a stable set of goroutines:
that any background goroutines from the test framework/previous tests
have run till they're blocked.

Add verification that stack.All() returns a "stable" set of stacks
before running these tests.
---
 internal/stack/stacks_test.go | 29 ++++++++++++++++++++++
 options_test.go               |  2 +-
 utils_test.go                 | 45 +++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/internal/stack/stacks_test.go b/internal/stack/stacks_test.go
index 074d1f0..646dd2a 100644
--- a/internal/stack/stacks_test.go
+++ b/internal/stack/stacks_test.go
@@ -21,7 +21,9 @@
 package stack
 
 import (
+	"runtime"
 	"sort"
+	"strings"
 	"sync"
 	"testing"
 
@@ -46,8 +48,18 @@ func TestAll(t *testing.T) {
 		go waitForDone()
 	}
 
+	cur := Current()
 	got := All()
 
+	// Retry until the background stacks are not runnable/running.
+	for {
+		if !isBackgroundRunning(cur, got) {
+			break
+		}
+		runtime.Gosched()
+		got = All()
+	}
+
 	// We have exactly 7 gorotuines:
 	// "main" goroutine
 	// test goroutine
@@ -104,6 +116,7 @@ func TestAllLargeStack(t *testing.T) {
 			if count == 0 {
 				started.Done()
 				<-done
+				return
 			}
 			f(count - 1)
 		}
@@ -126,3 +139,19 @@ type byGoroutineID []Stack
 func (ss byGoroutineID) Len() int           { return len(ss) }
 func (ss byGoroutineID) Less(i, j int) bool { return ss[i].ID() < ss[j].ID() }
 func (ss byGoroutineID) Swap(i, j int)      { ss[i], ss[j] = ss[j], ss[i] }
+
+// Note: This is the same logic as in ../../utils_test.go
+// Copy+pasted to avoid dependency loops and exporting this test-helper.
+func isBackgroundRunning(cur Stack, stacks []Stack) bool {
+	for _, s := range stacks {
+		if cur.ID() == s.ID() {
+			continue
+		}
+
+		if strings.Contains(s.State(), "run") {
+			return true
+		}
+	}
+
+	return false
+}
diff --git a/options_test.go b/options_test.go
index dd41ee2..6bdec34 100644
--- a/options_test.go
+++ b/options_test.go
@@ -32,7 +32,7 @@ import (
 func TestOptionsFilters(t *testing.T) {
 	opts := buildOpts()
 	cur := stack.Current()
-	all := stack.All()
+	all := getStableAll(t, cur)
 
 	// At least one of these should be the same as current, the others should be filtered out.
 	for _, s := range all {
diff --git a/utils_test.go b/utils_test.go
index 4f8d331..7d6b0f1 100644
--- a/utils_test.go
+++ b/utils_test.go
@@ -20,6 +20,14 @@
 
 package goleak
 
+import (
+	"runtime"
+	"strings"
+	"testing"
+
+	"go.uber.org/goleak/internal/stack"
+)
+
 type blockedG struct {
 	started chan struct{}
 	wait    chan struct{}
@@ -43,3 +51,40 @@ func (bg *blockedG) run() {
 func (bg *blockedG) unblock() {
 	close(bg.wait)
 }
+
+func getStableAll(t *testing.T, cur stack.Stack) []stack.Stack {
+	all := stack.All()
+
+	// There may be running goroutines that were just scheduled or finishing up
+	// from previous tests, so reduce flakiness by waiting till no other goroutines
+	// are runnable or running except the current goroutine.
+	for retry := 0; true; retry++ {
+		if !isBackgroundRunning(cur, all) {
+			break
+		}
+
+		if retry >= 100 {
+			t.Fatalf("background goroutines are possibly running, %v", all)
+		}
+
+		runtime.Gosched()
+		all = stack.All()
+	}
+
+	return all
+}
+
+// Note: This is the same logic as in internal/stacks/stacks_test.go
+func isBackgroundRunning(cur stack.Stack, stacks []stack.Stack) bool {
+	for _, s := range stacks {
+		if cur.ID() == s.ID() {
+			continue
+		}
+
+		if strings.Contains(s.State(), "run") {
+			return true
+		}
+	}
+
+	return false
+}