Convert uses of runtime.SetFinalizer to runtime.AddCleanup

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)
 }

link/perf_event.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"runtime"
 	"unsafe"
 
 	"github.com/cilium/ebpf"
@@ -61,15 +60,11 @@ type perfEvent struct {
 
 func newPerfEvent(fd *sys.FD, event *tracefs.Event) *perfEvent {
 	pe := &perfEvent{event, fd}
-	// Both event and fd have their own finalizer, but we want to
-	// guarantee that they are closed in a certain order.
-	runtime.SetFinalizer(pe, (*perfEvent).Close)
 	return pe
 }
 
 func (pe *perfEvent) Close() error {
-	runtime.SetFinalizer(pe, nil)
-
+	// We close the perf event before attempting to remove the tracefs event.
 	if err := pe.fd.Close(); err != nil {
 		return fmt.Errorf("closing perf event fd: %w", err)
 	}

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/reader.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"runtime"
 	"sync"
 	"time"
 
@@ -261,7 +260,6 @@ func NewReaderWithOptions(array *ebpf.Map, perCPUBuffer int, opts ReaderOptions)
 	if err = pr.Resume(); err != nil {
 		return nil, err
 	}
-	runtime.SetFinalizer(pr, (*Reader).Close)
 	return pr, 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)

ringbuf/ring_other.go

@@ -16,6 +16,7 @@ type ringbufEventRing struct {
 	prod []byte
 	cons []byte
 	*ringReader
+	cleanup runtime.Cleanup
 }
 
 func newRingBufEventRing(mapFD, size int) (*ringbufEventRing, error) {
@@ -38,13 +39,16 @@ func newRingBufEventRing(mapFD, size int) (*ringbufEventRing, error) {
 		cons:       cons,
 		ringReader: newRingReader(cons_pos, prod_pos, prod[os.Getpagesize():]),
 	}
-	runtime.SetFinalizer(ring, (*ringbufEventRing).Close)
+	ring.cleanup = runtime.AddCleanup(ring, func(*byte) {
+		_ = unix.Munmap(prod)
+		_ = unix.Munmap(cons)
+	}, nil)
 
 	return ring, nil
 }
 
 func (ring *ringbufEventRing) Close() error {
-	runtime.SetFinalizer(ring, nil)
+	ring.cleanup.Stop()
 
 	prod, cons := ring.prod, ring.cons
 	ring.prod, ring.cons = nil, nil

ringbuf/ring_windows.go

@@ -14,6 +14,8 @@ type ringbufEventRing struct {
 	mapFd            *sys.FD
 	cons, prod, data *uint8
 	*ringReader
+
+	cleanup runtime.Cleanup
 }
 
 func newRingBufEventRing(mapFD, size int) (*ringbufEventRing, error) {
@@ -51,13 +53,15 @@ func newRingBufEventRing(mapFD, size int) (*ringbufEventRing, error) {
 		data:       dataPtr,
 		ringReader: newRingReader(consPos, prodPos, data),
 	}
-	runtime.SetFinalizer(ring, (*ringbufEventRing).Close)
+	ring.cleanup = runtime.AddCleanup(ring, func(*byte) {
+		efw.EbpfRingBufferMapUnmapBuffer(fd.Int(), consPtr, prodPtr, dataPtr)
+	}, nil)
 
 	return ring, nil
 }
 
 func (ring *ringbufEventRing) Close() error {
-	runtime.SetFinalizer(ring, nil)
+	ring.cleanup.Stop()
 
 	return errors.Join(
 		efw.EbpfRingBufferMapUnmapBuffer(ring.mapFd.Int(), ring.cons, ring.prod, ring.data),