all: use runtime.AddCleanup

Convert almost all uses of runtime.SetFinalizer to runtime.AddCleanup.
The exception is the unsafe memory code, which doesn't seem that easy
to convert.

Signed-off-by: Lorenz Bauer <[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) {
+		runtime.AddCleanup(spec.decoder.sharedBuf, func(raw []byte) {
 			_ = unix.Munmap(raw)
-		})
+		}, 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(epollFd int) {
+		_ = unix.Close(epollFd)
+	}, 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

@@ -199,7 +199,8 @@ type Event struct {
 	typ         ProbeType
 	group, name string
 	// event id allocated by the kernel. 0 if the event has already been removed.
-	id uint64
+	id      uint64
+	cleanup runtime.Cleanup
 }
 
 // NewEvent creates a new ephemeral trace event.
@@ -316,8 +317,17 @@ 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}
+
+	type cleanupData struct {
+		typ ProbeType
+		pe  string
+	}
+
+	evt.cleanup = runtime.AddCleanup(evt, func(cleanupData cleanupData) {
+		_ = removeEvent(cleanupData.typ, pe)
+	}, cleanupData{args.Type, fmt.Sprintf("%s/%s", evt.group, evt.name)})
+
 	return evt, nil
 }
 
@@ -330,7 +340,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)
 }

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

@@ -19,6 +19,7 @@ type ringbufEventRing struct {
 	prod []byte
 	cons []byte
 	*ringReader
+	cleanup runtime.Cleanup
 }
 
 func newRingBufEventRing(mapFD, size int) (*ringbufEventRing, error) {
@@ -41,13 +42,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(mappings struct{ prod, cons []byte }) {
+		_ = unix.Munmap(mappings.prod)
+		_ = unix.Munmap(mappings.cons)
+	}, struct{ prod, cons []byte }{prod, cons})
 
 	return ring, nil
 }
 
 func (ring *ringbufEventRing) Close() {
-	runtime.SetFinalizer(ring, nil)
+	ring.cleanup.Stop()
 
 	_ = unix.Munmap(ring.prod)
 	_ = unix.Munmap(ring.cons)