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.
15 "golang.org/x/sync/errgroup"
16 "golang.org/x/sync/semaphore"
19 const maxSleep = 1 * time.Millisecond
21 func HammerWeighted(sem *semaphore.Weighted, n int64, loops int) {
22 for i := 0; i < loops; i++ {
23 sem.Acquire(context.Background(), n)
24 time.Sleep(time.Duration(rand.Int63n(int64(maxSleep/time.Nanosecond))) * time.Nanosecond)
29 func TestWeighted(t *testing.T) {
32 n := runtime.GOMAXPROCS(0)
34 sem := semaphore.NewWeighted(int64(n))
37 for i := 0; i < n; i++ {
41 HammerWeighted(sem, int64(i), loops)
47 func TestWeightedPanic(t *testing.T) {
52 t.Fatal("release of an unacquired weighted semaphore did not panic")
55 w := semaphore.NewWeighted(1)
59 func TestWeightedTryAcquire(t *testing.T) {
62 ctx := context.Background()
63 sem := semaphore.NewWeighted(2)
66 tries = append(tries, sem.TryAcquire(1))
67 tries = append(tries, sem.TryAcquire(1))
71 tries = append(tries, sem.TryAcquire(1))
73 tries = append(tries, sem.TryAcquire(1))
75 want := []bool{true, false, true, false}
76 for i := range tries {
77 if tries[i] != want[i] {
78 t.Errorf("tries[%d]: got %t, want %t", i, tries[i], want[i])
83 func TestWeightedAcquire(t *testing.T) {
86 ctx := context.Background()
87 sem := semaphore.NewWeighted(2)
88 tryAcquire := func(n int64) bool {
89 ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
91 return sem.Acquire(ctx, n) == nil
96 tries = append(tries, tryAcquire(1))
97 tries = append(tries, tryAcquire(1))
101 tries = append(tries, tryAcquire(1))
103 tries = append(tries, tryAcquire(1))
105 want := []bool{true, false, true, false}
106 for i := range tries {
107 if tries[i] != want[i] {
108 t.Errorf("tries[%d]: got %t, want %t", i, tries[i], want[i])
113 func TestWeightedDoesntBlockIfTooBig(t *testing.T) {
117 sem := semaphore.NewWeighted(n)
119 ctx, cancel := context.WithCancel(context.Background())
121 go sem.Acquire(ctx, n+1)
124 g, ctx := errgroup.WithContext(context.Background())
125 for i := n * 3; i > 0; i-- {
127 err := sem.Acquire(ctx, 1)
129 time.Sleep(1 * time.Millisecond)
135 if err := g.Wait(); err != nil {
136 t.Errorf("semaphore.NewWeighted(%v) failed to AcquireCtx(_, 1) with AcquireCtx(_, %v) pending", n, n+1)
140 // TestLargeAcquireDoesntStarve times out if a large call to Acquire starves.
141 // Merely returning from the test function indicates success.
142 func TestLargeAcquireDoesntStarve(t *testing.T) {
145 ctx := context.Background()
146 n := int64(runtime.GOMAXPROCS(0))
147 sem := semaphore.NewWeighted(n)
150 var wg sync.WaitGroup
152 for i := n; i > 0; i-- {
160 time.Sleep(1 * time.Millisecond)
173 // translated from https://github.com/zhiqiangxu/util/blob/master/mutex/crwmutex_test.go#L43
174 func TestAllocCancelDoesntStarve(t *testing.T) {
175 sem := semaphore.NewWeighted(10)
177 // Block off a portion of the semaphore so that Acquire(_, 10) can eventually succeed.
178 sem.Acquire(context.Background(), 1)
180 // In the background, Acquire(_, 10).
181 ctx, cancel := context.WithCancel(context.Background())
187 // Wait until the Acquire(_, 10) call blocks.
188 for sem.TryAcquire(1) {
193 // Now try to grab a read lock, and simultaneously unblock the Acquire(_, 10) call.
194 // Both Acquire calls should unblock and return, in either order.
197 err := sem.Acquire(context.Background(), 1)
199 t.Fatalf("Acquire(_, 1) failed unexpectedly: %v", err)