sync: implement OnceFunc, OnceValue, and OnceValues

This adds the three functions from #56102 to the sync package. These
provide a convenient API for the most common uses of sync.Once.

The performance of these is comparable to direct use of sync.Once:

$ go test -run ^$ -bench OnceFunc\|OnceVal -count 20 | benchstat -row .name -col /v
goos: linux
goarch: amd64
pkg: sync
cpu: 11th Gen Intel(R) Core(TM) i7-1185G7 @ 3.00GHz
          │     Once     │                Global                 │                Local                 │
          │    sec/op    │    sec/op     vs base                 │    sec/op     vs base                │
OnceFunc    1.3500n ± 6%   2.7030n ± 1%  +100.22% (p=0.000 n=20)   0.3935n ± 0%  -70.86% (p=0.000 n=20)
OnceValue   1.3155n ± 0%   2.7460n ± 1%  +108.74% (p=0.000 n=20)   0.5478n ± 1%  -58.35% (p=0.000 n=20)

The "Once" column represents the baseline of how code would typically
express these patterns using sync.Once. "Global" binds the closure
returned by OnceFunc/OnceValue to global, which is how I expect these
to be used most of the time. Currently, this defeats some inlining
opportunities, which roughly doubles the cost over sync.Once; however,
it's still *extremely* fast. Finally, "Local" binds the returned
closure to a local variable. This unlocks several levels of inlining
and represents pretty much the best possible case for these APIs, but
is also unlikely to happen in practice. In principle the compiler
could recognize that the global in the "Global" case is initialized in
place and never mutated and do the same optimizations it does in the
"Local" case, but it currently does not.

Fixes #56102

Change-Id: If7355eccd7c8de7288d89a4282ff15ab1469e420
Reviewed-on: https://go-review.googlesource.com/c/go/+/451356
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Austin Clements <austin@google.com>
Reviewed-by: Andrew Gerrand <adg@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Caleb Spare <cespare@gmail.com>
Auto-Submit: Austin Clements <austin@google.com>

2 files changed, 362 insertions, 0 deletions

diff --git a/src/sync/oncefunc.go b/src/sync/oncefunc.go
new file mode 100644
index 0000000000..9ef8344132
--- /dev/null
+++ b/src/sync/oncefunc.go
@@ -0,0 +1,97 @@
+// 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.
+
+package 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  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()
+		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   Once
+		valid  bool
+		p      any
+		result T
+	)
+	g := func() {
+		defer func() {
+			p = recover()
+			if !valid {
+				panic(p)
+			}
+		}()
+		result = f()
+		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  Once
+		valid bool
+		p     any
+		r1    T1
+		r2    T2
+	)
+	g := func() {
+		defer func() {
+			p = recover()
+			if !valid {
+				panic(p)
+			}
+		}()
+		r1, r2 = f()
+		valid = true
+	}
+	return func() (T1, T2) {
+		once.Do(g)
+		if !valid {
+			panic(p)
+		}
+		return r1, r2
+	}
+}
diff --git a/src/sync/oncefunc_test.go b/src/sync/oncefunc_test.go
new file mode 100644
index 0000000000..3c523a5b62
--- /dev/null
+++ b/src/sync/oncefunc_test.go
@@ -0,0 +1,265 @@
+// 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.
+
+package sync_test
+
+import (
+	"bytes"
+	"runtime"
+	"runtime/debug"
+	"sync"
+	"testing"
+)
+
+// We assume that the Once.Do tests have already covered parallelism.
+
+func TestOnceFunc(t *testing.T) {
+	calls := 0
+	f := sync.OnceFunc(func() { calls++ })
+	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)
+	}
+}
+
+func TestOnceValue(t *testing.T) {
+	calls := 0
+	f := sync.OnceValue(func() int {
+		calls++
+		return calls
+	})
+	allocs := testing.AllocsPerRun(10, func() { f() })
+	value := f()
+	if calls != 1 {
+		t.Errorf("want calls==1, got %d", calls)
+	}
+	if value != 1 {
+		t.Errorf("want value==1, got %d", value)
+	}
+	if allocs != 0 {
+		t.Errorf("want 0 allocations per call, got %v", allocs)
+	}
+}
+
+func TestOnceValues(t *testing.T) {
+	calls := 0
+	f := sync.OnceValues(func() (int, int) {
+		calls++
+		return calls, calls + 1
+	})
+	allocs := testing.AllocsPerRun(10, func() { f() })
+	v1, v2 := f()
+	if calls != 1 {
+		t.Errorf("want calls==1, got %d", calls)
+	}
+	if v1 != 1 || v2 != 2 {
+		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)
+	}
+}
+
+func testOncePanicX(t *testing.T, calls *int, f func()) {
+	testOncePanicWith(t, calls, f, func(label string, p any) {
+		if p != "x" {
+			t.Fatalf("%s: want panic %v, got %v", label, "x", p)
+		}
+	})
+}
+
+func testOncePanicWith(t *testing.T, calls *int, f func(), check func(label string, p any)) {
+	// Check that the each call to f panics with the same value, but the
+	// underlying function is only called once.
+	for _, label := range []string{"first time", "second time"} {
+		var p any
+		panicked := true
+		func() {
+			defer func() {
+				p = recover()
+			}()
+			f()
+			panicked = false
+		}()
+		if !panicked {
+			t.Fatalf("%s: f did not panic", label)
+		}
+		check(label, p)
+	}
+	if *calls != 1 {
+		t.Errorf("want calls==1, got %d", *calls)
+	}
+}
+
+func TestOnceFuncPanic(t *testing.T) {
+	calls := 0
+	f := sync.OnceFunc(func() {
+		calls++
+		panic("x")
+	})
+	testOncePanicX(t, &calls, f)
+}
+
+func TestOnceValuePanic(t *testing.T) {
+	calls := 0
+	f := sync.OnceValue(func() int {
+		calls++
+		panic("x")
+	})
+	testOncePanicX(t, &calls, func() { f() })
+}
+
+func TestOnceValuesPanic(t *testing.T) {
+	calls := 0
+	f := sync.OnceValues(func() (int, int) {
+		calls++
+		panic("x")
+	})
+	testOncePanicX(t, &calls, func() { f() })
+}
+
+func TestOnceFuncPanicNil(t *testing.T) {
+	calls := 0
+	f := sync.OnceFunc(func() {
+		calls++
+		panic(nil)
+	})
+	testOncePanicWith(t, &calls, f, func(label string, p any) {
+		switch p.(type) {
+		case nil, *runtime.PanicNilError:
+			return
+		}
+		t.Fatalf("%s: want nil panic, got %v", label, p)
+	})
+}
+
+func TestOnceFuncGoexit(t *testing.T) {
+	// If f calls Goexit, the results are unspecified. But check that f doesn't
+	// get called twice.
+	calls := 0
+	f := sync.OnceFunc(func() {
+		calls++
+		runtime.Goexit()
+	})
+	var wg sync.WaitGroup
+	for i := 0; i < 2; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer func() { recover() }()
+			f()
+		}()
+		wg.Wait()
+	}
+	if calls != 1 {
+		t.Errorf("want calls==1, got %d", calls)
+	}
+}
+
+func TestOnceFuncPanicTraceback(t *testing.T) {
+	// Test that on the first invocation of a OnceFunc, the stack trace goes all
+	// the way to the origin of the panic.
+	f := sync.OnceFunc(onceFuncPanic)
+
+	defer func() {
+		if p := recover(); p != "x" {
+			t.Fatalf("want panic %v, got %v", "x", p)
+		}
+		stack := debug.Stack()
+		want := "sync_test.onceFuncPanic"
+		if !bytes.Contains(stack, []byte(want)) {
+			t.Fatalf("want stack containing %v, got:\n%s", want, string(stack))
+		}
+	}()
+	f()
+}
+
+func onceFuncPanic() {
+	panic("x")
+}
+
+var (
+	onceFunc = sync.OnceFunc(func() {})
+
+	onceFuncOnce sync.Once
+)
+
+func doOnceFunc() {
+	onceFuncOnce.Do(func() {})
+}
+
+func BenchmarkOnceFunc(b *testing.B) {
+	b.Run("v=Once", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			// The baseline is direct use of sync.Once.
+			doOnceFunc()
+		}
+	})
+	b.Run("v=Global", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			// As of 3/2023, the compiler doesn't recognize that onceFunc is
+			// never mutated and is a closure that could be inlined.
+			// Too bad, because this is how OnceFunc will usually be used.
+			onceFunc()
+		}
+	})
+	b.Run("v=Local", func(b *testing.B) {
+		b.ReportAllocs()
+		// As of 3/2023, the compiler *does* recognize this local binding as an
+		// inlinable closure. This is the best case for OnceFunc, but probably
+		// not typical usage.
+		f := sync.OnceFunc(func() {})
+		for i := 0; i < b.N; i++ {
+			f()
+		}
+	})
+}
+
+var (
+	onceValue = sync.OnceValue(func() int { return 42 })
+
+	onceValueOnce  sync.Once
+	onceValueValue int
+)
+
+func doOnceValue() int {
+	onceValueOnce.Do(func() {
+		onceValueValue = 42
+	})
+	return onceValueValue
+}
+
+func BenchmarkOnceValue(b *testing.B) {
+	// See BenchmarkOnceFunc
+	b.Run("v=Once", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if want, got := 42, doOnceValue(); want != got {
+				b.Fatalf("want %d, got %d", want, got)
+			}
+		}
+	})
+	b.Run("v=Global", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			if want, got := 42, onceValue(); want != got {
+				b.Fatalf("want %d, got %d", want, got)
+			}
+		}
+	})
+	b.Run("v=Local", func(b *testing.B) {
+		b.ReportAllocs()
+		onceValue := sync.OnceValue(func() int { return 42 })
+		for i := 0; i < b.N; i++ {
+			if want, got := 42, onceValue(); want != got {
+				b.Fatalf("want %d, got %d", want, got)
+			}
+		}
+	})
+}