diff options
Diffstat (limited to 'src/internal/singleflight/singleflight.go')
-rw-r--r-- | src/internal/singleflight/singleflight.go | 123 |
1 files changed, 123 insertions, 0 deletions
diff --git a/src/internal/singleflight/singleflight.go b/src/internal/singleflight/singleflight.go new file mode 100644 index 0000000..d0e6d2f --- /dev/null +++ b/src/internal/singleflight/singleflight.go @@ -0,0 +1,123 @@ +// 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. +package singleflight + +import "sync" + +// call is an in-flight or completed singleflight.Do call +type call struct { + wg sync.WaitGroup + + // These fields are written once before the WaitGroup is done + // and are only read after the WaitGroup is done. + val any + err error + + // 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 +} + +// Group represents a class of work and forms a namespace in +// which units of work can be executed with duplicate suppression. +type Group struct { + mu sync.Mutex // protects m + m map[string]*call // lazily initialized +} + +// Result holds the results of Do, so they can be passed +// on a channel. +type Result struct { + Val any + 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) Do(key string, fn func() (any, error)) (v any, err error, shared bool) { + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*call) + } + if c, ok := g.m[key]; ok { + c.dups++ + g.mu.Unlock() + c.wg.Wait() + return c.val, c.err, true + } + c := new(call) + 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. +func (g *Group) DoChan(key string, fn func() (any, error)) <-chan Result { + ch := make(chan Result, 1) + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*call) + } + if c, ok := g.m[key]; ok { + c.dups++ + c.chans = append(c.chans, ch) + g.mu.Unlock() + return ch + } + c := &call{chans: []chan<- Result{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) doCall(c *call, key string, fn func() (any, error)) { + c.val, c.err = fn() + + g.mu.Lock() + c.wg.Done() + if g.m[key] == c { + delete(g.m, key) + } + for _, ch := range c.chans { + ch <- Result{c.val, c.err, c.dups > 0} + } + g.mu.Unlock() +} + +// ForgetUnshared tells the singleflight to forget about a key if it is not +// shared with any other goroutines. Future calls to Do for a forgotten key +// will call the function rather than waiting for an earlier call to complete. +// Returns whether the key was forgotten or unknown--that is, whether no +// other goroutines are waiting for the result. +func (g *Group) ForgetUnshared(key string) bool { + g.mu.Lock() + defer g.mu.Unlock() + c, ok := g.m[key] + if !ok { + return true + } + if c.dups == 0 { + delete(g.m, key) + return true + } + return false +} |