src/runtime/mcleanup_test.go - go - Git at Google

 // Copyright 2024 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package runtime_test

 import (
 	"internal/runtime/atomic"
 	"runtime"
 	"sync"
 	"testing"
 	"time"
 	"unsafe"
 )

 func TestCleanup(t *testing.T) {
 	ch := make(chan bool, 1)
 	done := make(chan bool, 1)
 	want := 97531
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			if x != want {
 				t.Errorf("cleanup %d, want %d", x, want)
 			}
 			ch <- true
 		}
 		runtime.AddCleanup(v, cleanup, 97531)
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	<-ch
 }

 func TestCleanupMultiple(t *testing.T) {
 	ch := make(chan bool, 3)
 	done := make(chan bool, 1)
 	want := 97531
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			if x != want {
 				t.Errorf("cleanup %d, want %d", x, want)
 			}
 			ch <- true
 		}
 		runtime.AddCleanup(v, cleanup, 97531)
 		runtime.AddCleanup(v, cleanup, 97531)
 		runtime.AddCleanup(v, cleanup, 97531)
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	<-ch
 	<-ch
 	<-ch
 }

 func TestCleanupZeroSizedStruct(t *testing.T) {
 	type Z struct{}
 	z := new(Z)
 	runtime.AddCleanup(z, func(s string) {}, "foo")
 }

 func TestCleanupAfterFinalizer(t *testing.T) {
 	ch := make(chan int, 2)
 	done := make(chan bool, 1)
 	want := 97531
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		finalizer := func(x *int) {
 			ch <- 1
 		}
 		cleanup := func(x int) {
 			if x != want {
 				t.Errorf("cleanup %d, want %d", x, want)
 			}
 			ch <- 2
 		}
 		runtime.AddCleanup(v, cleanup, 97531)
 		runtime.SetFinalizer(v, finalizer)
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	var result int
 	result = <-ch
 	if result != 1 {
 		t.Errorf("result %d, want 1", result)
 	}
 	runtime.GC()
 	result = <-ch
 	if result != 2 {
 		t.Errorf("result %d, want 2", result)
 	}
 }

 func TestCleanupInteriorPointer(t *testing.T) {
 	ch := make(chan bool, 3)
 	done := make(chan bool, 1)
 	want := 97531
 	go func() {
 		// Allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			p unsafe.Pointer
 			i int
 			a int
 			b int
 			c int
 		}
 		ts := new(T)
 		ts.a = 97531
 		ts.b = 97531
 		ts.c = 97531
 		cleanup := func(x int) {
 			if x != want {
 				t.Errorf("cleanup %d, want %d", x, want)
 			}
 			ch <- true
 		}
 		runtime.AddCleanup(&ts.a, cleanup, 97531)
 		runtime.AddCleanup(&ts.b, cleanup, 97531)
 		runtime.AddCleanup(&ts.c, cleanup, 97531)
 		ts = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	<-ch
 	<-ch
 	<-ch
 }

 func TestCleanupStop(t *testing.T) {
 	done := make(chan bool, 1)
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			t.Error("cleanup called, want no cleanup called")
 		}
 		c := runtime.AddCleanup(v, cleanup, 97531)
 		c.Stop()
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 }

 func TestCleanupStopMultiple(t *testing.T) {
 	done := make(chan bool, 1)
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			t.Error("cleanup called, want no cleanup called")
 		}
 		c := runtime.AddCleanup(v, cleanup, 97531)
 		c.Stop()
 		c.Stop()
 		c.Stop()
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 }

 func TestCleanupStopinterleavedMultiple(t *testing.T) {
 	ch := make(chan bool, 3)
 	done := make(chan bool, 1)
 	go func() {
 		// allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			if x != 1 {
 				t.Error("cleanup called, want no cleanup called")
 			}
 			ch <- true
 		}
 		runtime.AddCleanup(v, cleanup, 1)
 		runtime.AddCleanup(v, cleanup, 2).Stop()
 		runtime.AddCleanup(v, cleanup, 1)
 		runtime.AddCleanup(v, cleanup, 2).Stop()
 		runtime.AddCleanup(v, cleanup, 1)
 		v = nil
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	<-ch
 	<-ch
 	<-ch
 }

 func TestCleanupStopAfterCleanupRuns(t *testing.T) {
 	ch := make(chan bool, 1)
 	done := make(chan bool, 1)
 	var stop func()
 	go func() {
 		// Allocate struct with pointer to avoid hitting tinyalloc.
 		// Otherwise we can't be sure when the allocation will
 		// be freed.
 		type T struct {
 			v int
 			p unsafe.Pointer
 		}
 		v := &new(T).v
 		*v = 97531
 		cleanup := func(x int) {
 			ch <- true
 		}
 		cl := runtime.AddCleanup(v, cleanup, 97531)
 		v = nil
 		stop = cl.Stop
 		done <- true
 	}()
 	<-done
 	runtime.GC()
 	<-ch
 	stop()
 }

 func TestCleanupPointerEqualsArg(t *testing.T) {
 	// See go.dev/issue/71316
 	defer func() {
 		want := "runtime.AddCleanup: ptr is equal to arg, cleanup will never run"
 		if r := recover(); r == nil {
 			t.Error("want panic, test did not panic")
 		} else if r == want {
 			// do nothing
 		} else {
 			t.Errorf("wrong panic: want=%q, got=%q", want, r)
 		}
 	}()

 	// allocate struct with pointer to avoid hitting tinyalloc.
 	// Otherwise we can't be sure when the allocation will
 	// be freed.
 	type T struct {
 		v int
 		p unsafe.Pointer
 	}
 	v := &new(T).v
 	*v = 97531
 	runtime.AddCleanup(v, func(x *int) {}, v)
 	v = nil
 	runtime.GC()
 }

 // Checks to make sure cleanups aren't lost when there are a lot of them.
 func TestCleanupLost(t *testing.T) {
 	type T struct {
 		v int
 		p unsafe.Pointer
 	}

 	cleanups := 10_000
 	if testing.Short() {
 		cleanups = 100
 	}
 	n := runtime.GOMAXPROCS(-1)
 	want := n * cleanups
 	var got atomic.Uint64
 	var wg sync.WaitGroup
 	for i := range n {
 		wg.Add(1)
 		go func(i int) {
 			defer wg.Done()

 			for range cleanups {
 				v := &new(T).v
 				*v = 97531
 				runtime.AddCleanup(v, func(_ int) {
 					got.Add(1)
 				}, 97531)
 			}
 		}(i)
 	}
 	wg.Wait()
 	runtime.GC()
 	timeout := 10 * time.Second
 	empty := runtime.BlockUntilEmptyCleanupQueue(int64(timeout))
 	if !empty {
 		t.Errorf("failed to drain cleanup queue within %s", timeout)
 	}

 	if got := int(got.Load()); got != want {
 		t.Errorf("%d cleanups executed, expected %d", got, want)
 	}
 }

 func TestCleanupUnreachablePointer(t *testing.T) {
 	defer func() {
 		if recover() == nil {
 			t.Error("AddCleanup failed to detect self-pointer")
 		}
 	}()

 	type T struct {
 		p *byte // use *byte to avoid tiny allocator
 		f int
 	}
 	v := &T{}
 	runtime.AddCleanup(v, func(*int) {
 		t.Error("cleanup ran unexpectedly")
 	}, &v.f)
 }

 func TestCleanupUnreachableClosure(t *testing.T) {
 	defer func() {
 		if recover() == nil {
 			t.Error("AddCleanup failed to detect closure pointer")
 		}
 	}()

 	type T struct {
 		p *byte // use *byte to avoid tiny allocator
 		f int
 	}
 	v := &T{}
 	runtime.AddCleanup(v, func(int) {
 		t.Log(v.f)
 		t.Error("cleanup ran unexpectedly")
 	}, 0)
 }

 // BenchmarkAddCleanupAndStop benchmarks adding and removing a cleanup
 // from the same allocation.
 //
 // At face value, this benchmark is unrealistic, since no program would
 // do this in practice. However, adding cleanups to new allocations in a
 // loop is also unrealistic. It adds additional unused allocations,
 // exercises uncommon performance pitfalls in AddCleanup (traversing the
 // specials list, which should just be its own benchmark), and executing
 // cleanups at a frequency that is unlikely to appear in real programs.
 //
 // This benchmark is still useful however, since we can get a low-noise
 // measurement of the cost of AddCleanup and Stop all in one without the
 // above pitfalls: we can measure the pure overhead. We can then separate
 // out the cost of each in CPU profiles if we so choose (they're not so
 // inexpensive as to make this infeasible).
 func BenchmarkAddCleanupAndStop(b *testing.B) {
 	b.ReportAllocs()

 	type T struct {
 		v int
 		p unsafe.Pointer
 	}
 	x := new(T)
 	for b.Loop() {
 		runtime.AddCleanup(x, func(int) {}, 14).Stop()
 	}
 }
	// Copyright 2024 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package runtime_test

	import (
	"internal/runtime/atomic"
	"runtime"
	"sync"
	"testing"
	"time"
	"unsafe"
	)

	func TestCleanup(t *testing.T) {
	ch := make(chan bool, 1)
	done := make(chan bool, 1)
	want := 97531
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	if x != want {
	t.Errorf("cleanup %d, want %d", x, want)
	}
	ch <- true
	}
	runtime.AddCleanup(v, cleanup, 97531)
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	<-ch
	}

	func TestCleanupMultiple(t *testing.T) {
	ch := make(chan bool, 3)
	done := make(chan bool, 1)
	want := 97531
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	if x != want {
	t.Errorf("cleanup %d, want %d", x, want)
	}
	ch <- true
	}
	runtime.AddCleanup(v, cleanup, 97531)
	runtime.AddCleanup(v, cleanup, 97531)
	runtime.AddCleanup(v, cleanup, 97531)
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	<-ch
	<-ch
	<-ch
	}

	func TestCleanupZeroSizedStruct(t *testing.T) {
	type Z struct{}
	z := new(Z)
	runtime.AddCleanup(z, func(s string) {}, "foo")
	}

	func TestCleanupAfterFinalizer(t *testing.T) {
	ch := make(chan int, 2)
	done := make(chan bool, 1)
	want := 97531
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	finalizer := func(x *int) {
	ch <- 1
	}
	cleanup := func(x int) {
	if x != want {
	t.Errorf("cleanup %d, want %d", x, want)
	}
	ch <- 2
	}
	runtime.AddCleanup(v, cleanup, 97531)
	runtime.SetFinalizer(v, finalizer)
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	var result int
	result = <-ch
	if result != 1 {
	t.Errorf("result %d, want 1", result)
	}
	runtime.GC()
	result = <-ch
	if result != 2 {
	t.Errorf("result %d, want 2", result)
	}
	}

	func TestCleanupInteriorPointer(t *testing.T) {
	ch := make(chan bool, 3)
	done := make(chan bool, 1)
	want := 97531
	go func() {
	// Allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	p unsafe.Pointer
	i int
	a int
	b int
	c int
	}
	ts := new(T)
	ts.a = 97531
	ts.b = 97531
	ts.c = 97531
	cleanup := func(x int) {
	if x != want {
	t.Errorf("cleanup %d, want %d", x, want)
	}
	ch <- true
	}
	runtime.AddCleanup(&ts.a, cleanup, 97531)
	runtime.AddCleanup(&ts.b, cleanup, 97531)
	runtime.AddCleanup(&ts.c, cleanup, 97531)
	ts = nil
	done <- true
	}()
	<-done
	runtime.GC()
	<-ch
	<-ch
	<-ch
	}

	func TestCleanupStop(t *testing.T) {
	done := make(chan bool, 1)
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	t.Error("cleanup called, want no cleanup called")
	}
	c := runtime.AddCleanup(v, cleanup, 97531)
	c.Stop()
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	}

	func TestCleanupStopMultiple(t *testing.T) {
	done := make(chan bool, 1)
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	t.Error("cleanup called, want no cleanup called")
	}
	c := runtime.AddCleanup(v, cleanup, 97531)
	c.Stop()
	c.Stop()
	c.Stop()
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	}

	func TestCleanupStopinterleavedMultiple(t *testing.T) {
	ch := make(chan bool, 3)
	done := make(chan bool, 1)
	go func() {
	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	if x != 1 {
	t.Error("cleanup called, want no cleanup called")
	}
	ch <- true
	}
	runtime.AddCleanup(v, cleanup, 1)
	runtime.AddCleanup(v, cleanup, 2).Stop()
	runtime.AddCleanup(v, cleanup, 1)
	runtime.AddCleanup(v, cleanup, 2).Stop()
	runtime.AddCleanup(v, cleanup, 1)
	v = nil
	done <- true
	}()
	<-done
	runtime.GC()
	<-ch
	<-ch
	<-ch
	}

	func TestCleanupStopAfterCleanupRuns(t *testing.T) {
	ch := make(chan bool, 1)
	done := make(chan bool, 1)
	var stop func()
	go func() {
	// Allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	cleanup := func(x int) {
	ch <- true
	}
	cl := runtime.AddCleanup(v, cleanup, 97531)
	v = nil
	stop = cl.Stop
	done <- true
	}()
	<-done
	runtime.GC()
	<-ch
	stop()
	}

	func TestCleanupPointerEqualsArg(t *testing.T) {
	// See go.dev/issue/71316
	defer func() {
	want := "runtime.AddCleanup: ptr is equal to arg, cleanup will never run"
	if r := recover(); r == nil {
	t.Error("want panic, test did not panic")
	} else if r == want {
	// do nothing
	} else {
	t.Errorf("wrong panic: want=%q, got=%q", want, r)
	}
	}()

	// allocate struct with pointer to avoid hitting tinyalloc.
	// Otherwise we can't be sure when the allocation will
	// be freed.
	type T struct {
	v int
	p unsafe.Pointer
	}
	v := &new(T).v
	*v = 97531
	runtime.AddCleanup(v, func(x *int) {}, v)
	v = nil
	runtime.GC()
	}

	// Checks to make sure cleanups aren't lost when there are a lot of them.
	func TestCleanupLost(t *testing.T) {
	type T struct {
	v int
	p unsafe.Pointer
	}

	cleanups := 10_000
	if testing.Short() {
	cleanups = 100
	}
	n := runtime.GOMAXPROCS(-1)
	want := n * cleanups
	var got atomic.Uint64
	var wg sync.WaitGroup
	for i := range n {
	wg.Add(1)
	go func(i int) {
	defer wg.Done()

	for range cleanups {
	v := &new(T).v
	*v = 97531
	runtime.AddCleanup(v, func(_ int) {
	got.Add(1)
	}, 97531)
	}
	}(i)
	}
	wg.Wait()
	runtime.GC()
	timeout := 10 * time.Second
	empty := runtime.BlockUntilEmptyCleanupQueue(int64(timeout))
	if !empty {
	t.Errorf("failed to drain cleanup queue within %s", timeout)
	}

	if got := int(got.Load()); got != want {
	t.Errorf("%d cleanups executed, expected %d", got, want)
	}
	}

	func TestCleanupUnreachablePointer(t *testing.T) {
	defer func() {
	if recover() == nil {
	t.Error("AddCleanup failed to detect self-pointer")
	}
	}()

	type T struct {
	p byte // use byte to avoid tiny allocator
	f int
	}
	v := &T{}
	runtime.AddCleanup(v, func(*int) {
	t.Error("cleanup ran unexpectedly")
	}, &v.f)
	}

	func TestCleanupUnreachableClosure(t *testing.T) {
	defer func() {
	if recover() == nil {
	t.Error("AddCleanup failed to detect closure pointer")
	}
	}()

	type T struct {
	p byte // use byte to avoid tiny allocator
	f int
	}
	v := &T{}
	runtime.AddCleanup(v, func(int) {
	t.Log(v.f)
	t.Error("cleanup ran unexpectedly")
	}, 0)
	}

	// BenchmarkAddCleanupAndStop benchmarks adding and removing a cleanup
	// from the same allocation.
	//
	// At face value, this benchmark is unrealistic, since no program would
	// do this in practice. However, adding cleanups to new allocations in a
	// loop is also unrealistic. It adds additional unused allocations,
	// exercises uncommon performance pitfalls in AddCleanup (traversing the
	// specials list, which should just be its own benchmark), and executing
	// cleanups at a frequency that is unlikely to appear in real programs.
	//
	// This benchmark is still useful however, since we can get a low-noise
	// measurement of the cost of AddCleanup and Stop all in one without the
	// above pitfalls: we can measure the pure overhead. We can then separate
	// out the cost of each in CPU profiles if we so choose (they're not so
	// inexpensive as to make this infeasible).
	func BenchmarkAddCleanupAndStop(b *testing.B) {
	b.ReportAllocs()

	type T struct {
	v int
	p unsafe.Pointer
	}
	x := new(T)
	for b.Loop() {
	runtime.AddCleanup(x, func(int) {}, 14).Stop()
	}
	}