Add package cleanup

Author: nekohasekai

common/cleanup/cleanup.go

@@ -0,0 +1,94 @@
+package cleanup
+
+import (
+	"runtime"
+	"sync"
+	_ "unsafe"
+)
+
+func init() {
+	registerPoolCleanup(myCleanup)
+}
+
+//go:linkname registerPoolCleanup sync.runtime_registerPoolCleanup
+func registerPoolCleanup(cleanup func())
+
+//go:linkname poolCleanup sync.poolCleanup
+func poolCleanup()
+
+var unsafeCleanupFuncs []func()
+
+func myCleanup() {
+	poolCleanup()
+	for _, cleanupFunc := range unsafeCleanupFuncs {
+		cleanupFunc()
+	}
+}
+
+// AddUnsafe must be called only in init {}
+// called in STW, must not allocate or hold a lock
+func AddUnsafe(cleanup func()) {
+	unsafeCleanupFuncs = append(unsafeCleanupFuncs, cleanup)
+}
+
+type Cleaner struct {
+	access sync.Mutex
+	state  *cleanupState
+}
+
+func Add(cleanup func()) *Cleaner {
+	state := &cleanupState{
+		cleanupFunc: cleanup,
+	}
+	newObject(state)
+	return &Cleaner{
+		state: state,
+	}
+}
+
+func (c *Cleaner) Close() {
+	c.access.Lock()
+	defer c.access.Unlock()
+	if c.state == nil {
+		return
+	}
+	c.state.close()
+	c.state = nil
+}
+
+type cleanupState struct {
+	access      sync.Mutex
+	cleanupFunc func()
+	closed      bool
+}
+
+func (s *cleanupState) close() {
+	s.access.Lock()
+	defer s.access.Unlock()
+	s.closed = true
+}
+
+type Object struct {
+	state *cleanupState
+}
+
+func newObject(state *cleanupState) {
+	object := &Object{
+		state: state,
+	}
+	runtime.SetFinalizer(object, (*Object).cleanup)
+}
+
+func (o *Object) cleanup() {
+	state := o.state
+	state.access.Lock()
+	if state.closed {
+		state.access.Unlock()
+		return
+	}
+
+	newObject(state)
+	cleanupFunc := state.cleanupFunc
+	state.access.Unlock()
+	cleanupFunc()
+}

common/cleanup/cleanup_test.go

@@ -0,0 +1,166 @@
+//nolint:paralleltest
+package cleanup
+
+import (
+	"runtime"
+	"runtime/debug"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCleanup(t *testing.T) {
+	var didCleanup atomic.Int32
+	Add(func() {
+		didCleanup.Add(1)
+	})
+	runCleanup(t, func() bool {
+		return didCleanup.Load() >= 1
+	})
+	runCleanup(t, func() bool {
+		return didCleanup.Load() >= 2
+	})
+}
+
+func TestCleanupUnsafe(t *testing.T) {
+	var didCleanup atomic.Int32
+	AddUnsafe(func() {
+		didCleanup.Add(1)
+	})
+	runCleanup(t, func() bool {
+		return didCleanup.Load() >= 1
+	})
+	runCleanup(t, func() bool {
+		return didCleanup.Load() >= 2
+	})
+}
+
+func TestCleanup1(t *testing.T) {
+	var didCleanup atomic.Bool
+	var didReset atomic.Bool
+	m := newMap(func() {
+		didCleanup.Store(true)
+	})
+	Add(func() {
+		*m = sync.Map{}
+		didReset.Store(true)
+	})
+	runtime.KeepAlive(&m)
+	runCleanup(t, func() bool {
+		return didReset.Load()
+	})
+	require.False(t, didCleanup.Load())
+	runCleanup(t, func() bool {
+		return didCleanup.Load()
+	})
+}
+
+func TestCleanup1Unsafe(t *testing.T) {
+	var didCleanup atomic.Bool
+	m := newMap(func() {
+		didCleanup.Store(true)
+	})
+	AddUnsafe(func() {
+		*m = sync.Map{}
+	})
+	runtime.KeepAlive(&m)
+	runCleanup(t, func() bool {
+		return didCleanup.Load()
+	})
+}
+
+type myObj struct {
+	_ string
+}
+
+func newMap(cleanup func()) *sync.Map {
+	obj := &myObj{}
+	runtime.SetFinalizer(obj, func(*myObj) {
+		cleanup()
+	})
+	var m sync.Map
+	m.Store("test", obj)
+	return &m
+}
+
+func TestSafeCleanup(t *testing.T) {
+	var didCleanup atomic.Int32
+	safeObject := Add(func() {
+		didCleanup.Add(1)
+	})
+	runCleanup(t, func() bool {
+		return didCleanup.Load() >= 1
+	})
+	safeObject.Close()
+	didCleanupAfterClose := didCleanup.Load()
+	debug.FreeOSMemory()
+	require.Never(t, func() bool {
+		return didCleanup.Load() != didCleanupAfterClose
+	}, cleanupWait, cleanupTick)
+}
+
+func TestCloseSkipsQueuedCleanup(t *testing.T) {
+	blockCleanup := make(chan struct{})
+	cleanupBlocked := make(chan struct{})
+	var blockOnce sync.Once
+	blocker := Add(func() {
+		blockOnce.Do(func() {
+			close(cleanupBlocked)
+			<-blockCleanup
+		})
+	})
+	runCleanup(t, func() bool {
+		select {
+		case <-cleanupBlocked:
+			return true
+		default:
+			return false
+		}
+	})
+
+	var didCleanup atomic.Bool
+	queued := Add(func() {
+		didCleanup.Store(true)
+	})
+	debug.FreeOSMemory()
+	queued.Close()
+
+	close(blockCleanup)
+	blocker.Close()
+	require.Never(t, didCleanup.Load, cleanupWait, cleanupTick)
+}
+
+func TestCleanupCanCloseItself(t *testing.T) {
+	ready := make(chan struct{})
+	done := make(chan struct{})
+	var cleaner *Cleaner
+	cleaner = Add(func() {
+		<-ready
+		cleaner.Close()
+		close(done)
+	})
+	close(ready)
+
+	runCleanup(t, func() bool {
+		select {
+		case <-done:
+			return true
+		default:
+			return false
+		}
+	})
+}
+
+const (
+	cleanupWait = time.Second
+	cleanupTick = time.Millisecond
+)
+
+func runCleanup(t *testing.T, condition func() bool) {
+	t.Helper()
+	debug.FreeOSMemory()
+	require.Eventually(t, condition, cleanupWait, cleanupTick)
+}