all: use runtime.AddCleanup

Convert all uses of runtime.SetFinalizer to runtime.AddCleanup.

The exception is the unsafe memory code, which cannot currently be converted
due to fundamental differences with finalizers. Cleanups cannot carry
references to the object in question, and cannot resurrect the object,
meaning the original object's allocation can be used to service new
allocations before the cleanup has run.

In our case, this means bpf map memory gets zeroed by the allocator, which is
obviously a show stopper. If we want to move off finalizers, we need to
implement golang/go#70224 first and switch to that
instead.

Signed-off-by: Timo Beckers <[email protected]>
Suggested-by: Lorenz Bauer <[email protected]>

Author: ti-mo

btf/kernel.go

@@ -151,9 +151,9 @@ func loadKernelSpec() (*Spec, error) {
 			return nil, fmt.Errorf("load vmlinux: %w", err)
 		}
 
-		runtime.SetFinalizer(spec.decoder.sharedBuf, func(_ *sharedBuf) {
-			_ = unix.Munmap(raw)
-		})
+		runtime.AddCleanup(spec.decoder.sharedBuf, func(b []byte) {
+			_ = unix.Munmap(b)
+		}, raw)
 
 		return spec, nil
 	}

internal/epoll/poller.go

@@ -34,6 +34,7 @@ type Poller struct {
 	eventMu    sync.Mutex
 	closeEvent *eventFd
 	flushEvent *eventFd
+	cleanup    runtime.Cleanup
 }
 
 func New() (_ *Poller, err error) {
@@ -75,7 +76,10 @@ func New() (_ *Poller, err error) {
 		return nil, fmt.Errorf("add flush eventfd: %w", err)
 	}
 
-	runtime.SetFinalizer(p, (*Poller).Close)
+	p.cleanup = runtime.AddCleanup(p, func(raw int) {
+		_ = unix.Close(raw)
+	}, p.epollFd)
+
 	return p, nil
 }
 
@@ -84,7 +88,7 @@ func New() (_ *Poller, err error) {
 // Interrupts any calls to Wait. Multiple calls to Close are valid, but subsequent
 // calls will return os.ErrClosed.
 func (p *Poller) Close() error {
-	runtime.SetFinalizer(p, nil)
+	p.cleanup.Stop()
 
 	// Interrupt Wait() via the closeEvent fd if it's currently blocked.
 	if err := p.wakeWaitForClose(); err != nil {

internal/sys/fd.go

@@ -21,23 +21,14 @@ const invalidFd = -1
 func newFD(value int) *FD {
 	testmain.TraceFD(value, 1)
 
-	fd := &FD{value}
-	runtime.SetFinalizer(fd, (*FD).finalize)
+	fd := &FD{raw: value}
+	fd.cleanup = runtime.AddCleanup(fd, func(raw int) {
+		testmain.LeakFD(raw)
+		_ = unix.Close(raw)
+	}, fd.raw)
 	return fd
 }
 
-// finalize is set as the FD's runtime finalizer and
-// sends a leak trace before calling FD.Close().
-func (fd *FD) finalize() {
-	if fd.raw == invalidFd {
-		return
-	}
-
-	testmain.LeakFD(fd.raw)
-
-	_ = fd.Close()
-}
-
 func (fd *FD) String() string {
 	return strconv.FormatInt(int64(fd.raw), 10)
 }
@@ -63,6 +54,6 @@ func (fd *FD) Disown() int {
 	testmain.ForgetFD(value)
 	fd.raw = invalidFd
 
-	runtime.SetFinalizer(fd, nil)
+	fd.cleanup.Stop()
 	return value
 }

internal/sys/fd_other.go

@@ -5,12 +5,14 @@ package sys
 import (
 	"fmt"
 	"os"
+	"runtime"
 
 	"github.com/cilium/ebpf/internal/unix"
 )
 
 type FD struct {
-	raw int
+	raw     int
+	cleanup runtime.Cleanup
 }
 
 // NewFD wraps a raw fd with a finalizer.

internal/sys/fd_windows.go

@@ -3,6 +3,7 @@ package sys
 import (
 	"fmt"
 	"os"
+	"runtime"
 
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/efw"
@@ -12,7 +13,8 @@ import (
 //
 // It is not equivalent to a real file descriptor or handle.
 type FD struct {
-	raw int
+	raw     int
+	cleanup runtime.Cleanup
 }
 
 // NewFD wraps a raw fd with a finalizer.

internal/testutils/chan.go

@@ -8,6 +8,8 @@ import (
 // WaitChan waits for a value to be sent on a channel, or for a timeout to
 // occur. If the timeout is reached, the test will fail.
 func WaitChan[T any](tb testing.TB, ch <-chan T, timeout time.Duration) {
+	tb.Helper()
+
 	select {
 	case <-ch:
 		return

internal/tracefs/kprobe.go

@@ -200,6 +200,8 @@ type Event struct {
 	group, name string
 	// event id allocated by the kernel. 0 if the event has already been removed.
 	id uint64
+
+	cleanup runtime.Cleanup
 }
 
 // NewEvent creates a new ephemeral trace event.
@@ -316,8 +318,11 @@ func NewEvent(args ProbeArgs) (*Event, error) {
 		return nil, fmt.Errorf("get trace event id: %w", err)
 	}
 
-	evt := &Event{args.Type, args.Group, eventName, tid}
-	runtime.SetFinalizer(evt, (*Event).Close)
+	evt := &Event{typ: args.Type, group: args.Group, name: eventName, id: tid}
+	evt.cleanup = runtime.AddCleanup(evt, func(*byte) {
+		_ = removeEvent(args.Type, fmt.Sprintf("%s/%s", args.Group, eventName))
+	}, nil)
+
 	return evt, nil
 }
 
@@ -330,7 +335,7 @@ func (evt *Event) Close() error {
 	}
 
 	evt.id = 0
-	runtime.SetFinalizer(evt, nil)
+	evt.cleanup.Stop()
 	pe := fmt.Sprintf("%s/%s", evt.group, evt.name)
 	return removeEvent(evt.typ, pe)
 }

memory.go

@@ -38,6 +38,8 @@ type Memory struct {
 	b    []byte
 	ro   bool
 	heap bool
+
+	cleanup runtime.Cleanup
 }
 
 func newMemory(fd, size int) (*Memory, error) {
@@ -62,20 +64,23 @@ func newMemory(fd, size int) (*Memory, error) {
 		return nil, fmt.Errorf("setting up memory-mapped region: %w", err)
 	}
 
-	mm := &Memory{
-		b,
-		ro,
-		false,
-	}
-	runtime.SetFinalizer(mm, (*Memory).close)
+	mm := &Memory{b: b, ro: ro, heap: false}
+	mm.cleanup = runtime.AddCleanup(&mm, memoryCleanupFunc(), b)
 
 	return mm, nil
 }
 
-func (mm *Memory) close() {
-	if err := unix.Munmap(mm.b); err != nil {
-		panic(fmt.Errorf("unmapping memory: %w", err))
+func memoryCleanupFunc() func([]byte) {
+	return func(b []byte) {
+		if err := unix.Munmap(b); err != nil {
+			panic(fmt.Errorf("unmapping memory: %w", err))
+		}
 	}
+}
+
+func (mm *Memory) close() {
+	mm.cleanup.Stop()
+	memoryCleanupFunc()(mm.b)
 	mm.b = nil
 }
 

memory_test.go

@@ -4,7 +4,6 @@ import (
 	"io"
 	"math"
 	"os"
-	"runtime"
 	"testing"
 
 	"github.com/go-quicktest/qt"
@@ -88,9 +87,6 @@ func TestMemoryClose(t *testing.T) {
 	mm, err := mustMmapableArray(t, 0).Memory()
 	qt.Assert(t, qt.IsNil(err))
 
-	// Avoid unmap running twice.
-	runtime.SetFinalizer(mm, nil)
-
 	// unmap panics if the operation fails.
 	mm.close()
 }

memory_unsafe.go

@@ -87,6 +87,7 @@ func newUnsafeMemory(fd, size int) (*Memory, error) {
 		unsafe.Slice((*byte)(alloc), size),
 		ro,
 		true,
+		runtime.Cleanup{},
 	}
 
 	return mm, nil