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 func (s *Weighted) Acquire(ctx context.Context, n int64) error { 39 done := ctx.Done() 40 41 s.mu.Lock() 42 select { 43 case <-done: 44 // ctx becoming done has "happened before" acquiring the semaphore, 45 // whether it became done before the call began or while we were 46 // waiting for the mutex. We prefer to fail even if we could acquire 47 // the mutex without blocking. 48 s.mu.Unlock() 49 return ctx.Err() 50 default: 51 } 52 if s.size-s.cur >= n && s.waiters.Len() == 0 { 53 // Since we hold s.mu and haven't synchronized since checking done, if 54 // ctx becomes done before we return here, it becoming done must have 55 // "happened concurrently" with this call - it cannot "happen before" 56 // we return in this branch. So, we're ok to always acquire here. 57 s.cur += n 58 s.mu.Unlock() 59 return nil 60 } 61 62 if n > s.size { 63 // Don't make other Acquire calls block on one that's doomed to fail. 64 s.mu.Unlock() 65 <-done 66 return ctx.Err() 67 } 68 69 ready := make(chan struct{}) 70 w := waiter{n: n, ready: ready} 71 elem := s.waiters.PushBack(w) 72 s.mu.Unlock() 73 74 select { 75 case <-done: 76 s.mu.Lock() 77 select { 78 case <-ready: 79 // Acquired the semaphore after we were canceled. 80 // Pretend we didn't and put the tokens back. 81 s.cur -= n 82 s.notifyWaiters() 83 default: 84 isFront := s.waiters.Front() == elem 85 s.waiters.Remove(elem) 86 // If we're at the front and there're extra tokens left, notify other waiters. 87 if isFront && s.size > s.cur { 88 s.notifyWaiters() 89 } 90 } 91 s.mu.Unlock() 92 return ctx.Err() 93 94 case <-ready: 95 // Acquired the semaphore. Check that ctx isn't already done. 96 // We check the done channel instead of calling ctx.Err because we 97 // already have the channel, and ctx.Err is O(n) with the nesting 98 // depth of ctx. 99 select { 100 case <-done: 101 s.Release(n) 102 return ctx.Err() 103 default: 104 } 105 return nil 106 } 107 } 108 109 // TryAcquire acquires the semaphore with a weight of n without blocking. 110 // On success, returns true. On failure, returns false and leaves the semaphore unchanged. 111 func (s *Weighted) TryAcquire(n int64) bool { 112 s.mu.Lock() 113 success := s.size-s.cur >= n && s.waiters.Len() == 0 114 if success { 115 s.cur += n 116 } 117 s.mu.Unlock() 118 return success 119 } 120 121 // Release releases the semaphore with a weight of n. 122 func (s *Weighted) Release(n int64) { 123 s.mu.Lock() 124 s.cur -= n 125 if s.cur < 0 { 126 s.mu.Unlock() 127 panic("semaphore: released more than held") 128 } 129 s.notifyWaiters() 130 s.mu.Unlock() 131 } 132 133 func (s *Weighted) notifyWaiters() { 134 for { 135 next := s.waiters.Front() 136 if next == nil { 137 break // No more waiters blocked. 138 } 139 140 w := next.Value.(waiter) 141 if s.size-s.cur < w.n { 142 // Not enough tokens for the next waiter. We could keep going (to try to 143 // find a waiter with a smaller request), but under load that could cause 144 // starvation for large requests; instead, we leave all remaining waiters 145 // blocked. 146 // 147 // Consider a semaphore used as a read-write lock, with N tokens, N 148 // readers, and one writer. Each reader can Acquire(1) to obtain a read 149 // lock. The writer can Acquire(N) to obtain a write lock, excluding all 150 // of the readers. If we allow the readers to jump ahead in the queue, 151 // the writer will starve — there is always one token available for every 152 // reader. 153 break 154 } 155 156 s.cur += w.n 157 s.waiters.Remove(next) 158 close(w.ready) 159 } 160 } 161