semaphore.go

     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
View as plain text