优化exsync.Once，使其性能接近sync.Once

Author: thinkeridea

exsync/benchmark/once_test.go

@@ -0,0 +1,50 @@
+// MIT License
+// Copyright (c) 2020 Qi Yin <[email protected]>
+
+package benchmark
+
+import (
+	"sync"
+	"testing"
+	"unsafe"
+
+	"github.com/thinkeridea/go-extend/exsync"
+)
+
+type one int
+
+func (o *one) Increment() {
+	*o++
+}
+
+func BenchmarkSyncOnce(b *testing.B) {
+	var once sync.Once
+	f := func() {
+		_ = new(one)
+	}
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			once.Do(f)
+		}
+	})
+}
+
+func BenchmarkOnce(b *testing.B) {
+	var once exsync.Once
+	f := func() interface{} { return new(one) }
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_ = once.Do(f).(*one)
+		}
+	})
+}
+
+func BenchmarkOncePointer(b *testing.B) {
+	var once exsync.OncePointer
+	f := func() unsafe.Pointer { return unsafe.Pointer(new(one)) }
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_ = (*one)(once.Do(f))
+		}
+	})
+}

exsync/once.go

@@ -5,6 +5,7 @@ package exsync
 
 import (
 	"sync"
+	"sync/atomic"
 	"unsafe"
 )
 
@@ -30,9 +31,15 @@ import (
 // 	return res[0].(*mysql.Client), res[1].(error)
 // }
 //
-// 使用该方法需要一些取舍，它简单实用，相比 sync.Once 性能有所下降，不过它依然很快，这不会形成性能问题。
+// 使用该方法需要一些取舍，它简单实用，性能无限接近 sync.Once。
 type Once struct {
-	once sync.Once
+	// done indicates whether the action has been performed.
+	// It is first in the struct because it is used in the hot path.
+	// The hot path is inlined at every call site.
+	// Placing done first allows more compact instructions on some architectures (amd64/x86),
+	// and fewer instructions (to calculate offset) on other architectures.
+	done uint32
+	m    sync.Mutex
 	v    interface{}
 }
 
@@ -55,16 +62,26 @@ type Once struct {
 // without calling f.
 //
 func (o *Once) Do(f func() interface{}) interface{} {
-	o.once.Do(func() {
-		o.v = f()
-	})
+	if atomic.LoadUint32(&o.done) == 0 {
+		o.doSlow(f)
+	}
 
 	return o.v
 }
 
-// OncePointer 性能方面略好于 Once，但不会有太大改善，依然落后于 sync.Once， 在某些场景下可以使用，更推荐使用 Once
+func (o *Once) doSlow(f func() interface{}) {
+	o.m.Lock()
+	defer o.m.Unlock()
+	if o.done == 0 {
+		defer atomic.StoreUint32(&o.done, 1)
+		o.v = f()
+	}
+}
+
+// OncePointer 性能方面略好于 Once，但不会有太大改善， 在某些场景下可以使用，更推荐使用 Once
 type OncePointer struct {
-	once sync.Once
+	done uint32
+	m    sync.Mutex
 	v    unsafe.Pointer
 }
 
@@ -87,9 +104,18 @@ type OncePointer struct {
 // without calling f.
 //
 func (o *OncePointer) Do(f func() unsafe.Pointer) unsafe.Pointer {
-	o.once.Do(func() {
-		o.v = f()
-	})
+	if atomic.LoadUint32(&o.done) == 0 {
+		o.doSlow(f)
+	}
 
 	return o.v
 }
+
+func (o *OncePointer) doSlow(f func() unsafe.Pointer) {
+	o.m.Lock()
+	defer o.m.Unlock()
+	if o.done == 0 {
+		defer atomic.StoreUint32(&o.done, 1)
+		o.v = f()
+	}
+}

exsync/once_test.go

@@ -63,17 +63,6 @@ func TestOncePanic(t *testing.T) {
 	})
 }
 
-func BenchmarkOnce(b *testing.B) {
-	var once Once
-	var o = new(one)
-	f := func() interface{} { return o }
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			_ = once.Do(f).(*one)
-		}
-	})
-}
-
 func runPointer(t *testing.T, once *OncePointer, o *one, c chan bool) {
 	v := once.Do(func() unsafe.Pointer {
 		o.Increment()
@@ -122,16 +111,3 @@ func TestOncePointerPanic(t *testing.T) {
 		return nil
 	})
 }
-
-func BenchmarkOncePointer(b *testing.B) {
-	var once OncePointer
-	var o = new(one)
-	f := func() unsafe.Pointer {
-		return unsafe.Pointer(o)
-	}
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			_ = (*one)(once.Do(f))
-		}
-	})
-}