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.
69func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
70	ch := make(chan Result, 1)
71	g.mu.Lock()
72	if g.m == nil {
73		g.m = make(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	go g.doCall(c, key, fn)
87
88	return ch
89}
90
91// doCall handles the single call for a key.
92func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
93	c.val, c.err = fn()
94	c.wg.Done()
95
96	g.mu.Lock()
97	delete(g.m, key)
98	for _, ch := range c.chans {
99		ch <- Result{c.val, c.err, c.dups > 0}
100	}
101	g.mu.Unlock()
102}
103
104// Forget tells the singleflight to forget about a key.  Future calls
105// to Do for this key will call the function rather than waiting for
106// an earlier call to complete.
107func (g *Group) Forget(key string) {
108	g.mu.Lock()
109	delete(g.m, key)
110	g.mu.Unlock()
111}
112