Simplify memlock and make it less racy.

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

Index: go/src/infra/gae/libs/memlock/memlock.go
diff --git a/go/src/infra/gae/libs/memlock/memlock.go b/go/src/infra/gae/libs/memlock/memlock.go
index fee7dcd2374ddd9690f0325b29c06f8321001e77..c3004c625ecb10d80181e364f337e037842a5700 100644
--- a/go/src/infra/gae/libs/memlock/memlock.go
+++ b/go/src/infra/gae/libs/memlock/memlock.go
@@ -12,7 +12,6 @@ import (
 	"bytes"
 	"errors"
 	"golang.org/x/net/context"
-	"sync/atomic"
 	"time"
 
 	"infra/gae/libs/gae"
@@ -58,7 +57,7 @@ var memcacheLockTime = 16 * time.Second
 // the same data must use the same key. clientID is the unique identifier for
 // this client (lock-holder). If it's empty then TryWithLock() will return
 // ErrEmptyClientID.

[dnj, 2015/07/17 18:01:35]

Document that the context that "f" is called with can be cancelled to release
the lock.

[iannucci, 2015/07/17 18:08:25]

I think you're misunderstanding how the cancelation works (f doesn't have a
handle to cancel the context, just TryWithLock), but I've updating the
documentation anyhow.

[dnj, 2015/07/17 18:10:29]

Ah bad comment on my part, I was mostly suggesting you document that the
context's "Done" channel is used to detect cancellation.

-func TryWithLock(ctx context.Context, key, clientID string, f func(check func() bool) error) error {
+func TryWithLock(ctx context.Context, key, clientID string, f func(context.Context) error) error {
 	if len(clientID) == 0 {
 		return ErrEmptyClientID
 	}
@@ -128,39 +127,36 @@ func TryWithLock(ctx context.Context, key, clientID string, f func(check func()
 	}
 
 	// At this point we nominally have the lock (at least for memcacheLockTime).
-
-	stopChan := make(chan struct{})
-	stoppedChan := make(chan struct{})
-	held := uint32(1)
-
+	finished := make(chan struct{})
+	subCtx, cancelFunc := context.WithCancel(ctx)

[dnj, 2015/07/17 18:01:35]

Main concern using contexts here is that if the parent cancels its context for
some other reason, that will also cancel this context and thus release the lock.
Is this desired?

Alternatively, you can go back to the "check()" function method and use a
channel instead of 'atomic' to do the check, replacing "subCtx/cancelFunc" with:

cancelC := make(chan struct{})
cancelOnce := sync.Once{}
cancelFunc := func() {
  cancelOnce.Do(func() {
    close(cancelC)
  })
}

[iannucci, 2015/07/17 18:08:25]

Yes

 	defer func() {
-		close(stopChan)
-		<-stoppedChan // this blocks TryWithLock until the goroutine below quits.
+		cancelFunc()
+		<-finished
 	}()
 
 	// This goroutine checks to see if we still posess the lock, and refreshes it
-	// if we do. It will stop doing this when either stopChan is activated (e.g.
-	// the user's function returns) or we lose the lock (memcache flake, etc.).
+	// if we do.
 	go func() {
-		defer close(stoppedChan)
+		defer func() {
+			cancelFunc()
+			close(finished)
+		}()
 
 	checkLoop:
 		for {
 			select {
-			case <-stopChan:
+			case <-subCtx.Done():
 				break checkLoop
 			case <-clock.Get(ctx).After(delay):
 			}
 			if !checkAnd(refresh) {
-				atomic.StoreUint32(&held, 0)
 				log.Warningf("lost lock: %s", err)
 				return
 			}
 		}
 
 		checkAnd(release)
-		atomic.StoreUint32(&held, 0)
 	}()
 
-	return f(func() bool { return atomic.LoadUint32(&held) == 1 })
+	return f(subCtx)
 }