sync: reduce OnceFunc (and variants) heap allocations

The lifetime of the variables are identical; capture
them in a single struct to avoid individual allocations.
The inner closure can also avoid allocation by using the
capture of the outer closure.

Escape analysis for OnceValues:

  /go/src/sync/oncefunc.go:74:29: moved to heap: sync.f
  /go/src/sync/oncefunc.go:76:3: moved to heap: sync.once
  /go/src/sync/oncefunc.go:77:3: moved to heap: sync.valid
  /go/src/sync/oncefunc.go:78:3: moved to heap: sync.p
  /go/src/sync/oncefunc.go:79:3: moved to heap: sync.r1
  /go/src/sync/oncefunc.go:80:3: moved to heap: sync.r2
  /go/src/sync/oncefunc.go:82:7: func literal escapes to heap
  /go/src/sync/oncefunc.go:83:9: func literal does not escape
  /go/src/sync/oncefunc.go:93:9: func literal escapes to heap

After provided changes:

  /go/src/sync/oncefunc.go:86:2: moved to heap: sync.d
  /go/src/sync/oncefunc.go:96:9: func literal escapes to heap
  /go/src/sync/oncefunc.go:99:13: func literal does not escape
  /go/src/sync/oncefunc.go:100:10: func literal does not escape

Change-Id: Ib06e650fd427b57e0bdbdf1fe759fe436104ff79
Reviewed-on: https://go-review.googlesource.com/c/go/+/601596
Auto-Submit: Austin Clements <austin@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Pratt <mpratt@google.com>
Reviewed-by: Austin Clements <austin@google.com>

2 files changed, 174 insertions, 58 deletions

diff --git a/src/sync/oncefunc.go b/src/sync/oncefunc.go
index db286283d1..2c49efeef8 100644
--- a/src/sync/oncefunc.go
+++ b/src/sync/oncefunc.go
@@ -9,29 +9,32 @@ package sync
 //
 // If f panics, the returned function will panic with the same value on every call.
 func OnceFunc(f func()) func() {
-	var (
+	// Use a struct so that there's a single heap allocation.
+	d := struct {
+		f     func()
 		once  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.
+	}{
+		f: f,
 	}
 	return func() {
-		once.Do(g)
-		if !valid {
-			panic(p)
+		d.once.Do(func() {
+			defer func() {
+				d.p = recover()
+				if !d.valid {
+					// Re-panic immediately so on the first
+					// call the user gets a complete stack
+					// trace into f.
+					panic(d.p)
+				}
+			}()
+			d.f()
+			d.f = nil      // Do not keep f alive after invoking it.
+			d.valid = true // Set only if f does not panic.
+		})
+		if !d.valid {
+			panic(d.p)
 		}
 	}
 }
@@ -41,29 +44,32 @@ func OnceFunc(f func()) func() {
 //
 // 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 (
+	// Use a struct so that there's a single heap allocation.
+	d := struct {
+		f      func() T
 		once   Once
 		valid  bool
 		p      any
 		result T
-	)
-	g := func() {
-		defer func() {
-			p = recover()
-			if !valid {
-				panic(p)
-			}
-		}()
-		result = f()
-		f = nil
-		valid = true
+	}{
+		f: f,
 	}
 	return func() T {
-		once.Do(g)
-		if !valid {
-			panic(p)
+		d.once.Do(func() {
+			defer func() {
+				d.p = recover()
+				if !d.valid {
+					panic(d.p)
+				}
+			}()
+			d.result = d.f()
+			d.f = nil
+			d.valid = true
+		})
+		if !d.valid {
+			panic(d.p)
 		}
-		return result
+		return d.result
 	}
 }
 
@@ -72,29 +78,32 @@ func OnceValue[T any](f func() T) func() T {
 //
 // 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 (
+	// Use a struct so that there's a single heap allocation.
+	d := struct {
+		f     func() (T1, T2)
 		once  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
+	}{
+		f: f,
 	}
 	return func() (T1, T2) {
-		once.Do(g)
-		if !valid {
-			panic(p)
+		d.once.Do(func() {
+			defer func() {
+				d.p = recover()
+				if !d.valid {
+					panic(d.p)
+				}
+			}()
+			d.r1, d.r2 = d.f()
+			d.f = nil
+			d.valid = true
+		})
+		if !d.valid {
+			panic(d.p)
 		}
-		return r1, r2
+		return d.r1, d.r2
 	}
 }
diff --git a/src/sync/oncefunc_test.go b/src/sync/oncefunc_test.go
index daf094571f..743a816b65 100644
--- a/src/sync/oncefunc_test.go
+++ b/src/sync/oncefunc_test.go
@@ -19,22 +19,30 @@ import (
 
 func TestOnceFunc(t *testing.T) {
 	calls := 0
-	f := sync.OnceFunc(func() { calls++ })
+	of := func() { calls++ }
+	f := sync.OnceFunc(of)
 	allocs := testing.AllocsPerRun(10, f)
 	if calls != 1 {
 		t.Errorf("want calls==1, got %d", calls)
 	}
 	if allocs != 0 {
-		t.Errorf("want 0 allocations per call, got %v", allocs)
+		t.Errorf("want 0 allocations per call to f, got %v", allocs)
+	}
+	allocs = testing.AllocsPerRun(10, func() {
+		f = sync.OnceFunc(of)
+	})
+	if allocs > 2 {
+		t.Errorf("want at most 2 allocations per call to OnceFunc, got %v", allocs)
 	}
 }
 
 func TestOnceValue(t *testing.T) {
 	calls := 0
-	f := sync.OnceValue(func() int {
+	of := func() int {
 		calls++
 		return calls
-	})
+	}
+	f := sync.OnceValue(of)
 	allocs := testing.AllocsPerRun(10, func() { f() })
 	value := f()
 	if calls != 1 {
@@ -44,16 +52,23 @@ func TestOnceValue(t *testing.T) {
 		t.Errorf("want value==1, got %d", value)
 	}
 	if allocs != 0 {
-		t.Errorf("want 0 allocations per call, got %v", allocs)
+		t.Errorf("want 0 allocations per call to f, got %v", allocs)
+	}
+	allocs = testing.AllocsPerRun(10, func() {
+		f = sync.OnceValue(of)
+	})
+	if allocs > 2 {
+		t.Errorf("want at most 2 allocations per call to OnceValue, got %v", allocs)
 	}
 }
 
 func TestOnceValues(t *testing.T) {
 	calls := 0
-	f := sync.OnceValues(func() (int, int) {
+	of := func() (int, int) {
 		calls++
 		return calls, calls + 1
-	})
+	}
+	f := sync.OnceValues(of)
 	allocs := testing.AllocsPerRun(10, func() { f() })
 	v1, v2 := f()
 	if calls != 1 {
@@ -63,7 +78,13 @@ func TestOnceValues(t *testing.T) {
 		t.Errorf("want v1==1 and v2==2, got %d and %d", v1, v2)
 	}
 	if allocs != 0 {
-		t.Errorf("want 0 allocations per call, got %v", allocs)
+		t.Errorf("want 0 allocations per call to f, got %v", allocs)
+	}
+	allocs = testing.AllocsPerRun(10, func() {
+		f = sync.OnceValues(of)
+	})
+	if allocs > 2 {
+		t.Errorf("want at most 2 allocations per call to OnceValues, got %v", allocs)
 	}
 }
 
@@ -234,6 +255,8 @@ var (
 	onceFunc = sync.OnceFunc(func() {})
 
 	onceFuncOnce sync.Once
+
+	onceFuncFunc func()
 )
 
 func doOnceFunc() {
@@ -267,6 +290,12 @@ func BenchmarkOnceFunc(b *testing.B) {
 			f()
 		}
 	})
+	b.Run("v=Make", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			onceFuncFunc = sync.OnceFunc(func() {})
+		}
+	})
 }
 
 var (
@@ -274,6 +303,8 @@ var (
 
 	onceValueOnce  sync.Once
 	onceValueValue int
+
+	onceValueFunc func() int
 )
 
 func doOnceValue() int {
@@ -310,4 +341,80 @@ func BenchmarkOnceValue(b *testing.B) {
 			}
 		}
 	})
+	b.Run("v=Make", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			onceValueFunc = sync.OnceValue(func() int { return 42 })
+		}
+	})
+}
+
+const (
+	onceValuesWant1 = 42
+	onceValuesWant2 = true
+)
+
+var (
+	onceValues = sync.OnceValues(func() (int, bool) {
+		return onceValuesWant1, onceValuesWant2
+	})
+
+	onceValuesOnce   sync.Once
+	onceValuesValue1 int
+	onceValuesValue2 bool
+
+	onceValuesFunc func() (int, bool)
+)
+
+func doOnceValues() (int, bool) {
+	onceValuesOnce.Do(func() {
+		onceValuesValue1 = onceValuesWant1
+		onceValuesValue2 = onceValuesWant2
+	})
+	return onceValuesValue1, onceValuesValue2
+}
+
+func BenchmarkOnceValues(b *testing.B) {
+	// See BenchmarkOnceFunc
+	b.Run("v=Once", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if got1, got2 := doOnceValues(); got1 != onceValuesWant1 {
+				b.Fatalf("value 1: got %d, want %d", got1, onceValuesWant1)
+			} else if got2 != onceValuesWant2 {
+				b.Fatalf("value 2: got %v, want %v", got2, onceValuesWant2)
+			}
+		}
+	})
+	b.Run("v=Global", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if got1, got2 := onceValues(); got1 != onceValuesWant1 {
+				b.Fatalf("value 1: got %d, want %d", got1, onceValuesWant1)
+			} else if got2 != onceValuesWant2 {
+				b.Fatalf("value 2: got %v, want %v", got2, onceValuesWant2)
+			}
+		}
+	})
+	b.Run("v=Local", func(b *testing.B) {
+		b.ReportAllocs()
+		onceValues := sync.OnceValues(func() (int, bool) {
+			return onceValuesWant1, onceValuesWant2
+		})
+		for i := 0; i < b.N; i++ {
+			if got1, got2 := onceValues(); got1 != onceValuesWant1 {
+				b.Fatalf("value 1: got %d, want %d", got1, onceValuesWant1)
+			} else if got2 != onceValuesWant2 {
+				b.Fatalf("value 2: got %v, want %v", got2, onceValuesWant2)
+			}
+		}
+	})
+	b.Run("v=Make", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			onceValuesFunc = sync.OnceValues(func() (int, bool) {
+				return onceValuesWant1, onceValuesWant2
+			})
+		}
+	})
 }