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: Lorenz Bauer <[email protected]>
Co-authored-by: Timo Beckers <[email protected]>

Author: lmb

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/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

perf/ring.go

@@ -22,6 +22,7 @@ type perfEventRing struct {
 	cpu  int
 	mmap []byte
 	ringReader
+	cleanup runtime.Cleanup
 }
 
 func newPerfEventRing(cpu, perCPUBuffer int, opts ReaderOptions) (_ *sys.FD, _ *perfEventRing, err error) {
@@ -73,7 +74,9 @@ func newPerfEventRing(cpu, perCPUBuffer int, opts ReaderOptions) (_ *sys.FD, _ *
 		mmap:       mmap,
 		ringReader: reader,
 	}
-	runtime.SetFinalizer(ring, (*perfEventRing).Close)
+	ring.cleanup = runtime.AddCleanup(ring, func(mmap []byte) {
+		_ = unix.Munmap(mmap)
+	}, ring.mmap)
 
 	return fd, ring, nil
 }
@@ -95,7 +98,7 @@ func perfBufferSize(perCPUBuffer int) int {
 }
 
 func (ring *perfEventRing) Close() error {
-	runtime.SetFinalizer(ring, nil)
+	ring.cleanup.Stop()
 	mmap := ring.mmap
 	ring.mmap = nil
 	return unix.Munmap(mmap)