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	// forgotten indicates whether Forget was called with this call's key
21	// while the call was still in flight.
22	forgotten bool
23
24	// These fields are read and written with the singleflight
25	// mutex held before the WaitGroup is done, and are read but
26	// not written after the WaitGroup is done.
27	dups  int
28	chans []chan<- Result
29}
30
31// Group represents a class of work and forms a namespace in
32// which units of work can be executed with duplicate suppression.
33type Group struct {
34	mu sync.Mutex       // protects m
35	m  map[string]*call // lazily initialized
36}
37
38// Result holds the results of Do, so they can be passed
39// on a channel.
40type Result struct {
41	Val    interface{}
42	Err    error
43	Shared bool
44}
45
46// Do executes and returns the results of the given function, making
47// sure that only one execution is in-flight for a given key at a
48// time. If a duplicate comes in, the duplicate caller waits for the
49// original to complete and receives the same results.
50// The return value shared indicates whether v was given to multiple callers.
51func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
52	g.mu.Lock()
53	if g.m == nil {
54		g.m = make(map[string]*call)
55	}
56	if c, ok := g.m[key]; ok {
57		c.dups++
58		g.mu.Unlock()
59		c.wg.Wait()
60		return c.val, c.err, true
61	}
62	c := new(call)
63	c.wg.Add(1)
64	g.m[key] = c
65	g.mu.Unlock()
66
67	g.doCall(c, key, fn)
68	return c.val, c.err, c.dups > 0
69}
70
71// DoChan is like Do but returns a channel that will receive the
72// results when they are ready.
73func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
74	ch := make(chan Result, 1)
75	g.mu.Lock()
76	if g.m == nil {
77		g.m = make(map[string]*call)
78	}
79	if c, ok := g.m[key]; ok {
80		c.dups++
81		c.chans = append(c.chans, ch)
82		g.mu.Unlock()
83		return ch
84	}
85	c := &call{chans: []chan<- Result{ch}}
86	c.wg.Add(1)
87	g.m[key] = c
88	g.mu.Unlock()
89
90	go g.doCall(c, key, fn)
91
92	return ch
93}
94
95// doCall handles the single call for a key.
96func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
97	c.val, c.err = fn()
98	c.wg.Done()
99
100	g.mu.Lock()
101	if !c.forgotten {
102		delete(g.m, key)
103	}
104	for _, ch := range c.chans {
105		ch <- Result{c.val, c.err, c.dups > 0}
106	}
107	g.mu.Unlock()
108}
109
110// Forget tells the singleflight to forget about a key.  Future calls
111// to Do for this key will call the function rather than waiting for
112// an earlier call to complete.
113func (g *Group) Forget(key string) {
114	g.mu.Lock()
115	if c, ok := g.m[key]; ok {
116		c.forgotten = true
117	}
118	delete(g.m, key)
119	g.mu.Unlock()
120}
121