Improve winpowrprof callback

Author: nekohasekai

common/oncefunc.go

@@ -0,0 +1,17 @@
+//go:build go1.21
+
+package common
+
+import "sync"
+
+func OnceFunc(f func()) func() {
+	return sync.OnceFunc(f)
+}
+
+func OnceValue[T any](f func() T) func() T {
+	return sync.OnceValue(f)
+}
+
+func OnceValues[T1, T2 any](f func() (T1, T2)) func() (T1, T2) {
+	return sync.OnceValues(f)
+}

common/oncefunc_compat.go

@@ -0,0 +1,106 @@
+//go:build !go1.21
+
+// Copyright 2022 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.
+
+// copied from go1.22.5
+
+package common
+
+import "sync"
+
+// OnceFunc returns a function that invokes f only once. The returned function
+// may be called concurrently.
+//
+// If f panics, the returned function will panic with the same value on every call.
+func OnceFunc(f func()) func() {
+	var (
+		once  sync.Once
+		valid bool
+		p     any
+	)
+	// Construct the inner closure just once to reduce costs on the fast path.
+	g := func() {
+		defer func() {
+			p = recover()
+			if !valid {
+				// Re-panic immediately so on the first call the user gets a
+				// complete stack trace into f.
+				panic(p)
+			}
+		}()
+		f()
+		f = nil      // Do not keep f alive after invoking it.
+		valid = true // Set only if f does not panic.
+	}
+	return func() {
+		once.Do(g)
+		if !valid {
+			panic(p)
+		}
+	}
+}
+
+// OnceValue returns a function that invokes f only once and returns the value
+// returned by f. The returned function may be called concurrently.
+//
+// If f panics, the returned function will panic with the same value on every call.
+func OnceValue[T any](f func() T) func() T {
+	var (
+		once   sync.Once
+		valid  bool
+		p      any
+		result T
+	)
+	g := func() {
+		defer func() {
+			p = recover()
+			if !valid {
+				panic(p)
+			}
+		}()
+		result = f()
+		f = nil
+		valid = true
+	}
+	return func() T {
+		once.Do(g)
+		if !valid {
+			panic(p)
+		}
+		return result
+	}
+}
+
+// OnceValues returns a function that invokes f only once and returns the values
+// returned by f. The returned function may be called concurrently.
+//
+// If f panics, the returned function will panic with the same value on every call.
+func OnceValues[T1, T2 any](f func() (T1, T2)) func() (T1, T2) {
+	var (
+		once  sync.Once
+		valid bool
+		p     any
+		r1    T1
+		r2    T2
+	)
+	g := func() {
+		defer func() {
+			p = recover()
+			if !valid {
+				panic(p)
+			}
+		}()
+		r1, r2 = f()
+		f = nil
+		valid = true
+	}
+	return func() (T1, T2) {
+		once.Do(g)
+		if !valid {
+			panic(p)
+		}
+		return r1, r2
+	}
+}

common/winpowrprof/event_windows.go

@@ -3,134 +3,89 @@ package winpowrprof
 // modify from https://github.com/golang/go/blob/b634f6fdcbebee23b7da709a243f3db217b64776/src/runtime/os_windows.go#L257
 
 import (
-	"sync"
 	"syscall"
 	"unsafe"
 
-	"github.com/sagernet/sing/common/x/list"
+	"github.com/sagernet/sing/common"
 
 	"golang.org/x/sys/windows"
 )
 
-type powerEventListener struct {
-	element *list.Element[EventCallback]
-}
-
-func NewEventListener(callback EventCallback) (EventListener, error) {
-	err := initCallback()
-	if err != nil {
-		return nil, err
-	}
-	access.Lock()
-	defer access.Unlock()
-	return &powerEventListener{
-		element: callbackList.PushBack(callback),
-	}, nil
-}
-
-func (l *powerEventListener) Start() error {
-	access.Lock()
-	defer access.Unlock()
-	if handle != 0 {
-		return nil
-	}
-	return startListener()
-}
-
-func (l *powerEventListener) Close() error {
-	access.Lock()
-	defer access.Unlock()
-	if l.element != nil {
-		callbackList.Remove(l.element)
-	}
-	if callbackList.Len() > 0 {
-		return nil
-	}
-	return closeListener()
-}
-
 var (
 	modpowerprof                                 = windows.NewLazySystemDLL("powrprof.dll")
 	procPowerRegisterSuspendResumeNotification   = modpowerprof.NewProc("PowerRegisterSuspendResumeNotification")
 	procPowerUnregisterSuspendResumeNotification = modpowerprof.NewProc("PowerUnregisterSuspendResumeNotification")
 )
 
-var (
-	access           sync.Mutex
-	callbackList     list.List[EventCallback]
-	initCallbackOnce sync.Once
-	rawCallback      uintptr
-	handle           uintptr
-)
+var suspendResumeNotificationCallback = common.OnceValue(func() uintptr {
+	return windows.NewCallback(func(context *EventCallback, changeType uint32, setting uintptr) uintptr {
+		callback := *context
+		const (
+			PBT_APMSUSPEND         uint32 = 4
+			PBT_APMRESUMESUSPEND   uint32 = 7
+			PBT_APMRESUMEAUTOMATIC uint32 = 18
+		)
+		var event int
+		switch changeType {
+		case PBT_APMSUSPEND:
+			event = EVENT_SUSPEND
+		case PBT_APMRESUMESUSPEND:
+			event = EVENT_RESUME
+		case PBT_APMRESUMEAUTOMATIC:
+			event = EVENT_RESUME_AUTOMATIC
+		default:
+			return 0
+		}
+		callback(event)
+		return 0
+	})
+})
 
-func initCallback() error {
+type powerEventListener struct {
+	callback EventCallback
+	handle   uintptr
+}
+
+func NewEventListener(callback EventCallback) (EventListener, error) {
 	err := procPowerRegisterSuspendResumeNotification.Find()
 	if err != nil {
-		return err // Running on Windows 7, where we don't need it anyway.
+		return nil, err // Running on Windows 7, where we don't need it anyway.
 	}
 	err = procPowerUnregisterSuspendResumeNotification.Find()
 	if err != nil {
-		return err // Running on Windows 7, where we don't need it anyway.
+		return nil, err // Running on Windows 7, where we don't need it anyway.
 	}
-	initCallbackOnce.Do(func() {
-		rawCallback = windows.NewCallback(func(context uintptr, changeType uint32, setting uintptr) uintptr {
-			const (
-				PBT_APMSUSPEND         uint32 = 4
-				PBT_APMRESUMESUSPEND   uint32 = 7
-				PBT_APMRESUMEAUTOMATIC uint32 = 18
-			)
-			var event int
-			switch changeType {
-			case PBT_APMSUSPEND:
-				event = EVENT_SUSPEND
-			case PBT_APMRESUMESUSPEND:
-				event = EVENT_RESUME
-			case PBT_APMRESUMEAUTOMATIC:
-				event = EVENT_RESUME_AUTOMATIC
-			default:
-				return 0
-			}
-			access.Lock()
-			callbacks := callbackList.Array()
-			access.Unlock()
-			for _, callback := range callbacks {
-				callback(event)
-			}
-			return 0
-		})
-	})
-	return nil
+	return &powerEventListener{
+		callback: callback,
+	}, nil
 }
 
-func startListener() error {
+func (l *powerEventListener) Start() error {
 	type DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS struct {
 		callback uintptr
-		context  uintptr
+		context  unsafe.Pointer
 	}
 	const DEVICE_NOTIFY_CALLBACK = 2
 	params := DEVICE_NOTIFY_SUBSCRIBE_PARAMETERS{
-		callback: rawCallback,
+		callback: suspendResumeNotificationCallback(),
+		context:  unsafe.Pointer(&l.callback),
 	}
 	_, _, errno := syscall.SyscallN(
 		procPowerRegisterSuspendResumeNotification.Addr(),
 		DEVICE_NOTIFY_CALLBACK,
 		uintptr(unsafe.Pointer(&params)),
-		uintptr(unsafe.Pointer(&handle)),
+		uintptr(unsafe.Pointer(&l.handle)),
 	)
 	if errno != 0 {
 		return errno
 	}
 	return nil
 }
 
-func closeListener() error {
-	if handle == 0 {
-		return nil
-	}
-	_, _, errno := syscall.SyscallN(procPowerUnregisterSuspendResumeNotification.Addr(), uintptr(unsafe.Pointer(&handle)))
+func (l *powerEventListener) Close() error {
+	_, _, errno := syscall.SyscallN(procPowerUnregisterSuspendResumeNotification.Addr(), uintptr(unsafe.Pointer(&l.handle)))
 	if errno != 0 {
 		return errno
 	}
-	handle = 0
 	return nil
 }