github.com/amazechain/amc@v0.1.3/utils/lock.go (about)

     1  package utils
     2  
     3  import (
     4  	"runtime"
     5  	"sort"
     6  )
     7  
     8  var locks = struct {
     9  	lock chan byte
    10  	list map[string]chan byte
    11  }{
    12  	lock: make(chan byte, 1),
    13  	list: make(map[string]chan byte),
    14  }
    15  
    16  type Lock struct {
    17  	keys   []string
    18  	chans  []chan byte
    19  	lock   chan byte
    20  	unlock chan byte
    21  }
    22  
    23  func (lk *Lock) Lock() {
    24  	lk.lock <- 1
    25  
    26  	// get the channels and attempt to acquire them
    27  	lk.chans = make([]chan byte, 0, len(lk.keys))
    28  	for i := 0; i < len(lk.keys); {
    29  		ch := getChan(lk.keys[i])
    30  		_, ok := <-ch
    31  		if ok {
    32  			lk.chans = append(lk.chans, ch)
    33  			i++
    34  		}
    35  	}
    36  
    37  	lk.unlock <- 1
    38  }
    39  
    40  // Unlock unlocks this lock. Must be called after Lock.
    41  // Can only be invoked if there is a previous call to Lock.
    42  func (lk *Lock) Unlock() {
    43  	<-lk.unlock
    44  
    45  	if lk.chans != nil {
    46  		for _, ch := range lk.chans {
    47  			ch <- 1
    48  		}
    49  		lk.chans = nil
    50  	}
    51  	// Clean unused channels after the unlock.
    52  	Clean()
    53  	<-lk.lock
    54  }
    55  
    56  // Yield temporarily unlocks, gives up the cpu time to other goroutine, and attempts to lock again.
    57  func (lk *Lock) Yield() {
    58  	lk.Unlock()
    59  	runtime.Gosched()
    60  	lk.Lock()
    61  }
    62  
    63  // NewMultilock creates a new multilock for the specified keys
    64  func NewMultilock(locks ...string) *Lock {
    65  	if len(locks) == 0 {
    66  		return nil
    67  	}
    68  
    69  	locks = unique(locks)
    70  	sort.Strings(locks)
    71  	return &Lock{
    72  		keys:   locks,
    73  		lock:   make(chan byte, 1),
    74  		unlock: make(chan byte, 1),
    75  	}
    76  }
    77  
    78  // Clean cleans old unused locks. Returns removed keys.
    79  func Clean() []string {
    80  	locks.lock <- 1
    81  	defer func() { <-locks.lock }()
    82  
    83  	toDelete := make([]string, 0, len(locks.list))
    84  	for key, ch := range locks.list {
    85  		select {
    86  		case <-ch:
    87  			close(ch)
    88  			toDelete = append(toDelete, key)
    89  		default:
    90  		}
    91  	}
    92  
    93  	for _, del := range toDelete {
    94  		delete(locks.list, del)
    95  	}
    96  
    97  	return toDelete
    98  }
    99  
   100  // Create and get the channel for the specified key.
   101  func getChan(key string) chan byte {
   102  	locks.lock <- 1
   103  	defer func() { <-locks.lock }()
   104  
   105  	if locks.list[key] == nil {
   106  		locks.list[key] = make(chan byte, 1)
   107  		locks.list[key] <- 1
   108  	}
   109  	return locks.list[key]
   110  }
   111  
   112  // Return a new string with unique elements.
   113  func unique(arr []string) []string {
   114  	if arr == nil || len(arr) <= 1 {
   115  		return arr
   116  	}
   117  
   118  	found := map[string]bool{}
   119  	result := make([]string, 0, len(arr))
   120  	for _, v := range arr {
   121  		if !found[v] {
   122  			found[v] = true
   123  			result = append(result, v)
   124  		}
   125  	}
   126  	return result
   127  }