1 // Copyright 2013 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 5 // Package singleflight provides a duplicate function call suppression
6 // mechanism.
7 package singleflight
8 9 import "sync"
10 11 // call is an in-flight or completed singleflight.Do call
12 type call struct {
13 wg sync.WaitGroup
14 15 // These fields are written once before the WaitGroup is done
16 // and are only read after the WaitGroup is done.
17 val any
18 err error
19 20 // These fields are read and written with the singleflight
21 // mutex held before the WaitGroup is done, and are read but
22 // not written after the WaitGroup is done.
23 dups int
24 chans []chan<- Result
25 }
26 27 // Group represents a class of work and forms a namespace in
28 // which units of work can be executed with duplicate suppression.
29 type Group struct {
30 mu sync.Mutex // protects m
31 m map[string]*call // lazily initialized
32 }
33 34 // Result holds the results of Do, so they can be passed
35 // on a channel.
36 type Result struct {
37 Val any
38 Err error
39 Shared bool
40 }
41 42 // Do executes and returns the results of the given function, making
43 // sure that only one execution is in-flight for a given key at a
44 // time. If a duplicate comes in, the duplicate caller waits for the
45 // original to complete and receives the same results.
46 // The return value shared indicates whether v was given to multiple callers.
47 func (g *Group) Do(key []byte, fn func() (any, error)) (v any, err error, shared bool) {
48 g.mu.Lock()
49 if g.m == nil {
50 g.m = map[string]*call{}
51 }
52 if c, ok := g.m[key]; ok {
53 c.dups++
54 g.mu.Unlock()
55 c.wg.Wait()
56 return c.val, c.err, true
57 }
58 c := &call{}
59 c.wg.Add(1)
60 g.m[key] = c
61 g.mu.Unlock()
62 63 g.doCall(c, key, fn)
64 return c.val, c.err, c.dups > 0
65 }
66 67 // DoChan is like Do but returns a channel that will receive the
68 // results when they are ready.
69 func (g *Group) DoChan(key []byte, fn func() (any, error)) <-chan Result {
70 ch := chan Result{1}
71 g.mu.Lock()
72 if g.m == nil {
73 g.m = map[string]*call{}
74 }
75 if c, ok := g.m[key]; ok {
76 c.dups++
77 c.chans = append(c.chans, ch)
78 g.mu.Unlock()
79 return ch
80 }
81 c := &call{chans: []chan<- Result{ch}}
82 c.wg.Add(1)
83 g.m[key] = c
84 g.mu.Unlock()
85 86 g.doCall(c, key, fn)
87 88 return ch
89 }
90 91 // doCall handles the single call for a key.
92 func (g *Group) doCall(c *call, key []byte, fn func() (any, error)) {
93 c.val, c.err = fn()
94 95 g.mu.Lock()
96 c.wg.Done()
97 if g.m[key] == c {
98 delete(g.m, key)
99 }
100 for _, ch := range c.chans {
101 ch <- Result{c.val, c.err, c.dups > 0}
102 }
103 g.mu.Unlock()
104 }
105 106 // ForgetUnshared tells the singleflight to forget about a key if it is not
107 // shared with any other goroutines. Future calls to Do for a forgotten key
108 // will call the function rather than waiting for an earlier call to complete.
109 // Returns whether the key was forgotten or unknown--that is, whether no
110 // other goroutines are waiting for the result.
111 func (g *Group) ForgetUnshared(key []byte) bool {
112 g.mu.Lock()
113 defer g.mu.Unlock()
114 c, ok := g.m[key]
115 if !ok {
116 return true
117 }
118 if c.dups == 0 {
119 delete(g.m, key)
120 return true
121 }
122 return false
123 }
124