From 54e18f1c2a44f9f2664486e8053c4ee40d41fb8a Mon Sep 17 00:00:00 2001
From: Ian Lance Taylor <iant@golang.org>
Date: Fri, 28 Aug 2020 10:34:16 -0700
Subject: [PATCH] testing: run a Cleanup registered by a Cleanup

Fixes #41085

Change-Id: Ieafc60cbc8e09f1935d38b1767b084d78dae5cb4
Reviewed-on: https://go-review.googlesource.com/c/go/+/251457
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>
---
 src/testing/sub_test.go | 27 +++++++++++++++
 src/testing/testing.go  | 76 +++++++++++++++++++++++++----------------
 2 files changed, 74 insertions(+), 29 deletions(-)

diff --git a/src/testing/sub_test.go b/src/testing/sub_test.go
index 8eb0084b1c..51fc0ccc39 100644
--- a/src/testing/sub_test.go
+++ b/src/testing/sub_test.go
@@ -928,3 +928,30 @@ func TestCleanupParallelSubtests(t *T) {
 		t.Errorf("unexpected cleanup count; got %d want 1", ranCleanup)
 	}
 }
+
+func TestNestedCleanup(t *T) {
+	ranCleanup := 0
+	t.Run("test", func(t *T) {
+		t.Cleanup(func() {
+			if ranCleanup != 2 {
+				t.Errorf("unexpected cleanup count in first cleanup: got %d want 2", ranCleanup)
+			}
+			ranCleanup++
+		})
+		t.Cleanup(func() {
+			if ranCleanup != 0 {
+				t.Errorf("unexpected cleanup count in second cleanup: got %d want 0", ranCleanup)
+			}
+			ranCleanup++
+			t.Cleanup(func() {
+				if ranCleanup != 1 {
+					t.Errorf("unexpected cleanup count in nested cleanup: got %d want 1", ranCleanup)
+				}
+				ranCleanup++
+			})
+		})
+	})
+	if ranCleanup != 3 {
+		t.Errorf("unexpected cleanup count: got %d want 3", ranCleanup)
+	}
+}
diff --git a/src/testing/testing.go b/src/testing/testing.go
index 01743969ee..f4f0060523 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -403,7 +403,7 @@ type common struct {
 	skipped     bool                // Test of benchmark has been skipped.
 	done        bool                // Test is finished and all subtests have completed.
 	helpers     map[string]struct{} // functions to be skipped when writing file/line info
-	cleanup     func()              // optional function to be called at the end of the test
+	cleanups    []func()            // optional functions to be called at the end of the test
 	cleanupName string              // Name of the cleanup function.
 	cleanupPc   []uintptr           // The stack trace at the point where Cleanup was called.
 
@@ -855,28 +855,31 @@ func (c *common) Helper() {
 // subtests complete. Cleanup functions will be called in last added,
 // first called order.
 func (c *common) Cleanup(f func()) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	oldCleanup := c.cleanup
-	oldCleanupPc := c.cleanupPc
-	c.cleanup = func() {
-		if oldCleanup != nil {
-			defer func() {
-				c.mu.Lock()
-				c.cleanupPc = oldCleanupPc
-				c.mu.Unlock()
-				oldCleanup()
-			}()
-		}
+	var pc [maxStackLen]uintptr
+	// Skip two extra frames to account for this function and runtime.Callers itself.
+	n := runtime.Callers(2, pc[:])
+	cleanupPc := pc[:n]
+
+	fn := func() {
+		defer func() {
+			c.mu.Lock()
+			defer c.mu.Unlock()
+			c.cleanupName = ""
+			c.cleanupPc = nil
+		}()
+
+		name := callerName(0)
 		c.mu.Lock()
-		c.cleanupName = callerName(0)
+		c.cleanupName = name
+		c.cleanupPc = cleanupPc
 		c.mu.Unlock()
+
 		f()
 	}
-	var pc [maxStackLen]uintptr
-	// Skip two extra frames to account for this function and runtime.Callers itself.
-	n := runtime.Callers(2, pc[:])
-	c.cleanupPc = pc[:n]
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.cleanups = append(c.cleanups, fn)
 }
 
 var tempDirReplacer struct {
@@ -934,22 +937,37 @@ const (
 // If catchPanic is true, this will catch panics, and return the recovered
 // value if any.
 func (c *common) runCleanup(ph panicHandling) (panicVal interface{}) {
-	c.mu.Lock()
-	cleanup := c.cleanup
-	c.cleanup = nil
-	c.mu.Unlock()
-	if cleanup == nil {
-		return nil
-	}
-
 	if ph == recoverAndReturnPanic {
 		defer func() {
 			panicVal = recover()
 		}()
 	}
 
-	cleanup()
-	return nil
+	// Make sure that if a cleanup function panics,
+	// we still run the remaining cleanup functions.
+	defer func() {
+		c.mu.Lock()
+		recur := len(c.cleanups) > 0
+		c.mu.Unlock()
+		if recur {
+			c.runCleanup(normalPanic)
+		}
+	}()
+
+	for {
+		var cleanup func()
+		c.mu.Lock()
+		if len(c.cleanups) > 0 {
+			last := len(c.cleanups) - 1
+			cleanup = c.cleanups[last]
+			c.cleanups = c.cleanups[:last]
+		}
+		c.mu.Unlock()
+		if cleanup == nil {
+			return nil
+		}
+		cleanup()
+	}
 }
 
 // callerName gives the function name (qualified with a package path)
-- 
2.48.1