semaphore: cancel acquisition with a done context

When acquiring from a semaphore could proceed without contention, the
previous behavior was to always do so, even when the provided context
was done. This was the documented behavior, but it could lead to
confusion. It isn't much more expensive to check the context error, so
cancel acquisition if it's done.

Fixes golang/go#63615.

goos: linux
goarch: amd64
pkg: golang.org/x/sync/semaphore
cpu: 12th Gen Intel(R) Core(TM) i5-1235U
                                        │  old.bench  │             new.bench              │
                                        │   sec/op    │   sec/op     vs base               │
AcquireSeq/Weighted-acquire-1-1-1-12      26.45n ± 2%   27.25n ± 3%  +3.04% (p=0.001 n=20)
AcquireSeq/Weighted-acquire-2-1-1-12      26.96n ± 1%   27.12n ± 1%       ~ (p=0.104 n=20)
AcquireSeq/Weighted-acquire-16-1-1-12     26.07n ± 3%   27.48n ± 1%  +5.45% (p=0.000 n=20)
AcquireSeq/Weighted-acquire-128-1-1-12    26.19n ± 2%   27.24n ± 1%  +4.01% (p=0.000 n=20)
AcquireSeq/Weighted-acquire-2-2-1-12      25.61n ± 1%   25.99n ± 2%       ~ (p=0.066 n=20)
AcquireSeq/Weighted-acquire-16-2-8-12     209.6n ± 2%   211.0n ± 3%       ~ (p=0.280 n=20)
AcquireSeq/Weighted-acquire-128-2-64-12   1.669µ ± 1%   1.721µ ± 2%  +3.09% (p=0.000 n=20)
AcquireSeq/Weighted-acquire-2-1-2-12      51.08n ± 1%   53.03n ± 2%  +3.82% (p=0.000 n=20)
AcquireSeq/Weighted-acquire-16-8-2-12     52.48n ± 2%   53.66n ± 2%  +2.26% (p=0.028 n=20)
AcquireSeq/Weighted-acquire-128-64-2-12   52.27n ± 1%   53.71n ± 2%  +2.75% (p=0.000 n=20)
geomean                                   60.06n        61.69n       +2.71%

Change-Id: I0ae1a0bb6c027461ac1a9ee71c51efd8427ab308
Reviewed-on: https://go-review.googlesource.com/c/sync/+/536275
Auto-Submit: Bryan Mills <bcmills@google.com>
Reviewed-by: Bryan Mills <bcmills@google.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

semaphore/semaphore_test.go

@@ -200,3 +200,38 @@ func TestAllocCancelDoesntStarve(t *testing.T) {
 	}
 	sem.Release(1)
 }
+
+func TestWeightedAcquireCanceled(t *testing.T) {
+	// https://go.dev/issue/63615
+	sem := semaphore.NewWeighted(2)
+	ctx, cancel := context.WithCancel(context.Background())
+	sem.Acquire(context.Background(), 1)
+	ch := make(chan struct{})
+	go func() {
+		// Synchronize with the Acquire(2) below.
+		for sem.TryAcquire(1) {
+			sem.Release(1)
+		}
+		// Now cancel ctx, and then release the token.
+		cancel()
+		sem.Release(1)
+		close(ch)
+	}()
+	// Since the context closing happens before enough tokens become available,
+	// this Acquire must fail.
+	if err := sem.Acquire(ctx, 2); err != context.Canceled {
+		t.Errorf("Acquire with canceled context returned wrong error: want context.Canceled, got %v", err)
+	}
+	// There must always be two tokens in the semaphore after the other
+	// goroutine releases the one we held at the start.
+	<-ch
+	if !sem.TryAcquire(2) {
+		t.Fatal("TryAcquire after canceled Acquire failed")
+	}
+	// Additionally verify that we don't acquire with a done context even when
+	// we wouldn't need to block to do so.
+	sem.Release(2)
+	if err := sem.Acquire(ctx, 1); err != context.Canceled {
+		t.Errorf("Acquire with canceled context returned wrong error: want context.Canceled, got %v", err)
+	}
+}