github.com/GitbookIO/syncgroup@v0.0.0-20200915204659-4f0b2961ab10/active.go (about)

     1  package syncgroup
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  )
     7  
     8  // ActiveGroup is like a smart grouped version of sync.WaitGroup
     9  type ActiveGroup struct {
    10  	groups map[string]*ActiveCounter
    11  	lock   sync.RWMutex
    12  }
    13  
    14  func NewActiveGroup() *ActiveGroup {
    15  	return &ActiveGroup{
    16  		groups: map[string]*ActiveCounter{},
    17  		lock:   sync.RWMutex{},
    18  	}
    19  }
    20  
    21  func (ag *ActiveGroup) IsActive(key string) bool {
    22  	// False for inexistant key
    23  	if !ag.Has(key) {
    24  		return false
    25  	}
    26  	return ag.get(key).IsActive()
    27  }
    28  
    29  func (ag *ActiveGroup) WaitUntilFree(key string) {
    30  	ag.get(key).WaitUntilFree()
    31  }
    32  
    33  func (ag *ActiveGroup) Inc(key string) int64 {
    34  	return ag.get(key).Inc()
    35  }
    36  
    37  func (ag *ActiveGroup) Dec(key string) int64 {
    38  	// Get counter
    39  	ac := ag.get(key)
    40  	// Decrement counter
    41  	count := ac.Dec()
    42  	if count == 0 {
    43  		return ag.del(key)
    44  	}
    45  	return count
    46  }
    47  
    48  func (ag *ActiveGroup) get(key string) *ActiveCounter {
    49  	ag.lock.Lock()
    50  	defer ag.lock.Unlock()
    51  
    52  	// Create if doesn't exist
    53  	if !ag.has(key) {
    54  		ag.groups[key] = &ActiveCounter{}
    55  	}
    56  
    57  	return ag.groups[key]
    58  }
    59  
    60  func (ag *ActiveGroup) Has(key string) bool {
    61  	ag.lock.RLock()
    62  	defer ag.lock.RUnlock()
    63  	return ag.has(key)
    64  }
    65  
    66  func (ag *ActiveGroup) has(key string) bool {
    67  	_, ok := ag.groups[key]
    68  	return ok
    69  }
    70  
    71  func (ag *ActiveGroup) del(key string) int64 {
    72  	ag.lock.Lock()
    73  	defer ag.lock.Unlock()
    74  
    75  	if counter, ok := ag.groups[key]; ok && counter.cnt == 0 {
    76  		delete(ag.groups, key)
    77  	} else if ok {
    78  		// Ok but not zero counter
    79  		return counter.cnt
    80  	}
    81  
    82  	// Already deleted
    83  	return 0
    84  }
    85  
    86  type ActiveCounter struct {
    87  	cnt int64
    88  	grp sync.WaitGroup
    89  }
    90  
    91  func (ac *ActiveCounter) Inc() int64 {
    92  	ac.grp.Add(1)
    93  	return atomic.AddInt64(&ac.cnt, 1)
    94  }
    95  
    96  func (ac *ActiveCounter) Dec() int64 {
    97  	count := atomic.AddInt64(&ac.cnt, -1)
    98  	ac.grp.Done()
    99  	return count
   100  }
   101  
   102  func (ac *ActiveCounter) IsActive() bool {
   103  	return atomic.LoadInt64(&ac.cnt) > 0
   104  }
   105  
   106  func (ac *ActiveCounter) WaitUntilFree() {
   107  	ac.grp.Wait()
   108  }