github.com/aavshr/aws-sdk-go@v1.41.3/internal/sync/singleflight/singleflight.go (about)

     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 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.
    33  type 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.
    40  type 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.
    51  func (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.
    73  func (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.
    96  func (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.
   113  func (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  }