support multiple cleanup functions (#2141)

* support multiple cleanup functions

* comment on the expected order

* add boilerplate

* poll while we wait for the test subprocess to setup

* use a tmp file to sychronize the two test processes
this should avoid the data race of reading the subprocess's stdout

* don't make this a breaking change

* fix lint issues

test/cleanup.go

@@ -14,27 +14,52 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// cleanup allows you to define a cleanup function that will be executed
-// if your test is interrupted.
-
 package test
 
 import (
 	"os"
 	"os/signal"
-
-	"knative.dev/pkg/test/logging"
+	"sync"
 )
 
-// CleanupOnInterrupt will execute the function cleanup if an interrupt signal is caught
-func CleanupOnInterrupt(cleanup func(), logf logging.FormatLogger) {
+type logFunc func(template string, args ...interface{})
+
+var cleanup struct {
+	once  sync.Once
+	mutex sync.RWMutex
+	funcs []func()
+}
+
+func waitForInterrupt() {
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt)
+
 	go func() {
-		for range c {
-			logf("Test interrupted, cleaning up.")
-			cleanup()
-			os.Exit(1)
+		<-c
+
+		cleanup.mutex.RLock()
+		defer cleanup.mutex.RUnlock()
+
+		for i := len(cleanup.funcs) - 1; i >= 0; i-- {
+			cleanup.funcs[i]()
 		}
+
+		os.Exit(1)
 	}()
 }
+
+// CleanupOnInterrupt will execute the function if an interrupt signal is caught
+// Deprecated - use OnInterrupt
+func CleanupOnInterrupt(f func(), log logFunc) {
+	OnInterrupt(f)
+}
+
+// OnInterrupt registers a cleanup function to run if an interrupt signal is caught
+func OnInterrupt(cleanupFunc func()) {
+	cleanup.once.Do(waitForInterrupt)
+
+	cleanup.mutex.Lock()
+	defer cleanup.mutex.Unlock()
+
+	cleanup.funcs = append(cleanup.funcs, cleanupFunc)
+}

test/cleanup_test.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2021 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package test
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"strings"
+	"testing"
+	"time"
+
+	"k8s.io/apimachinery/pkg/util/wait"
+)
+
+func TestCleanupOnInterrupt(t *testing.T) {
+	if os.Getenv("CLEANUP") == "1" {
+		OnInterrupt(func() { fmt.Println("cleanup 1") })
+		OnInterrupt(func() { fmt.Println("cleanup 2") })
+		OnInterrupt(func() { fmt.Println("cleanup 3") })
+
+		// This signals to the parent test that it should proceed
+		os.Remove(os.Getenv("READY_FILE"))
+
+		time.Sleep(5 * time.Second)
+		return
+	}
+
+	// TODO: Move to os.CreateTemp when we adopt 1.16 more widely
+	readyFile, err := ioutil.TempFile("", "")
+	if err != nil {
+		t.Fatalf("failed to setup tests")
+	}
+	readyFile.Close()
+
+	cmd := exec.Command(os.Args[0], "-test.run=TestCleanupOnInterrupt", "-test.v=true")
+	cmd.Env = append(os.Environ(), "CLEANUP=1", "READY_FILE="+readyFile.Name())
+
+	var output bytes.Buffer
+	cmd.Stdout = &output
+	cmd.Stderr = &output
+
+	if err := cmd.Start(); err != nil {
+		t.Fatal("Running test failed", err)
+	}
+
+	p, err := os.FindProcess(cmd.Process.Pid)
+	if err != nil {
+		t.Fatal("Failed to find process", err)
+	}
+
+	// poll until the ready file is gone - indicating the subtest has been set up
+	// with the cleanup functions
+	err = wait.PollImmediate(100*time.Millisecond, 2*time.Second, func() (bool, error) {
+		_, err := os.Stat(readyFile.Name())
+		if os.IsNotExist(err) {
+			return true, nil
+		}
+		return false, err
+	})
+
+	if err != nil {
+		t.Fatal("Test subprocess never became ready", err)
+	}
+
+	if err := p.Signal(os.Interrupt); err != nil {
+		t.Fatal("Failed to interrupt", err)
+	}
+
+	err = cmd.Wait()
+	var exitErr *exec.ExitError
+	if ok := errors.As(err, &exitErr); err != nil && !ok {
+		t.Fatal("Running test had abnormal exit", err)
+	}
+
+	testOutput := output.String()
+
+	idx1 := strings.Index(testOutput, "cleanup 1")
+	idx2 := strings.Index(testOutput, "cleanup 2")
+	idx3 := strings.Index(testOutput, "cleanup 3")
+
+	// Order is first in first out (3, 2, 1)
+	if idx3 > idx2 || idx2 > idx1 || idx1 == -1 {
+		t.Errorf("Cleanup functions were not invoked in the proper order")
+	}
+}