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 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.
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    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.
47func (g *Group) Do(key string, fn func() (any, error)) (v any, 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() (any, 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() (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.
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