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

     1  package syncx
     2  
     3  import "sync"
     4  
     5  type (
     6  	// LockedCalls makes sure the calls with the same key to be called sequentially.
     7  	// For example, A called F, before it's done, B called F, then B's call would not blocked,
     8  	// after A's call finished, B's call got executed.
     9  	// The calls with the same key are independent, not sharing the returned values.
    10  	// A ------->calls F with key and executes<------->returns
    11  	// B ------------------>calls F with key<--------->executes<---->returns
    12  	LockedCalls interface {
    13  		Do(key string, fn func() (interface{}, error)) (interface{}, error)
    14  	}
    15  
    16  	lockedGroup struct {
    17  		mu sync.Mutex
    18  		m  map[string]*sync.WaitGroup
    19  	}
    20  )
    21  
    22  // NewLockedCalls returns a LockedCalls.
    23  func NewLockedCalls() LockedCalls {
    24  	return &lockedGroup{
    25  		m: make(map[string]*sync.WaitGroup),
    26  	}
    27  }
    28  
    29  func (lg *lockedGroup) Do(key string, fn func() (interface{}, error)) (interface{}, error) {
    30  begin:
    31  	lg.mu.Lock()
    32  	if wg, ok := lg.m[key]; ok {
    33  		lg.mu.Unlock()
    34  		wg.Wait()
    35  		goto begin
    36  	}
    37  
    38  	return lg.makeCall(key, fn)
    39  }
    40  
    41  func (lg *lockedGroup) makeCall(key string, fn func() (interface{}, error)) (interface{}, error) {
    42  	var wg sync.WaitGroup
    43  	wg.Add(1)
    44  	lg.m[key] = &wg
    45  	lg.mu.Unlock()
    46  
    47  	defer func() {
    48  		// delete key first, done later. can't reverse the order, because if reverse,
    49  		// another Do call might wg.Wait() without get notified with wg.Done()
    50  		lg.mu.Lock()
    51  		delete(lg.m, key)
    52  		lg.mu.Unlock()
    53  		wg.Done()
    54  	}()
    55  
    56  	return fn()
    57  }