wranger: Clean up leak check and use existing version (#15334)

Signed-off-by: Dirkjan Bussink <[email protected]>

go/vt/wrangler/materializer_env_test.go

@@ -19,19 +19,14 @@ package wrangler
 import (
 	"context"
 	"fmt"
-	"os"
 	"regexp"
-	"runtime"
 	"strconv"
 	"strings"
 	"sync"
 	"testing"
-	"time"
-
-	"go.uber.org/goleak"
 
 	"vitess.io/vitess/go/sqltypes"
-	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/test/utils"
 	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/mysqlctl/tmutils"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -40,7 +35,6 @@ import (
 	"vitess.io/vitess/go/vt/vtenv"
 	"vitess.io/vitess/go/vt/vttablet/tmclient"
 
-	_flag "vitess.io/vitess/go/internal/flag"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -61,66 +55,9 @@ type testMaterializerEnv struct {
 //----------------------------------------------
 // testMaterializerEnv
 
-// EnsureNoLeaks is a helper function to fail tests if there are goroutine leaks.
-// At this moment we still have a lot of goroutine leaks in the unit tests in this package.
-// So we only use this while debugging and fixing the leaks. Once fixed we will use this
-// in TestMain instead of just logging the number of leaked goroutines.
-func EnsureNoLeaks(t testing.TB) {
-	if t.Failed() {
-		return
-	}
-	err := ensureNoGoroutines()
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
-func ensureNoGoroutines() error {
-	// These goroutines have been found to stay around.
-	// Need to investigate and fix the Vitess ones at some point, if we indeed find out that they are unintended leaks.
-	var leaksToIgnore = []goleak.Option{
-		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
-		goleak.IgnoreTopFunction("github.com/golang/glog.(*loggingT).flushDaemon"),
-		goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/dbconfigs.init.0.func1"),
-		goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.resetAggregators"),
-		goleak.IgnoreTopFunction("vitess.io/vitess/go/vt/vtgate.processQueryInfo"),
-		goleak.IgnoreTopFunction("github.com/patrickmn/go-cache.(*janitor).Run"),
-	}
-
-	const (
-		// give ample time for the goroutines to exit in CI.
-		waitTime      = 100 * time.Millisecond
-		numIterations = 50 // 5 seconds
-	)
-	var err error
-	for i := 0; i < numIterations; i++ {
-		err = goleak.Find(leaksToIgnore...)
-		if err == nil {
-			return nil
-		}
-		time.Sleep(waitTime)
-	}
-	return err
-}
-
-func testMainWrapper(m *testing.M) int {
-	startingNumGoRoutines := runtime.NumGoroutine()
-	defer func() {
-		numGoroutines := runtime.NumGoroutine()
-		if numGoroutines > startingNumGoRoutines {
-			log.Infof("!!!!!!!!!!!! Wrangler unit tests Leaked %d goroutines", numGoroutines-startingNumGoRoutines)
-		}
-	}()
-	_flag.ParseFlagsForTest()
-	return m.Run()
-}
-
-func TestMain(m *testing.M) {
-	os.Exit(testMainWrapper(m))
-}
-
-func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.MaterializeSettings, sources, targets []string) *testMaterializerEnv {
+func newTestMaterializerEnv(t *testing.T, ms *vtctldatapb.MaterializeSettings, sources, targets []string) (*testMaterializerEnv, context.Context) {
 	t.Helper()
+	ctx, cancel := context.WithCancel(context.Background())
 	env := &testMaterializerEnv{
 		ms:       ms,
 		sources:  sources,
@@ -167,7 +104,12 @@ func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.M
 	if ms.Workflow != "" {
 		env.expectValidation()
 	}
-	return env
+	t.Cleanup(func() {
+		defer utils.EnsureNoLeaks(t)
+		env.close()
+		cancel()
+	})
+	return env, ctx
 }
 
 func (env *testMaterializerEnv) expectValidation() {