Non racy tests for memlock.

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"
         "infra/gae/libs/wrapper"
+        "infra/libs/clock"
         "sync/atomic"
         "time"
 
         "github.com/luci/luci-go/common/logging"
         "golang.org/x/net/context"
 
         "appengine/memcache"
 )
 
 // ErrFailedToLock is returned from TryWithLock when it fails to obtain a lock
@@ skipping 26 matching lines @@
 // 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.
-func TryWithLock(c context.Context, key, clientID string, f func(check func() bool) error) error {
+func TryWithLock(ctx context.Context, key, clientID string, f func(check func() bool) error) error {
         if len(clientID) == 0 {
                 return ErrEmptyClientID
         }
 
-        c = logging.SetField(c, "key", key)
-        c = logging.SetField(c, "clientID", clientID)
-        log := logging.Get(c)
-        mc := wrapper.GetMC(c)
+        ctx = logging.SetField(ctx, "key", key)
+        ctx = logging.SetField(ctx, "clientID", clientID)
+        log := logging.Get(ctx)
+        mc := wrapper.GetMC(ctx)
 
         key = memlockKeyPrefix + key
         cid := []byte(clientID)
 
         // checkAnd gets the current value from memcache, and then attempts to do the
         // checkOp (which can either be `refresh` or `release`). These pieces of
         // functionality are necessarially intertwined, because CAS only works with
         // the exact-same *Item which was returned from a Get.
         //
         // refresh will attempt to CAS the item with the same content to reset it's
@@ skipping 64 matching lines @@
         // 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.).
         go func() {
                 defer close(stoppedChan)
 
         checkLoop:
                 for {
                         select {
                         case <-stopChan:
                                 break checkLoop
-                        case <-time.After(delay):
+                        case <-clock.Get(ctx).After(delay):
                         }
                         if !checkAnd(refresh) {
                                 atomic.StoreUint32(&held, 0)
                                 log.Warningf("lost lock: %s", err)
-                                break
+                                return
                         }
                 }
 
                 checkAnd(release)
                 atomic.StoreUint32(&held, 0)
         }()
 
         return f(func() bool { return atomic.LoadUint32(&held) == 1 })
 }