cloud: test server detects leaked storage instances

pkg/cloud/cloudtestutils/BUILD.bazel

@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 
 go_library(
     name = "cloudtestutils",
-    srcs = ["cloud_test_helpers.go"],
+    srcs = [
+        "cloud_test_helpers.go",
+        "leak_detector_storage_factory.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/cloud/cloudtestutils",
     visibility = ["//visibility:public"],
     deps = [
@@ -14,6 +17,7 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/sql/isql",
         "//pkg/testutils",
+        "//pkg/util/closetest",
         "//pkg/util/ioctx",
         "//pkg/util/randutil",
         "//pkg/util/sysutil",

pkg/cloud/cloudtestutils/leak_detector_storage_factory.go

@@ -0,0 +1,39 @@
+package cloudtestutils
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/cloud"
+	"github.com/cockroachdb/cockroach/pkg/util/closetest"
+)
+
+type StorageLeakDetector struct {
+	tracker *closetest.AllocationTracker
+}
+
+func NewStorageLeakDetector() *StorageLeakDetector {
+	return &StorageLeakDetector { closetest.NewTracker("cloud.ExternalStorage") }
+}
+
+func (s *StorageLeakDetector) Check() error {
+	return s.tracker.CheckLeaks()
+}
+
+func (s *StorageLeakDetector) Wrap(storage cloud.ExternalStorage) cloud.ExternalStorage {
+	if storage == nil {
+		return nil
+	}
+	return &wrappedStorage{
+		storage,
+		s.tracker.TrackAllocation(3),
+	}
+}
+
+type wrappedStorage struct {
+	cloud.ExternalStorage
+	allocation *closetest.Allocation
+}
+
+// Close implements cloud.ExternalStorage.
+func (w *wrappedStorage) Close() error {
+	w.allocation.Close()
+	return w.ExternalStorage.Close()
+}

pkg/server/BUILD.bazel

@@ -92,6 +92,7 @@ go_library(
         "//pkg/build",
         "//pkg/cloud",
         "//pkg/cloud/cloudpb",
+        "//pkg/cloud/cloudtestutils",
         "//pkg/cloud/externalconn",
         "//pkg/clusterversion",
         "//pkg/config",

pkg/server/server.go

@@ -20,6 +20,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/blobs"
 	"github.com/cockroachdb/cockroach/pkg/blobs/blobspb"
+	"github.com/cockroachdb/cockroach/pkg/cloud"
+	"github.com/cockroachdb/cockroach/pkg/cloud/cloudpb"
+	"github.com/cockroachdb/cockroach/pkg/cloud/cloudtestutils"
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/gossip"
 	"github.com/cockroachdb/cockroach/pkg/inspectz"
@@ -671,9 +674,22 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 
 	// Create an ExternalStorageBuilder. This is only usable after Start() where
 	// we initialize all the configuration params.
+	leakChecker := cloudtestutils.NewStorageLeakDetector()
 	externalStorageBuilder := &externalStorageBuilder{}
-	externalStorage := externalStorageBuilder.makeExternalStorage
+	externalStorage := func(ctx context.Context, dest cloudpb.ExternalStorage, opts ...cloud.ExternalStorageOption) (cloud.ExternalStorage, error) {
+		storage, err := externalStorageBuilder.makeExternalStorage(ctx, dest, opts...)
+		return leakChecker.Wrap(storage), err
+	}
 	externalStorageFromURI := externalStorageBuilder.makeExternalStorageFromURI
+	externalStorageFromURI = func(ctx context.Context, uri string, user username.SQLUsername, opts ...cloud.ExternalStorageOption) (cloud.ExternalStorage, error) {
+		storage, err := externalStorageBuilder.makeExternalStorageFromURI(ctx, uri, user, opts...)
+		return leakChecker.Wrap(storage), err
+	}
+	stopper.AddCloser(stop.CloserFn(func() {
+		if err := leakChecker.Check(); err != nil {
+			panic(err)
+		}
+	}))
 
 	protectedtsKnobs, _ := cfg.TestingKnobs.ProtectedTS.(*protectedts.TestingKnobs)
 	protectedtsProvider, err := ptprovider.New(ptprovider.Config{

pkg/testutils/serverutils/test_server_shim.go

@@ -242,6 +242,7 @@ type TestFataler interface {
 	Fatalf(format string, args ...interface{})
 	Errorf(format string, args ...interface{})
 	FailNow()
+	Cleanup(func())
 }
 
 // StartServerOnlyE is like StartServerOnly() but it lets

pkg/util/closetest/BUILD.bazel

@@ -0,0 +1,9 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "closetest",
+    srcs = ["close_tracker.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/util/closetest",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_cockroachdb_errors//:errors"],
+)

pkg/util/closetest/close_tracker.go

@@ -0,0 +1,95 @@
+package closetest
+
+import (
+	"fmt"
+	"path"
+	"runtime"
+	"strings"
+	"sync"
+
+	"github.com/cockroachdb/errors"
+)
+
+type AllocationTracker struct {
+	name string
+	mu struct {
+		sync.Mutex
+		total int
+		allocations map[int]*Allocation
+	}
+}
+
+func NewTracker(name string) *AllocationTracker {
+	tracker := &AllocationTracker{
+		name: name,
+	}
+	tracker.mu.allocations = map[int]*Allocation{}
+	return tracker
+}
+
+func (a *AllocationTracker) CheckLeaks() error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	if len(a.mu.allocations) == 0 {
+		return nil
+	}
+
+	leakingStacks := map[stack]int {}
+	for _, leakedAllocation := range a.mu.allocations {
+		leakingStacks[leakedAllocation.stack] = leakingStacks[leakedAllocation.stack] + 1
+	}
+
+	var builder strings.Builder
+	for leakedStack, count := range leakingStacks {
+		fmt.Fprintf(&builder, "found %d copies of stack: \n", count)
+
+		frames := runtime.CallersFrames(leakedStack.callers[:leakedStack.size])
+		for index := 0; ; index++ {
+			frame, more := frames.Next()
+			fmt.Fprintf(&builder, "%d: %s %s:%d\n", index, path.Base(frame.Function), path.Base(frame.File), frame.Line)
+			if !more {
+				break
+			}
+		}
+	}
+
+	return errors.Newf(
+		"leaked %d of %d %s instances from %d unique allocation stacks:\n%s",
+		len(a.mu.allocations), a.mu.total, a.name, len(leakingStacks), builder.String())
+}
+
+func (a *AllocationTracker) TrackAllocation(ignoreCallers int) *Allocation {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	id := a.mu.total
+	a.mu.total += 1
+
+	allocation := &Allocation {
+		tracker: a,
+		id: id,
+	}
+	allocation.stack.size = runtime.Callers(ignoreCallers + 1, allocation.stack.callers[:])
+
+	a.mu.allocations[id] = allocation
+
+	return allocation
+}
+
+type stack struct {
+	callers [32]uintptr
+	size    int
+}
+
+type Allocation struct {
+	tracker *AllocationTracker
+	stack
+	id      int
+}
+
+func (a *Allocation) Close() {
+	a.tracker.mu.Lock()
+	defer a.tracker.mu.Unlock()
+	delete(a.tracker.mu.allocations, a.id)
+}