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..ddf20972d2e9bd58f7da7db1b1e798a0358b6f8d 100644
--- a/go/src/infra/gae/libs/memlock/memlock_test.go
+++ b/go/src/infra/gae/libs/memlock/memlock_test.go
@@ -6,7 +6,7 @@ package memlock
 
 import (
 	"fmt"
-	"runtime"
+	"sync"
 	"testing"
 	"time"
 
@@ -25,6 +25,17 @@ func init() {
 	memcacheLockTime = time.Millisecond * 4
 }
 
+type getBlockerFilter struct {
+	gae.Memcache
+	sync.Mutex
+}
+
+func (f *getBlockerFilter) Get(key string) (gae.MCItem, error) {
+	f.Lock()
+	defer f.Unlock()
+	return f.Memcache.Get(key)
+}
+
 func TestSimple(t *testing.T) {
 	// TODO(riannucci): Mock time.After so that we don't have to delay for real.
 
@@ -41,12 +52,24 @@ func TestSimple(t *testing.T) {
 			default:
 			}
 		})
+		waitFalse := func(ctx context.Context) {
+		loop:
+			for {
+				select {
+				case <-blocker:
+					continue
+				case <-ctx.Done():
+					break loop
+				}
+			}
+		}
+
 		ctx, fb := featureBreaker.FilterMC(memory.Use(ctx), nil)
 		mc := gae.GetMC(ctx)
 
 		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 +79,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,26 +87,25 @@ 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)
-
-				waitFalse := func() {
-					<-blocker
-					for i := 0; i < 3; i++ {
-						if check() {
-							runtime.Gosched()
-						}
+			err := TryWithLock(ctx, "testkey", "id", func(ctx context.Context) error {
+				isDone := func() bool {
+					select {
+					case <-ctx.Done():
+						return true
+					default:
+						return false
 					}
-					So(check(), ShouldBeFalse)
 				}
 
+				So(isDone(), ShouldBeFalse)
+
 				Convey("waiting for a while keeps refreshing the lock", func() {
 					// simulate waiting for 64*delay time, and ensuring that checkLoop
 					// runs that many times.
@@ -91,24 +113,19 @@ func TestSimple(t *testing.T) {
 						<-blocker
 						clk.Add(delay)
 					}
-					So(check(), ShouldBeTrue)
+					So(isDone(), ShouldBeFalse)
 				})
 
 				Convey("but sometimes we might lose it", func() {
 					Convey("because it was evicted", func() {
 						mc.Delete(key)
 						clk.Add(memcacheLockTime)
-						waitFalse()
-					})
-
-					Convey("or because it was stolen", func() {
-						mc.Set(mc.NewItem(key).SetValue([]byte("wat")))
-						waitFalse()
+						waitFalse(ctx)
 					})
 
 					Convey("or because of service issues", func() {
 						fb.BreakFeatures(nil, "CompareAndSwap")
-						waitFalse()
+						waitFalse(ctx)
 					})
 				})
 				return nil
@@ -116,6 +133,29 @@ func TestSimple(t *testing.T) {
 			So(err, ShouldBeNil)
 		})
 
+		Convey("can lose it when it gets stolen", func() {
+			gbf := &getBlockerFilter{}
+			ctx = gae.AddMCFilters(ctx, func(_ context.Context, mc gae.Memcache) gae.Memcache {
+				gbf.Memcache = mc
+				return gbf
+			})
+			mc = gae.GetMC(ctx)
+			err := TryWithLock(ctx, "testkey", "id", func(ctx context.Context) error {
+				// simulate waiting for 64*delay time, and ensuring that checkLoop
+				// runs that many times.
+				for i := 0; i < 64; i++ {
+					<-blocker
+					clk.Add(delay)
+				}
+				gbf.Lock()
+				mc.Set(mc.NewItem(key).SetValue([]byte("wat")))
+				gbf.Unlock()
+				waitFalse(ctx)
+				return nil
+			})
+			So(err, ShouldBeNil)
+		})
+
 		Convey("an empty context id is an error", func() {
 			So(TryWithLock(ctx, "testkey", "", nil), ShouldEqual, ErrEmptyClientID)
 		})