sync2: Add AcquireContext to sema

And fix flaky test

Signed-off-by: Sugu Sougoumarane <ssougou@gmail.com>

go/sync2/semaphore.go

@@ -21,6 +21,7 @@ package sync2
 // cases, you just want a familiar API.
 
 import (
+	"context"
 	"time"
 )
 
@@ -61,6 +62,19 @@ func (sem *Semaphore) Acquire() bool {
 	}
 }
 
+// AcquireContext returns true on successful acquisition, and
+// false on context expiry. Timeout is ignored.
+func (sem *Semaphore) AcquireContext(ctx context.Context) bool {
+	tm := time.NewTimer(sem.timeout)
+	defer tm.Stop()
+	select {
+	case <-sem.slots:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
 // TryAcquire acquires a semaphore if it's immediately available.
 // It returns false otherwise.
 func (sem *Semaphore) TryAcquire() bool {

go/sync2/semaphore_test.go

@@ -17,8 +17,11 @@ limitations under the License.
 package sync2
 
 import (
+	"context"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestSemaNoTimeout(t *testing.T) {
@@ -26,42 +29,51 @@ func TestSemaNoTimeout(t *testing.T) {
 	s.Acquire()
 	released := false
 	go func() {
-		time.Sleep(10 * time.Millisecond)
 		released = true
 		s.Release()
 	}()
 	s.Acquire()
-	if !released {
-		t.Errorf("release: false, want true")
-	}
+	assert.True(t, released)
 }
 
 func TestSemaTimeout(t *testing.T) {
-	s := NewSemaphore(1, 5*time.Millisecond)
+	s := NewSemaphore(1, 1*time.Millisecond)
 	s.Acquire()
+	release := make(chan struct{})
+	released := make(chan struct{})
 	go func() {
-		time.Sleep(10 * time.Millisecond)
+		<-release
 		s.Release()
+		released <- struct{}{}
 	}()
-	if s.Acquire() {
-		t.Errorf("Acquire: true, want false")
-	}
-	time.Sleep(10 * time.Millisecond)
-	if !s.Acquire() {
-		t.Errorf("Acquire: false, want true")
-	}
+	assert.False(t, s.Acquire())
+	release <- struct{}{}
+	<-released
+	assert.True(t, s.Acquire())
+}
+
+func TestSemaAcquireContext(t *testing.T) {
+	s := NewSemaphore(1, 0)
+	s.Acquire()
+	release := make(chan struct{})
+	released := make(chan struct{})
+	go func() {
+		<-release
+		s.Release()
+		released <- struct{}{}
+	}()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	assert.False(t, s.AcquireContext(ctx))
+	release <- struct{}{}
+	<-released
+	assert.True(t, s.AcquireContext(context.Background()))
 }
 
 func TestSemaTryAcquire(t *testing.T) {
 	s := NewSemaphore(1, 0)
-	if !s.TryAcquire() {
-		t.Errorf("TryAcquire: false, want true")
-	}
-	if s.TryAcquire() {
-		t.Errorf("TryAcquire: true, want false")
-	}
+	assert.True(t, s.TryAcquire())
+	assert.False(t, s.TryAcquire())
 	s.Release()
-	if !s.TryAcquire() {
-		t.Errorf("TryAcquire: false, want true")
-	}
+	assert.True(t, s.TryAcquire())
 }