Simplify memlock and make it less racy.

go/src/infra/gae/libs/memlock/memlock_test.go

Index: go/src/infra/gae/libs/memlock/memlock_test.go
diff --git a/go/src/infra/gae/libs/memlock/memlock_test.go b/go/src/infra/gae/libs/memlock/memlock_test.go
index e6de4fb4b1a56eb81bda9c46d09d5d906df3cee0..4051ba6117982f4909d3bfde4533cf51d37d0402 100644
--- a/go/src/infra/gae/libs/memlock/memlock_test.go
+++ b/go/src/infra/gae/libs/memlock/memlock_test.go
@@ -6,7 +6,6 @@ package memlock
 
 import (
 	"fmt"
-	"runtime"
 	"testing"
 	"time"
 
@@ -46,7 +45,7 @@ func TestSimple(t *testing.T) {
 
 		Convey("fails to acquire when memcache is down", func() {
 			fb.BreakFeatures(nil, "Add")
-			err := TryWithLock(ctx, "testkey", "id", func(check func() bool) error {
+			err := TryWithLock(ctx, "testkey", "id", func(context.Context) error {
 				// should never reach here
 				So(false, ShouldBeTrue)
 				return nil
@@ -56,7 +55,7 @@ func TestSimple(t *testing.T) {
 
 		Convey("returns the inner error", func() {
 			toRet := fmt.Errorf("sup")
-			err := TryWithLock(ctx, "testkey", "id", func(check func() bool) error {
+			err := TryWithLock(ctx, "testkey", "id", func(context.Context) error {
 				return toRet
 			})
 			So(err, ShouldEqual, toRet)
@@ -64,24 +63,36 @@ func TestSimple(t *testing.T) {
 
 		Convey("returns the error", func() {
 			toRet := fmt.Errorf("sup")
-			err := TryWithLock(ctx, "testkey", "id", func(check func() bool) error {
+			err := TryWithLock(ctx, "testkey", "id", func(context.Context) error {
 				return toRet
 			})
 			So(err, ShouldEqual, toRet)
 		})
 
 		Convey("can acquire when empty", func() {
-			err := TryWithLock(ctx, "testkey", "id", func(check func() bool) error {
-				So(check(), ShouldBeTrue)
+			err := TryWithLock(ctx, "testkey", "id", func(ctx context.Context) error {
+				isDone := func() bool {
+					select {
+					case <-ctx.Done():
+						return true
+					default:
+						return false
+					}
+				}
+
+				So(isDone(), ShouldBeFalse)
 
 				waitFalse := func() {
-					<-blocker
-					for i := 0; i < 3; i++ {
-						if check() {
-							runtime.Gosched()
+				loop:
+					for {
+						select {
+						case <-blocker:
+							continue
+						case <-ctx.Done():
+							break loop
 						}
 					}
-					So(check(), ShouldBeFalse)
+					So(isDone(), ShouldBeTrue)
 				}
 
 				Convey("waiting for a while keeps refreshing the lock", func() {
@@ -91,7 +102,7 @@ func TestSimple(t *testing.T) {
 						<-blocker
 						clk.Add(delay)
 					}
-					So(check(), ShouldBeTrue)
+					So(isDone(), ShouldBeFalse)
 				})
 
 				Convey("but sometimes we might lose it", func() {