singleflight.go

     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 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. 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).
    71  func (g *Group) DoChan(key string, fn func() (any, 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.
    94  func (g *Group) doCall(c *call, key string, fn func() (any, 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.
   111  func (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
View as plain text