Simplify memlock and make it less racy.

go/src/infra/gae/libs/memlock/memlock.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 allows multiple appengine handlers to coordinate best-effort
 // mutual execution via memcache. "best-effort" here means "best-effort"...
 // memcache is not reliable. However, colliding on memcache is a lot cheaper
 // than, for example, colliding with datastore transactions.
 package memlock
 
 import (
         "bytes"
         "errors"
         "golang.org/x/net/context"
-        "sync/atomic"
         "time"
 
         "infra/gae/libs/gae"
 
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/logging"
 )
 
 // ErrFailedToLock is returned from TryWithLock when it fails to obtain a lock
 // prior to invoking the user-supplied function.
@@ skipping 24 matching lines @@
 // TryWithLock attempts to obtains the lock once, and then invokes f if
 // sucessful. The `check` function can be used within f to see if the lock is
 // still held.
 //
 // TryWithLock function returns ErrFailedToLock if it fails to obtain the lock,
 // otherwise returns the error that f returns.
 //
 // `key` is the memcache key to use (i.e. the name of the lock). Clients locking
 // 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
         }
 
         ctx = logging.SetField(ctx, "key", key)
         ctx = logging.SetField(ctx, "clientID", clientID)
         log := logging.Get(ctx)
         mc := gae.GetMC(ctx)
 
         key = memlockKeyPrefix + key
@@ skipping 49 matching lines @@
         if err != nil {
                 if err != gae.ErrMCNotStored {
                         log.Warningf("error adding: %s", err)
                 }
                 if !checkAnd(refresh) {
                         return ErrFailedToLock
                 }
         }
 
         // 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)
 }