github.com/lingyao2333/mo-zero@v1.4.1/core/syncx/singleflight.go (about)

     1  package syncx
     2  
     3  import "sync"
     4  
     5  type (
     6  	// SingleFlight lets the concurrent calls with the same key to share the call result.
     7  	// For example, A called F, before it's done, B called F. Then B would not execute F,
     8  	// and shared the result returned by F which called by A.
     9  	// The calls with the same key are dependent, concurrent calls share the returned values.
    10  	// A ------->calls F with key<------------------->returns val
    11  	// B --------------------->calls F with key------>returns val
    12  	SingleFlight interface {
    13  		Do(key string, fn func() (interface{}, error)) (interface{}, error)
    14  		DoEx(key string, fn func() (interface{}, error)) (interface{}, bool, error)
    15  	}
    16  
    17  	call struct {
    18  		wg  sync.WaitGroup
    19  		val interface{}
    20  		err error
    21  	}
    22  
    23  	flightGroup struct {
    24  		calls map[string]*call
    25  		lock  sync.Mutex
    26  	}
    27  )
    28  
    29  // NewSingleFlight returns a SingleFlight.
    30  func NewSingleFlight() SingleFlight {
    31  	return &flightGroup{
    32  		calls: make(map[string]*call),
    33  	}
    34  }
    35  
    36  func (g *flightGroup) Do(key string, fn func() (interface{}, error)) (interface{}, error) {
    37  	c, done := g.createCall(key)
    38  	if done {
    39  		return c.val, c.err
    40  	}
    41  
    42  	g.makeCall(c, key, fn)
    43  	return c.val, c.err
    44  }
    45  
    46  func (g *flightGroup) DoEx(key string, fn func() (interface{}, error)) (val interface{}, fresh bool, err error) {
    47  	c, done := g.createCall(key)
    48  	if done {
    49  		return c.val, false, c.err
    50  	}
    51  
    52  	g.makeCall(c, key, fn)
    53  	return c.val, true, c.err
    54  }
    55  
    56  func (g *flightGroup) createCall(key string) (c *call, done bool) {
    57  	g.lock.Lock()
    58  	if c, ok := g.calls[key]; ok {
    59  		g.lock.Unlock()
    60  		c.wg.Wait()
    61  		return c, true
    62  	}
    63  
    64  	c = new(call)
    65  	c.wg.Add(1)
    66  	g.calls[key] = c
    67  	g.lock.Unlock()
    68  
    69  	return c, false
    70  }
    71  
    72  func (g *flightGroup) makeCall(c *call, key string, fn func() (interface{}, error)) {
    73  	defer func() {
    74  		g.lock.Lock()
    75  		delete(g.calls, key)
    76  		g.lock.Unlock()
    77  		c.wg.Done()
    78  	}()
    79  
    80  	c.val, c.err = fn()
    81  }