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.
38func (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.
111func (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.
122func (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
133func (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