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 <[email protected]>
Reviewed-by: Damien Neil <[email protected]>
Run-TryBot: Ian Lance Taylor <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

Author: jba

semaphore/semaphore.go

@@ -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()
+}

semaphore/semaphore_bench_test.go

@@ -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)
+			})
+		}
+	}
+}