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.
6package semaphore // import "golang.org/x/sync/semaphore"
7
8import (
9	"container/list"
10	"context"
11	"sync"
12)
13
14type 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.
21func 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.
28type 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.
40func (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.
87func (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.
98func (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
109func (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