fix: race in Single.Do

This commit is contained in:
wwqgtxx 2025-04-04 10:55:16 +08:00
parent e81f3a97af
commit eaaccffcef

View File

@ -13,25 +13,26 @@ type call[T any] struct {
type Single[T any] struct { type Single[T any] struct {
mux sync.Mutex mux sync.Mutex
last time.Time
wait time.Duration wait time.Duration
call *call[T] call *call[T]
result *Result[T] result *Result[T]
} }
type Result[T any] struct { type Result[T any] struct {
Val T Val T
Err error Err error
Time time.Time
} }
// Do single.Do likes sync.singleFlight // Do single.Do likes sync.singleFlight
func (s *Single[T]) Do(fn func() (T, error)) (v T, err error, shared bool) { func (s *Single[T]) Do(fn func() (T, error)) (v T, err error, shared bool) {
s.mux.Lock() s.mux.Lock()
now := time.Now() result := s.result
if now.Before(s.last.Add(s.wait)) { if result != nil && time.Since(result.Time) < s.wait {
s.mux.Unlock() s.mux.Unlock()
return s.result.Val, s.result.Err, true return result.Val, result.Err, true
} }
s.result = nil // The result has expired, clear it
if callM := s.call; callM != nil { if callM := s.call; callM != nil {
s.mux.Unlock() s.mux.Unlock()
@ -47,15 +48,19 @@ func (s *Single[T]) Do(fn func() (T, error)) (v T, err error, shared bool) {
callM.wg.Done() callM.wg.Done()
s.mux.Lock() s.mux.Lock()
s.call = nil if s.call == callM { // maybe reset when fn is running
s.result = &Result[T]{callM.val, callM.err} s.call = nil
s.last = now s.result = &Result[T]{callM.val, callM.err, time.Now()}
}
s.mux.Unlock() s.mux.Unlock()
return callM.val, callM.err, false return callM.val, callM.err, false
} }
func (s *Single[T]) Reset() { func (s *Single[T]) Reset() {
s.last = time.Time{} s.mux.Lock()
s.call = nil
s.result = nil
s.mux.Unlock()
} }
func NewSingle[T any](wait time.Duration) *Single[T] { func NewSingle[T any](wait time.Duration) *Single[T] {