1 // Copyright 2017 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 semaphore provides a weighted semaphore implementation. 6 package semaphore // import "golang.org/x/sync/semaphore" 7 8 import ( 9 "container/list" 10 "context" 11 "sync" 12 ) 13 14 type waiter struct { 15 n int64 16 ready chan<- struct{} // Closed when semaphore acquired. 17 } 18 19 // NewWeighted creates a new weighted semaphore with the given 20 // maximum combined weight for concurrent access. 21 func NewWeighted(n int64) *Weighted { 22 w := &Weighted{size: n} 23 return w 24 } 25 26 // Weighted provides a way to bound concurrent access to a resource. 27 // The callers can request access with a given weight. 28 type Weighted struct { 29 size int64 30 cur int64 31 mu sync.Mutex 32 waiters list.List 33 } 34 35 // Acquire acquires the semaphore with a weight of n, blocking until resources 36 // are available or ctx is done. On success, returns nil. On failure, returns 37 // ctx.Err() and leaves the semaphore unchanged. 38 // 39 // If ctx is already done, Acquire may still succeed without blocking. 40 func (s *Weighted) Acquire(ctx context.Context, n int64) error { 41 s.mu.Lock() 42 if s.size-s.cur >= n && s.waiters.Len() == 0 { 43 s.cur += n 44 s.mu.Unlock() 45 return nil 46 } 47 48 if n > s.size { 49 // Don't make other Acquire calls block on one that's doomed to fail. 50 s.mu.Unlock() 51 <-ctx.Done() 52 return ctx.Err() 53 } 54 55 ready := make(chan struct{}) 56 w := waiter{n: n, ready: ready} 57 elem := s.waiters.PushBack(w) 58 s.mu.Unlock() 59 60 select { 61 case <-ctx.Done(): 62 err := ctx.Err() 63 s.mu.Lock() 64 select { 65 case <-ready: 66 // Acquired the semaphore after we were canceled. Rather than trying to 67 // fix up the queue, just pretend we didn't notice the cancelation. 68 err = nil 69 default: 70 isFront := s.waiters.Front() == elem 71 s.waiters.Remove(elem) 72 // If we're at the front and there're extra tokens left, notify other waiters. 73 if isFront && s.size > s.cur { 74 s.notifyWaiters() 75 } 76 } 77 s.mu.Unlock() 78 return err 79 80 case <-ready: 81 return nil 82 } 83 } 84 85 // TryAcquire acquires the semaphore with a weight of n without blocking. 86 // On success, returns true. On failure, returns false and leaves the semaphore unchanged. 87 func (s *Weighted) TryAcquire(n int64) bool { 88 s.mu.Lock() 89 success := s.size-s.cur >= n && s.waiters.Len() == 0 90 if success { 91 s.cur += n 92 } 93 s.mu.Unlock() 94 return success 95 } 96 97 // Release releases the semaphore with a weight of n. 98 func (s *Weighted) Release(n int64) { 99 s.mu.Lock() 100 s.cur -= n 101 if s.cur < 0 { 102 s.mu.Unlock() 103 panic("semaphore: released more than held") 104 } 105 s.notifyWaiters() 106 s.mu.Unlock() 107 } 108 109 func (s *Weighted) notifyWaiters() { 110 for { 111 next := s.waiters.Front() 112 if next == nil { 113 break // No more waiters blocked. 114 } 115 116 w := next.Value.(waiter) 117 if s.size-s.cur < w.n { 118 // Not enough tokens for the next waiter. We could keep going (to try to 119 // find a waiter with a smaller request), but under load that could cause 120 // starvation for large requests; instead, we leave all remaining waiters 121 // blocked. 122 // 123 // Consider a semaphore used as a read-write lock, with N tokens, N 124 // readers, and one writer. Each reader can Acquire(1) to obtain a read 125 // lock. The writer can Acquire(N) to obtain a write lock, excluding all 126 // of the readers. If we allow the readers to jump ahead in the queue, 127 // the writer will starve — there is always one token available for every 128 // reader. 129 break 130 } 131 132 s.cur += w.n 133 s.waiters.Remove(next) 134 close(w.ready) 135 } 136 } 137