semaphore: add a weighted semaphore implementation
This package provides a weighted semaphore that is context-aware. The code is derived from a similar package inside Google. Change-Id: Id1dad96d79e8ccfd289e4299e8265aa5bdad3a5b Reviewed-on: https://go-review.googlesource.com/38298 Reviewed-by: Ian Lance Taylor <iant@golang.org> Reviewed-by: Damien Neil <dneil@google.com> Run-TryBot: Ian Lance Taylor <iant@golang.org> TryBot-Result: Gobot Gobot <gobot@golang.org>
This commit is contained in:
Jonathan Amsterdam
Ian Lance Taylor
parent
a60ad46e0e
commit
5a06fca2c3
3 changed files with 429 additions and 0 deletions
131
semaphore/semaphore.go
Normal file
131
semaphore/semaphore.go
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
// Copyright 2017 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package semaphore provides a weighted semaphore implementation.
|
||||
package semaphore // import "golang.org/x/sync/semaphore"
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"sync"
|
||||
|
||||
// Use the old context because packages that depend on this one
|
||||
// (e.g. cloud.google.com/go/...) must run on Go 1.6.
|
||||
// TODO(jba): update to "context" when possible.
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
type waiter struct {
|
||||
n int64
|
||||
ready chan<- struct{} // Closed when semaphore acquired.
|
||||
}
|
||||
|
||||
// NewWeighted creates a new weighted semaphore with the given
|
||||
// maximum combined weight for concurrent access.
|
||||
func NewWeighted(n int64) *Weighted {
|
||||
w := &Weighted{size: n}
|
||||
return w
|
||||
}
|
||||
|
||||
// Weighted provides a way to bound concurrent access to a resource.
|
||||
// The callers can request access with a given weight.
|
||||
type Weighted struct {
|
||||
size int64
|
||||
cur int64
|
||||
mu sync.Mutex
|
||||
waiters list.List
|
||||
}
|
||||
|
||||
// Acquire acquires the semaphore with a weight of n, blocking only until ctx
|
||||
// is done. On success, returns nil. On failure, returns ctx.Err() and leaves
|
||||
// the semaphore unchanged.
|
||||
//
|
||||
// If ctx is already done, Acquire may still succeed without blocking.
|
||||
func (s *Weighted) Acquire(ctx context.Context, n int64) error {
|
||||
s.mu.Lock()
|
||||
if s.size-s.cur >= n && s.waiters.Len() == 0 {
|
||||
s.cur += n
|
||||
s.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
if n > s.size {
|
||||
// Don't make other Acquire calls block on one that's doomed to fail.
|
||||
s.mu.Unlock()
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
ready := make(chan struct{})
|
||||
w := waiter{n: n, ready: ready}
|
||||
elem := s.waiters.PushBack(w)
|
||||
s.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
err := ctx.Err()
|
||||
s.mu.Lock()
|
||||
select {
|
||||
case <-ready:
|
||||
// Acquired the semaphore after we were canceled. Rather than trying to
|
||||
// fix up the queue, just pretend we didn't notice the cancelation.
|
||||
err = nil
|
||||
default:
|
||||
s.waiters.Remove(elem)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
return err
|
||||
|
||||
case <-ready:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire acquires the semaphore with a weight of n without blocking.
|
||||
// On success, returns true. On failure, returns false and leaves the semaphore unchanged.
|
||||
func (s *Weighted) TryAcquire(n int64) bool {
|
||||
s.mu.Lock()
|
||||
success := s.size-s.cur >= n && s.waiters.Len() == 0
|
||||
if success {
|
||||
s.cur += n
|
||||
}
|
||||
s.mu.Unlock()
|
||||
return success
|
||||
}
|
||||
|
||||
// Release releases the semaphore with a weight of n.
|
||||
func (s *Weighted) Release(n int64) {
|
||||
s.mu.Lock()
|
||||
s.cur -= n
|
||||
if s.cur < 0 {
|
||||
s.mu.Unlock()
|
||||
panic("semaphore: bad release")
|
||||
}
|
||||
for {
|
||||
next := s.waiters.Front()
|
||||
if next == nil {
|
||||
break // No more waiters blocked.
|
||||
}
|
||||
|
||||
w := next.Value.(waiter)
|
||||
if s.size-s.cur < w.n {
|
||||
// Not enough tokens for the next waiter. We could keep going (to try to
|
||||
// find a waiter with a smaller request), but under load that could cause
|
||||
// starvation for large requests; instead, we leave all remaining waiters
|
||||
// blocked.
|
||||
//
|
||||
// Consider a semaphore used as a read-write lock, with N tokens, N
|
||||
// readers, and one writer. Each reader can Acquire(1) to obtain a read
|
||||
// lock. The writer can Acquire(N) to obtain a write lock, excluding all
|
||||
// of the readers. If we allow the readers to jump ahead in the queue,
|
||||
// the writer will starve — there is always one token available for every
|
||||
// reader.
|
||||
break
|
||||
}
|
||||
|
||||
s.cur += w.n
|
||||
s.waiters.Remove(next)
|
||||
close(w.ready)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
130
semaphore/semaphore_bench_test.go
Normal file
130
semaphore/semaphore_bench_test.go
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
// Copyright 2017 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// +build go1.7
|
||||
|
||||
package semaphore
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// weighted is an interface matching a subset of *Weighted. It allows
|
||||
// alternate implementations for testing and benchmarking.
|
||||
type weighted interface {
|
||||
Acquire(context.Context, int64) error
|
||||
TryAcquire(int64) bool
|
||||
Release(int64)
|
||||
}
|
||||
|
||||
// semChan implements Weighted using a channel for
|
||||
// comparing against the condition variable-based implementation.
|
||||
type semChan chan struct{}
|
||||
|
||||
func newSemChan(n int64) semChan {
|
||||
return semChan(make(chan struct{}, n))
|
||||
}
|
||||
|
||||
func (s semChan) Acquire(_ context.Context, n int64) error {
|
||||
for i := int64(0); i < n; i++ {
|
||||
s <- struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s semChan) TryAcquire(n int64) bool {
|
||||
if int64(len(s))+n > int64(cap(s)) {
|
||||
return false
|
||||
}
|
||||
|
||||
for i := int64(0); i < n; i++ {
|
||||
s <- struct{}{}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s semChan) Release(n int64) {
|
||||
for i := int64(0); i < n; i++ {
|
||||
<-s
|
||||
}
|
||||
}
|
||||
|
||||
// acquireN calls Acquire(size) on sem N times and then calls Release(size) N times.
|
||||
func acquireN(b *testing.B, sem weighted, size int64, N int) {
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
for j := 0; j < N; j++ {
|
||||
sem.Acquire(context.Background(), size)
|
||||
}
|
||||
for j := 0; j < N; j++ {
|
||||
sem.Release(size)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// tryAcquireN calls TryAcquire(size) on sem N times and then calls Release(size) N times.
|
||||
func tryAcquireN(b *testing.B, sem weighted, size int64, N int) {
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
for j := 0; j < N; j++ {
|
||||
if !sem.TryAcquire(size) {
|
||||
b.Fatalf("TryAcquire(%v) = false, want true", size)
|
||||
}
|
||||
}
|
||||
for j := 0; j < N; j++ {
|
||||
sem.Release(size)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkNewSeq(b *testing.B) {
|
||||
for _, cap := range []int64{1, 128} {
|
||||
b.Run(fmt.Sprintf("Weighted-%d", cap), func(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
_ = NewWeighted(cap)
|
||||
}
|
||||
})
|
||||
b.Run(fmt.Sprintf("semChan-%d", cap), func(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
_ = newSemChan(cap)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkAcquireSeq(b *testing.B) {
|
||||
for _, c := range []struct {
|
||||
cap, size int64
|
||||
N int
|
||||
}{
|
||||
{1, 1, 1},
|
||||
{2, 1, 1},
|
||||
{16, 1, 1},
|
||||
{128, 1, 1},
|
||||
{2, 2, 1},
|
||||
{16, 2, 8},
|
||||
{128, 2, 64},
|
||||
{2, 1, 2},
|
||||
{16, 8, 2},
|
||||
{128, 64, 2},
|
||||
} {
|
||||
for _, w := range []struct {
|
||||
name string
|
||||
w weighted
|
||||
}{
|
||||
{"Weighted", NewWeighted(c.cap)},
|
||||
{"semChan", newSemChan(c.cap)},
|
||||
} {
|
||||
b.Run(fmt.Sprintf("%s-acquire-%d-%d-%d", w.name, c.cap, c.size, c.N), func(b *testing.B) {
|
||||
acquireN(b, w.w, c.size, c.N)
|
||||
})
|
||||
b.Run(fmt.Sprintf("%s-tryAcquire-%d-%d-%d", w.name, c.cap, c.size, c.N), func(b *testing.B) {
|
||||
tryAcquireN(b, w.w, c.size, c.N)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
168
semaphore/semaphore_test.go
Normal file
168
semaphore/semaphore_test.go
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
// Copyright 2017 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package semaphore
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
const maxSleep = 1 * time.Millisecond
|
||||
|
||||
func HammerWeighted(sem *Weighted, n int64, loops int) {
|
||||
for i := 0; i < loops; i++ {
|
||||
sem.Acquire(context.Background(), n)
|
||||
time.Sleep(time.Duration(rand.Int63n(int64(maxSleep/time.Nanosecond))) * time.Nanosecond)
|
||||
sem.Release(n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWeighted(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
n := runtime.GOMAXPROCS(0)
|
||||
sem := NewWeighted(int64(n))
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(n)
|
||||
for i := 0; i < n; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
HammerWeighted(sem, int64(i), 1000)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestWeightedPanic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
defer func() {
|
||||
if recover() == nil {
|
||||
t.Fatal("release of an unacquired weighted semaphore did not panic")
|
||||
}
|
||||
}()
|
||||
w := NewWeighted(1)
|
||||
w.Release(1)
|
||||
}
|
||||
|
||||
func TestWeightedTryAcquire(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
sem := NewWeighted(2)
|
||||
tries := []bool{}
|
||||
sem.Acquire(ctx, 1)
|
||||
tries = append(tries, sem.TryAcquire(1))
|
||||
tries = append(tries, sem.TryAcquire(1))
|
||||
|
||||
sem.Release(2)
|
||||
|
||||
tries = append(tries, sem.TryAcquire(1))
|
||||
sem.Acquire(ctx, 1)
|
||||
tries = append(tries, sem.TryAcquire(1))
|
||||
|
||||
want := []bool{true, false, true, false}
|
||||
for i := range tries {
|
||||
if tries[i] != want[i] {
|
||||
t.Errorf("tries[%d]: got %t, want %t", i, tries[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWeightedAcquire(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
sem := NewWeighted(2)
|
||||
tryAcquire := func(n int64) bool {
|
||||
ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
|
||||
defer cancel()
|
||||
return sem.Acquire(ctx, n) == nil
|
||||
}
|
||||
|
||||
tries := []bool{}
|
||||
sem.Acquire(ctx, 1)
|
||||
tries = append(tries, tryAcquire(1))
|
||||
tries = append(tries, tryAcquire(1))
|
||||
|
||||
sem.Release(2)
|
||||
|
||||
tries = append(tries, tryAcquire(1))
|
||||
sem.Acquire(ctx, 1)
|
||||
tries = append(tries, tryAcquire(1))
|
||||
|
||||
want := []bool{true, false, true, false}
|
||||
for i := range tries {
|
||||
if tries[i] != want[i] {
|
||||
t.Errorf("tries[%d]: got %t, want %t", i, tries[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWeightedDoesntBlockIfTooBig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const n = 2
|
||||
sem := NewWeighted(n)
|
||||
{
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go sem.Acquire(ctx, n+1)
|
||||
}
|
||||
|
||||
g, ctx := errgroup.WithContext(context.Background())
|
||||
for i := n * 3; i > 0; i-- {
|
||||
g.Go(func() error {
|
||||
err := sem.Acquire(ctx, 1)
|
||||
if err == nil {
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
sem.Release(1)
|
||||
}
|
||||
return err
|
||||
})
|
||||
}
|
||||
if err := g.Wait(); err != nil {
|
||||
t.Errorf("NewWeighted(%v) failed to AcquireCtx(_, 1) with AcquireCtx(_, %v) pending", n, n+1)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLargeAcquireDoesntStarve times out if a large call to Acquire starves.
|
||||
// Merely returning from the test function indicates success.
|
||||
func TestLargeAcquireDoesntStarve(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
n := int64(runtime.GOMAXPROCS(0))
|
||||
sem := NewWeighted(n)
|
||||
running := true
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(int(n))
|
||||
for i := n; i > 0; i-- {
|
||||
sem.Acquire(ctx, 1)
|
||||
go func() {
|
||||
defer func() {
|
||||
sem.Release(1)
|
||||
wg.Done()
|
||||
}()
|
||||
for running {
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
sem.Release(1)
|
||||
sem.Acquire(ctx, 1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
sem.Acquire(ctx, n)
|
||||
running = false
|
||||
sem.Release(n)
|
||||
wg.Wait()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue