internal/fdtrace: allow tracing of sys package

internal/sys/fd.go

@@ -7,6 +7,7 @@ import (
 	"runtime"
 	"strconv"
 
+	"github.com/cilium/ebpf/internal/testutils/fdtrace"
 	"github.com/cilium/ebpf/internal/unix"
 )
 
@@ -17,15 +18,7 @@ type FD struct {
 }
 
 func newFD(value int) *FD {
-	if onLeakFD != nil {
-		// Attempt to store the caller's stack for the given fd value.
-		// Panic if fds contains an existing stack for the fd.
-		old, exist := fds.LoadOrStore(value, callersFrames())
-		if exist {
-			f := old.(*runtime.Frames)
-			panic(fmt.Sprintf("found existing stack for fd %d:\n%s", value, FormatFrames(f)))
-		}
-	}
+	fdtrace.TraceFD(value, 1)
 
 	fd := &FD{value}
 	runtime.SetFinalizer(fd, (*FD).finalize)
@@ -39,13 +32,7 @@ func (fd *FD) finalize() {
 		return
 	}
 
-	// Invoke the fd leak callback. Calls LoadAndDelete to guarantee the callback
-	// is invoked at most once for one sys.FD allocation, runtime.Frames can only
-	// be unwound once.
-	f, ok := fds.LoadAndDelete(fd.Int())
-	if ok && onLeakFD != nil {
-		onLeakFD(f.(*runtime.Frames))
-	}
+	fdtrace.LeakFD(fd.raw)
 
 	_ = fd.Close()
 }
@@ -96,8 +83,8 @@ func (fd *FD) Close() error {
 }
 
 func (fd *FD) disown() int {
-	value := int(fd.raw)
-	fds.Delete(int(value))
+	value := fd.raw
+	fdtrace.ForgetFD(value)
 	fd.raw = -1
 
 	runtime.SetFinalizer(fd, nil)

internal/sys/fd_test.go

@@ -41,6 +41,7 @@ func TestFD(t *testing.T) {
 	fd, err := NewFD(0)
 	qt.Assert(t, qt.IsNil(err))
 	qt.Assert(t, qt.Not(qt.Equals(fd.Int(), 0)), qt.Commentf("fd value should not be zero"))
+	qt.Assert(t, qt.IsNil(fd.Close()))
 
 	var stat unix.Stat_t
 	err = unix.Fstat(0, &stat)

internal/sys/syscall_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/cilium/ebpf/internal/testutils/fdtrace"
 	"github.com/cilium/ebpf/internal/unix"
 
 	"github.com/go-quicktest/qt"
@@ -59,3 +60,7 @@ func TestSyscallError(t *testing.T) {
 		t.Error("Error is the SyscallError")
 	}
 }
+
+func TestMain(m *testing.M) {
+	fdtrace.TestMain(m)
+}

internal/testutils/fdtrace/fd_trace.go

@@ -0,0 +1,103 @@
+package fdtrace
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"runtime"
+	"sync"
+	"sync/atomic"
+)
+
+// foundLeak is atomic since the GC may collect objects in parallel.
+var foundLeak atomic.Bool
+
+func onLeakFD(fs *runtime.Frames) {
+	foundLeak.Store(true)
+	fmt.Fprintln(os.Stderr, "leaked fd created at:")
+	fmt.Fprintln(os.Stderr, formatFrames(fs))
+}
+
+// fds is a registry of all file descriptors wrapped into sys.fds that were
+// created while an fd tracer was active.
+var fds *sync.Map // map[int]*runtime.Frames
+
+// TraceFD associates raw with the current execution stack.
+//
+// skip controls how many entries of the stack the function should skip.
+func TraceFD(raw int, skip int) {
+	if fds == nil {
+		return
+	}
+
+	// Attempt to store the caller's stack for the given fd value.
+	// Panic if fds contains an existing stack for the fd.
+	old, exist := fds.LoadOrStore(raw, callersFrames(skip))
+	if exist {
+		f := old.(*runtime.Frames)
+		panic(fmt.Sprintf("found existing stack for fd %d:\n%s", raw, formatFrames(f)))
+	}
+}
+
+// ForgetFD removes any existing association for raw.
+func ForgetFD(raw int) {
+	if fds != nil {
+		fds.Delete(raw)
+	}
+}
+
+// LeakFD indicates that raw was leaked.
+//
+// Calling the function with a value that was not passed to [TraceFD] before
+// is undefined.
+func LeakFD(raw int) {
+	if fds == nil {
+		return
+	}
+
+	// Invoke the fd leak callback. Calls LoadAndDelete to guarantee the callback
+	// is invoked at most once for one sys.FD allocation, runtime.Frames can only
+	// be unwound once.
+	f, ok := fds.LoadAndDelete(raw)
+	if ok {
+		onLeakFD(f.(*runtime.Frames))
+	}
+}
+
+// flushFrames removes all elements from fds and returns them as a slice. This
+// deals with the fact that a runtime.Frames can only be unwound once using
+// Next().
+func flushFrames() []*runtime.Frames {
+	var frames []*runtime.Frames
+	fds.Range(func(key, value any) bool {
+		frames = append(frames, value.(*runtime.Frames))
+		fds.Delete(key)
+		return true
+	})
+	return frames
+}
+
+func callersFrames(skip int) *runtime.Frames {
+	c := make([]uintptr, 32)
+
+	// Skip runtime.Callers and this function.
+	i := runtime.Callers(skip+2, c)
+	if i == 0 {
+		return nil
+	}
+
+	return runtime.CallersFrames(c)
+}
+
+// formatFrames formats a runtime.Frames as a human-readable string.
+func formatFrames(fs *runtime.Frames) string {
+	var b bytes.Buffer
+	for {
+		f, more := fs.Next()
+		b.WriteString(fmt.Sprintf("\t%s+%#x\n\t\t%s:%d\n", f.Function, f.PC-f.Entry, f.File, f.Line))
+		if !more {
+			break
+		}
+	}
+	return b.String()
+}

internal/testutils/fdtrace/main.go

@@ -0,0 +1,31 @@
+package fdtrace
+
+import (
+	"os"
+	"sync"
+)
+
+type testingM interface {
+	Run() int
+}
+
+// TestMain runs m with fd tracing enabled.
+//
+// The function calls [os.Exit] and does not return.
+func TestMain(m testingM) {
+	fds = new(sync.Map)
+
+	ret := m.Run()
+
+	if fs := flushFrames(); len(fs) != 0 {
+		for _, f := range fs {
+			onLeakFD(f)
+		}
+	}
+
+	if foundLeak.Load() {
+		ret = 99
+	}
+
+	os.Exit(ret)
+}