mirror of https://github.com/tailscale/tailscale/
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>pull/4877/head
Brad Fitzpatrick
2 years ago
committed by
Brad Fitzpatrick
parent
757ecf7e80
commit
59ed846277
@ -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()
|
||||
}
|
||||
@ -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")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue