util/singleflight: add fork of singleflight with generics

Forked from golang.org/x/sync/singleflight at
the x/sync repo's commit 67f06af15bc961c363a7260195bcd53487529a21

Updates golang/go#53427

Change-Id: Iec2b47b7777940017bb9b3db9bd7d93ba4a2e394
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

util/singleflight/singleflight.go

@@ -0,0 +1,225 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Copyright 2013 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 singleflight provides a duplicate function call suppression
+// mechanism.
+//
+// This is a Tailscale fork of Go's singleflight package which has had several
+// homes in the past:
+//
+//    * https://github.com/golang/go/commit/61d3b2db6292581fc07a3767ec23ec94ad6100d1
+//    * https://github.com/golang/groupcache/tree/master/singleflight
+//    * https://pkg.go.dev/golang.org/x/sync/singleflight
+//
+// This fork adds generics.
+package singleflight // import "tailscale.com/util/singleflight"
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"runtime"
+	"runtime/debug"
+	"sync"
+)
+
+// errGoexit indicates the runtime.Goexit was called in
+// the user given function.
+var errGoexit = errors.New("runtime.Goexit was called")
+
+// A panicError is an arbitrary value recovered from a panic
+// with the stack trace during the execution of given function.
+type panicError struct {
+	value interface{}
+	stack []byte
+}
+
+// Error implements error interface.
+func (p *panicError) Error() string {
+	return fmt.Sprintf("%v\n\n%s", p.value, p.stack)
+}
+
+func newPanicError(v interface{}) error {
+	stack := debug.Stack()
+
+	// The first line of the stack trace is of the form "goroutine N [status]:"
+	// but by the time the panic reaches Do the goroutine may no longer exist
+	// and its status will have changed. Trim out the misleading line.
+	if line := bytes.IndexByte(stack[:], '\n'); line >= 0 {
+		stack = stack[line+1:]
+	}
+	return &panicError{value: v, stack: stack}
+}
+
+// call is an in-flight or completed singleflight.Do call
+type call[V any] struct {
+	wg sync.WaitGroup
+
+	// These fields are written once before the WaitGroup is done
+	// and are only read after the WaitGroup is done.
+	val V
+	err error
+
+	// forgotten indicates whether Forget was called with this call's key
+	// while the call was still in flight.
+	forgotten bool
+
+	// These fields are read and written with the singleflight
+	// mutex held before the WaitGroup is done, and are read but
+	// not written after the WaitGroup is done.
+	dups  int
+	chans []chan<- Result[V]
+}
+
+// Group represents a class of work and forms a namespace in
+// which units of work can be executed with duplicate suppression.
+type Group[K comparable, V any] struct {
+	mu sync.Mutex     // protects m
+	m  map[K]*call[V] // lazily initialized
+}
+
+// Result holds the results of Do, so they can be passed
+// on a channel.
+type Result[V any] struct {
+	Val    V
+	Err    error
+	Shared bool
+}
+
+// Do executes and returns the results of the given function, making
+// sure that only one execution is in-flight for a given key at a
+// time. If a duplicate comes in, the duplicate caller waits for the
+// original to complete and receives the same results.
+// The return value shared indicates whether v was given to multiple callers.
+func (g *Group[K, V]) Do(key K, fn func() (V, error)) (v V, err error, shared bool) {
+	g.mu.Lock()
+	if g.m == nil {
+		g.m = make(map[K]*call[V])
+	}
+	if c, ok := g.m[key]; ok {
+		c.dups++
+		g.mu.Unlock()
+		c.wg.Wait()
+
+		if e, ok := c.err.(*panicError); ok {
+			panic(e)
+		} else if c.err == errGoexit {
+			runtime.Goexit()
+		}
+		return c.val, c.err, true
+	}
+	c := new(call[V])
+	c.wg.Add(1)
+	g.m[key] = c
+	g.mu.Unlock()
+
+	g.doCall(c, key, fn)
+	return c.val, c.err, c.dups > 0
+}
+
+// DoChan is like Do but returns a channel that will receive the
+// results when they are ready.
+//
+// The returned channel will not be closed.
+func (g *Group[K, V]) DoChan(key K, fn func() (V, error)) <-chan Result[V] {
+	ch := make(chan Result[V], 1)
+	g.mu.Lock()
+	if g.m == nil {
+		g.m = make(map[K]*call[V])
+	}
+	if c, ok := g.m[key]; ok {
+		c.dups++
+		c.chans = append(c.chans, ch)
+		g.mu.Unlock()
+		return ch
+	}
+	c := &call[V]{chans: []chan<- Result[V]{ch}}
+	c.wg.Add(1)
+	g.m[key] = c
+	g.mu.Unlock()
+
+	go g.doCall(c, key, fn)
+
+	return ch
+}
+
+// doCall handles the single call for a key.
+func (g *Group[K, V]) doCall(c *call[V], key K, fn func() (V, error)) {
+	normalReturn := false
+	recovered := false
+
+	// use double-defer to distinguish panic from runtime.Goexit,
+	// more details see https://golang.org/cl/134395
+	defer func() {
+		// the given function invoked runtime.Goexit
+		if !normalReturn && !recovered {
+			c.err = errGoexit
+		}
+
+		c.wg.Done()
+		g.mu.Lock()
+		defer g.mu.Unlock()
+		if !c.forgotten {
+			delete(g.m, key)
+		}
+
+		if e, ok := c.err.(*panicError); ok {
+			// In order to prevent the waiting channels from being blocked forever,
+			// needs to ensure that this panic cannot be recovered.
+			if len(c.chans) > 0 {
+				go panic(e)
+				select {} // Keep this goroutine around so that it will appear in the crash dump.
+			} else {
+				panic(e)
+			}
+		} else if c.err == errGoexit {
+			// Already in the process of goexit, no need to call again
+		} else {
+			// Normal return
+			for _, ch := range c.chans {
+				ch <- Result[V]{c.val, c.err, c.dups > 0}
+			}
+		}
+	}()
+
+	func() {
+		defer func() {
+			if !normalReturn {
+				// Ideally, we would wait to take a stack trace until we've determined
+				// whether this is a panic or a runtime.Goexit.
+				//
+				// Unfortunately, the only way we can distinguish the two is to see
+				// whether the recover stopped the goroutine from terminating, and by
+				// the time we know that, the part of the stack trace relevant to the
+				// panic has been discarded.
+				if r := recover(); r != nil {
+					c.err = newPanicError(r)
+				}
+			}
+		}()
+
+		c.val, c.err = fn()
+		normalReturn = true
+	}()
+
+	if !normalReturn {
+		recovered = true
+	}
+}
+
+// Forget tells the singleflight to forget about a key.  Future calls
+// to Do for this key will call the function rather than waiting for
+// an earlier call to complete.
+func (g *Group[K, V]) Forget(key K) {
+	g.mu.Lock()
+	if c, ok := g.m[key]; ok {
+		c.forgotten = true
+	}
+	delete(g.m, key)
+	g.mu.Unlock()
+}

util/singleflight/singleflight_test.go

@@ -0,0 +1,324 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Copyright 2013 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 singleflight
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"runtime"
+	"runtime/debug"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestDo(t *testing.T) {
+	var g Group[string, any]
+	v, err, _ := g.Do("key", func() (interface{}, error) {
+		return "bar", nil
+	})
+	if got, want := fmt.Sprintf("%v (%T)", v, v), "bar (string)"; got != want {
+		t.Errorf("Do = %v; want %v", got, want)
+	}
+	if err != nil {
+		t.Errorf("Do error = %v", err)
+	}
+}
+
+func TestDoErr(t *testing.T) {
+	var g Group[string, any]
+	someErr := errors.New("Some error")
+	v, err, _ := g.Do("key", func() (interface{}, error) {
+		return nil, someErr
+	})
+	if err != someErr {
+		t.Errorf("Do error = %v; want someErr %v", err, someErr)
+	}
+	if v != nil {
+		t.Errorf("unexpected non-nil value %#v", v)
+	}
+}
+
+func TestDoDupSuppress(t *testing.T) {
+	var g Group[string, any]
+	var wg1, wg2 sync.WaitGroup
+	c := make(chan string, 1)
+	var calls int32
+	fn := func() (interface{}, error) {
+		if atomic.AddInt32(&calls, 1) == 1 {
+			// First invocation.
+			wg1.Done()
+		}
+		v := <-c
+		c <- v // pump; make available for any future calls
+
+		time.Sleep(10 * time.Millisecond) // let more goroutines enter Do
+
+		return v, nil
+	}
+
+	const n = 10
+	wg1.Add(1)
+	for i := 0; i < n; i++ {
+		wg1.Add(1)
+		wg2.Add(1)
+		go func() {
+			defer wg2.Done()
+			wg1.Done()
+			v, err, _ := g.Do("key", fn)
+			if err != nil {
+				t.Errorf("Do error: %v", err)
+				return
+			}
+			if s, _ := v.(string); s != "bar" {
+				t.Errorf("Do = %T %v; want %q", v, v, "bar")
+			}
+		}()
+	}
+	wg1.Wait()
+	// At least one goroutine is in fn now and all of them have at
+	// least reached the line before the Do.
+	c <- "bar"
+	wg2.Wait()
+	if got := atomic.LoadInt32(&calls); got <= 0 || got >= n {
+		t.Errorf("number of calls = %d; want over 0 and less than %d", got, n)
+	}
+}
+
+// Test that singleflight behaves correctly after Forget called.
+// See https://github.com/golang/go/issues/31420
+func TestForget(t *testing.T) {
+	var g Group[string, any]
+
+	var (
+		firstStarted  = make(chan struct{})
+		unblockFirst  = make(chan struct{})
+		firstFinished = make(chan struct{})
+	)
+
+	go func() {
+		g.Do("key", func() (i interface{}, e error) {
+			close(firstStarted)
+			<-unblockFirst
+			close(firstFinished)
+			return
+		})
+	}()
+	<-firstStarted
+	g.Forget("key")
+
+	unblockSecond := make(chan struct{})
+	secondResult := g.DoChan("key", func() (i interface{}, e error) {
+		<-unblockSecond
+		return 2, nil
+	})
+
+	close(unblockFirst)
+	<-firstFinished
+
+	thirdResult := g.DoChan("key", func() (i interface{}, e error) {
+		return 3, nil
+	})
+
+	close(unblockSecond)
+	<-secondResult
+	r := <-thirdResult
+	if r.Val != 2 {
+		t.Errorf("We should receive result produced by second call, expected: 2, got %d", r.Val)
+	}
+}
+
+func TestDoChan(t *testing.T) {
+	var g Group[string, any]
+	ch := g.DoChan("key", func() (interface{}, error) {
+		return "bar", nil
+	})
+
+	res := <-ch
+	v := res.Val
+	err := res.Err
+	if got, want := fmt.Sprintf("%v (%T)", v, v), "bar (string)"; got != want {
+		t.Errorf("Do = %v; want %v", got, want)
+	}
+	if err != nil {
+		t.Errorf("Do error = %v", err)
+	}
+}
+
+// Test singleflight behaves correctly after Do panic.
+// See https://github.com/golang/go/issues/41133
+func TestPanicDo(t *testing.T) {
+	var g Group[string, any]
+	fn := func() (interface{}, error) {
+		panic("invalid memory address or nil pointer dereference")
+	}
+
+	const n = 5
+	waited := int32(n)
+	panicCount := int32(0)
+	done := make(chan struct{})
+	for i := 0; i < n; i++ {
+		go func() {
+			defer func() {
+				if err := recover(); err != nil {
+					t.Logf("Got panic: %v\n%s", err, debug.Stack())
+					atomic.AddInt32(&panicCount, 1)
+				}
+
+				if atomic.AddInt32(&waited, -1) == 0 {
+					close(done)
+				}
+			}()
+
+			g.Do("key", fn)
+		}()
+	}
+
+	select {
+	case <-done:
+		if panicCount != n {
+			t.Errorf("Expect %d panic, but got %d", n, panicCount)
+		}
+	case <-time.After(time.Second):
+		t.Fatalf("Do hangs")
+	}
+}
+
+func TestGoexitDo(t *testing.T) {
+	var g Group[string, any]
+	fn := func() (interface{}, error) {
+		runtime.Goexit()
+		return nil, nil
+	}
+
+	const n = 5
+	waited := int32(n)
+	done := make(chan struct{})
+	for i := 0; i < n; i++ {
+		go func() {
+			var err error
+			defer func() {
+				if err != nil {
+					t.Errorf("Error should be nil, but got: %v", err)
+				}
+				if atomic.AddInt32(&waited, -1) == 0 {
+					close(done)
+				}
+			}()
+			_, err, _ = g.Do("key", fn)
+		}()
+	}
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatalf("Do hangs")
+	}
+}
+
+func TestPanicDoChan(t *testing.T) {
+	if runtime.GOOS == "js" {
+		t.Skipf("js does not support exec")
+	}
+
+	if os.Getenv("TEST_PANIC_DOCHAN") != "" {
+		defer func() {
+			recover()
+		}()
+
+		g := new(Group[string, any])
+		ch := g.DoChan("", func() (interface{}, error) {
+			panic("Panicking in DoChan")
+		})
+		<-ch
+		t.Fatalf("DoChan unexpectedly returned")
+	}
+
+	t.Parallel()
+
+	cmd := exec.Command(os.Args[0], "-test.run="+t.Name(), "-test.v")
+	cmd.Env = append(os.Environ(), "TEST_PANIC_DOCHAN=1")
+	out := new(bytes.Buffer)
+	cmd.Stdout = out
+	cmd.Stderr = out
+	if err := cmd.Start(); err != nil {
+		t.Fatal(err)
+	}
+
+	err := cmd.Wait()
+	t.Logf("%s:\n%s", strings.Join(cmd.Args, " "), out)
+	if err == nil {
+		t.Errorf("Test subprocess passed; want a crash due to panic in DoChan")
+	}
+	if bytes.Contains(out.Bytes(), []byte("DoChan unexpectedly")) {
+		t.Errorf("Test subprocess failed with an unexpected failure mode.")
+	}
+	if !bytes.Contains(out.Bytes(), []byte("Panicking in DoChan")) {
+		t.Errorf("Test subprocess failed, but the crash isn't caused by panicking in DoChan")
+	}
+}
+
+func TestPanicDoSharedByDoChan(t *testing.T) {
+	if runtime.GOOS == "js" {
+		t.Skipf("js does not support exec")
+	}
+
+	if os.Getenv("TEST_PANIC_DOCHAN") != "" {
+		blocked := make(chan struct{})
+		unblock := make(chan struct{})
+
+		g := new(Group[string, any])
+		go func() {
+			defer func() {
+				recover()
+			}()
+			g.Do("", func() (interface{}, error) {
+				close(blocked)
+				<-unblock
+				panic("Panicking in Do")
+			})
+		}()
+
+		<-blocked
+		ch := g.DoChan("", func() (interface{}, error) {
+			panic("DoChan unexpectedly executed callback")
+		})
+		close(unblock)
+		<-ch
+		t.Fatalf("DoChan unexpectedly returned")
+	}
+
+	t.Parallel()
+
+	cmd := exec.Command(os.Args[0], "-test.run="+t.Name(), "-test.v")
+	cmd.Env = append(os.Environ(), "TEST_PANIC_DOCHAN=1")
+	out := new(bytes.Buffer)
+	cmd.Stdout = out
+	cmd.Stderr = out
+	if err := cmd.Start(); err != nil {
+		t.Fatal(err)
+	}
+
+	err := cmd.Wait()
+	t.Logf("%s:\n%s", strings.Join(cmd.Args, " "), out)
+	if err == nil {
+		t.Errorf("Test subprocess passed; want a crash due to panic in Do shared by DoChan")
+	}
+	if bytes.Contains(out.Bytes(), []byte("DoChan unexpectedly")) {
+		t.Errorf("Test subprocess failed with an unexpected failure mode.")
+	}
+	if !bytes.Contains(out.Bytes(), []byte("Panicking in Do")) {
+		t.Errorf("Test subprocess failed, but the crash isn't caused by panicking in Do")
+	}
+}