Simplify memlock and make it less racy.

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

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package memlock
 
 import (
         "fmt"
-        "runtime"
         "testing"
         "time"
 
         "infra/gae/libs/gae"
         "infra/gae/libs/gae/filters/featureBreaker"
         "infra/gae/libs/gae/memory"
 
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/clock/testclock"
         . "github.com/smartystreets/goconvey/convey"
@@ skipping 19 matching lines @@
                         select {
                         case blocker <- struct{}{}:
                         default:
                         }
                 })
                 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
                         })
                         So(err, ShouldEqual, ErrFailedToLock)
                 })
 
                 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)
                 })
 
                 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() {
                                         // simulate waiting for 64*delay time, and ensuring that checkLoop
                                         // runs that many times.
                                         for i := 0; i < 64; i++ {
                                                 <-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()
                                         })
 
                                         Convey("or because of service issues", func() {
                                                 fb.BreakFeatures(nil, "CompareAndSwap")
                                                 waitFalse()
                                         })
                                 })
                                 return nil
                         })
                         So(err, ShouldBeNil)
                 })
 
                 Convey("an empty context id is an error", func() {
                         So(TryWithLock(ctx, "testkey", "", nil), ShouldEqual, ErrEmptyClientID)
                 })
         })
 }