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 15 matching lines @@
         release checkOp = "release"
         refresh         = "refresh"
 )
 
 // memcacheLockTime is the expiration time of the memcache entry. If the lock
 // is correctly released, then it will be released before this time. It's a
 // var so we can override it in the tests.
 var memcacheLockTime = 16 * time.Second
 
 // 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.
+// sucessful. The context provided to f will be canceled (e.g. ctx.Done() will
+// be closed) if memlock detects that we've lost the lock.
 //
 // 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.
-func TryWithLock(ctx context.Context, key, clientID string, f func(check func() bool) error) error {
+//
+// Note that the lock provided by TryWithLock is a best-effort lock... some
+// other form of locking or synchronization should be used inside of f (such as
+// Datastore transactions) to ensure that f is, in fact, operating exclusively.
+// The purpose of TryWithLock is to have a cheap filter to prevent unnecessary
+// contention on heavier synchronization primitives like transactions.
+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)
         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)
 }