A simple memcache lock for appengine.

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"
+        "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
+// prior to invoking the user-supplied function.
+var ErrFailedToLock = errors.New("memlock: failed to obtain lock")
+
+// ErrEmptyClientID is returned from TryWithLock when you specify an empty
+// clientID.
+var ErrEmptyClientID = errors.New("memlock: empty clientID")
+
+// memlockKeyPrefix is the memcache Key prefix for all user-supplied keys.
+const memlockKeyPrefix = "memlock:"
+
+type checkOp string
+
+// var so we can override it in the tests

[M-A Ruel, 2015/05/31 20:35:10]

delay is a var so ...

+var delay = time.Second
+
+const (
+        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.
+//
+// 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 {

[M-A Ruel, 2015/05/31 20:35:11]

A nicer way IMHO would be to have this signature instead:

// Returns nil if failed to get the lock.
func WithLock(c context.Context, key, clientID string) func() {
  ...
}

so at call sites, it'd do:

if release := WithLock(...); release != nil {
  defer release()
  .. run with lock
}

I don't see what's the usecase of the check function.

[iannucci, 2015/05/31 21:07:57]

The reason for the check function is so the code running with the lock 'held'
can be informed as early as possible that the lock has been lost (e.g. due to
memcache flush or whatever). Some of the users of this lock do a couple loops
over the datastore while holding the lock (e.g. in a backend), and so it's handy
to check if we still have it (and safer than giving up the lock and then trying
to re-acquire it).

This is also why I can't implement sync.Locker safely, because unlike a regular
in-memory mutex, losing the lock during execution is actually pretty common.
This thing is less like a lock and more like a 'probably locked, but maybe not'
thing.

But if it were really reliable, then memcache would be really really slow :).
So... tradeoffs.

+        if len(clientID) == 0 {
+                return ErrEmptyClientID
+        }
+
+        log := logging.Get(c)
+        mc := wrapper.GetMC(c)
+
+        key = memlockKeyPrefix + key
+        cid := []byte(clientID)
+
+        // logf is a logging helper which includes the key and clientID in the log
+        // message automatically. The function f is meant to be one of the methods
+        // in logging.Logger (e.g. `Logger.Warningf`).
+        logf := func(f func(string, ...interface{}), fmt string, args ...interface{}) {
+                args = append([]interface{}{key, clientID}, args...)

[M-A Ruel, 2015/05/31 20:35:11]

wouldn't it be simpler to do:

return f(fmt.Sprintf("%s:%q):", key, clientID) + fmt.Sprintf(fmt, args...))

Go strings are pascal style so concatenating strings is faster than in C.

[iannucci, 2015/05/31 21:07:57]

This function is actually going away... I'm standardizing a way to add 'context
fields' to a Logger and let the actual Logger implementation handle it. I'll
post a CL for that todayish.

+                f("memlock(%s:%q): "+fmt, args...)
+        }
+
+        // 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
+        // timeout.
+        //
+        // release will attempt to CAS the item to remove it's contents (clientID).
+        // another lock observing an empty clientID will know that the lock is
+        // obtainable.
+        checkAnd := func(op checkOp) bool {
+                itm, err := mc.Get(key)
+                if err != nil {
+                        logf(log.Warningf, "error getting: %s", err)
+                        return false
+                }
+
+                if len(itm.Value) > 0 && !bytes.Equal(itm.Value, cid) {
+                        logf(log.Infof, "lock owned by %q", string(itm.Value))
+                        return false
+                }
+
+                if op == refresh {
+                        itm.Value = cid
+                        itm.Expiration = memcacheLockTime
+                } else {
+                        if len(itm.Value) == 0 {
+                                // it's already unlocked, no need to CAS
+                                logf(log.Infof, "lock already released")
+                                return true
+                        }
+                        itm.Value = []byte{}
+                        itm.Expiration = delay
+                }
+
+                err = mc.CompareAndSwap(itm)
+                if err != nil {
+                        logf(log.Warningf, "failed to %s lock: %q", op, err)
+                        return false
+                }
+
+                return true
+        }
+
+        // Now the actual logic begins. First we 'Add' the item, which will set it if
+        // it's not present in the memcache, otherwise leaves it alone.
+        err := mc.Add(&memcache.Item{
+                Key: key, Value: cid, Expiration: memcacheLockTime})
+        if err != nil {
+                if err != memcache.ErrNotStored {
+                        logf(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)
+
+        defer func() {
+                close(stopChan)
+                <-stoppedChan // this blocks TryWithLock until the goroutine below quits.
+        }()
+
+        // 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.).
+        go func() {
+                defer close(stoppedChan)
+
+        checkLoop:
+                for {
+                        select {
+                        case <-stopChan:
+                                break checkLoop
+                        case <-time.After(delay):
+                        }
+                        if !checkAnd(refresh) {
+                                atomic.StoreUint32(&held, 0)
+                                logf(log.Warningf, "lost lock: %s", err)
+                                break
+                        }
+                }
+
+                checkAnd(release)
+                atomic.StoreUint32(&held, 0)
+        }()
+
+        return f(func() bool { return atomic.LoadUint32(&held) == 1 })
+}