/
singleflight.go
233 lines (202 loc) · 6.18 KB
/
singleflight.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
// SPDX-License-Identifier: MIT
//
// Copyright 2024 Andrew Bursavich. All rights reserved.
// Use of this source code is governed by The MIT License
// which can be found in the LICENSE file.
// Copyright (c) Tailscale Inc & AUTHOR
// BSD-3-Clause: https://github.com/tailscale/tailscale/blob/v1.32.0/LICENSE
// Copyright 2013 The Go Authors. All rights reserved.
// BSD-3-Clause: https://cs.opensource.google/go/x/sync/+/refs/tags/v0.7.0:LICENSE
// Package singleflight provides a duplicate function call suppression mechanism.
//
// This is a fork of Tailscale's fork of Go's singleflight package which has had several
// homes in the past:
//
// - https://pkg.go.dev/tailscale.com/util/singleflight
// - https://github.com/golang/go/commit/61d3b2db6292581fc07a3767ec23ec94ad6100d1
// - https://github.com/golang/groupcache/tree/master/singleflight
// - https://pkg.go.dev/golang.org/x/sync/singleflight
//
// The tailscale fork adds generics.
//
// This fork adds contexts.
package singleflight
import (
"bytes"
"context"
"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 {
done chan (struct{})
// These fields are written once before the done channel is closed
// and are only read after the done channel is closed.
val V
err error
// These fields are read and written with the singleflight
// mutex held before the done channel is closed, and are read
// but not written after the done channel is closed.
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(ctx context.Context, key K, fn func(context.Context) (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()
select {
case <-c.done:
case <-ctx.Done():
return v, ctx.Err(), false
}
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])
g.m[key] = c
g.mu.Unlock()
g.doCall(ctx, 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(ctx context.Context, key K, fn func(context.Context) (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]{
done: make(chan struct{}),
chans: []chan<- Result[V]{ch},
}
g.m[key] = c
g.mu.Unlock()
go g.doCall(ctx, c, key, fn)
return ch
}
// doCall handles the single call for a key.
func (g *Group[K, V]) doCall(ctx context.Context, c *call[V], key K, fn func(context.Context) (V, error)) {
normalReturn := false
recovered := false
// Use double-defer to distinguish panic from runtime.Goexit,
// For more details see https://golang.org/cl/134395.
defer func() {
// The given function invoked runtime.Goexit.
if !normalReturn && !recovered {
c.err = errGoexit
}
g.mu.Lock()
if g.m[key] == c {
delete(g.m, key)
}
g.mu.Unlock()
close(c.done)
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.
}
panic(e)
}
if c.err == errGoexit {
// Already in the process of goexit, no need to call again.
// This is may leak goroutines waiting on a result from DoChan.
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(ctx)
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()
delete(g.m, key)
g.mu.Unlock()
}