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.
7package singleflight
8
9import "sync"
10
11// call is an in-flight or completed singleflight.Do call
12type 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 interface{}
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.
29type 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.
36type Result struct {
37	Val    interface{}
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.
47func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
48	g.mu.Lock()
49	if g.m == nil {
50		g.m = make(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 := new(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. The second result is true if the function
69// will eventually be called, false if it will not (because there is
70// a pending request with this key).
71func (g *Group) DoChan(key string, fn func() (interface{}, error)) (<-chan Result, bool) {
72	ch := make(chan Result, 1)
73	g.mu.Lock()
74	if g.m == nil {
75		g.m = make(map[string]*call)
76	}
77	if c, ok := g.m[key]; ok {
78		c.dups++
79		c.chans = append(c.chans, ch)
80		g.mu.Unlock()
81		return ch, false
82	}
83	c := &call{chans: []chan<- Result{ch}}
84	c.wg.Add(1)
85	g.m[key] = c
86	g.mu.Unlock()
87
88	go g.doCall(c, key, fn)
89
90	return ch, true
91}
92
93// doCall handles the single call for a key.
94func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
95	c.val, c.err = fn()
96	c.wg.Done()
97
98	g.mu.Lock()
99	delete(g.m, key)
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.
111func (g *Group) ForgetUnshared(key string) 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