A simple memcache lock for appengine.

go/src/infra/gae/libs/mlock/mlock.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 mlock
+
+import (
+        "bytes"
+        "crypto/rand"
+        "fmt"
+        "sync/atomic"
+        "time"
+
+        "appengine"
+        "appengine/memcache"
+)
+
+var (
+        // ErrUnableToLock we failed to obtain the lock.
+        ErrUnableToLock = fmt.Errorf("mlock: unable to lock")

[Vadim Sh., 2015/03/06 03:39:26]

mlock is also syscall... maybe pick some different name? memlock

[iannucci, 2015/03/07 02:27:52]

done

+)
+
+// MemcacheLockTime is the expiration time of the memcache entry. If the lock
+// is correctly released, then it will be released before this time. It is set
+// to 64 seconds.
+const MemcacheLockTime = (time.Second << 6) // 64 seconds

[Vadim Sh., 2015/03/06 03:39:26]

omg, why not "64 * time.Second" and remove two comments about 64 seconds?

also it seems way too much, considering you ping the lock each second. Though
stupid RPCs sometimes hang for extended periods of time.

[iannucci, 2015/03/07 02:27:52]

haha :P. It used to be `time.Duration(numSecs) * time.Second` which was
annoying. And now it's a constant, so it doesn't matter.
Yeah, hanging and also goroutine scheduling starvation. Maybe I'll make it 16
secs?

+
+// MemcacheLock 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.
+//
+// Construct a new MemcacheLock with the New() method.
+type MemcacheLock struct {
+        ctx      appengine.Context
+        key      string
+        clientID []byte
+
+        running uint32
+}
+
+// New creates a MemcacheLock. `key` is the memcache key to use. 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 New() will generate a random
+// one.
+//
+// Using a deterministid clientID (like the taskqueue task name, for example) has
+// the benefit that on an error the re-tried taskqueue handler may be able to
+// re-obtain the lock, assuming it wasn't released properly.
+func New(c appengine.Context, key, clientID string) (*MemcacheLock, error) {
+        clientIDbytes := []byte(clientID)
+        if len(clientIDbytes) == 0 {
+                clientIDbytes = make([]byte, 32)
+                _, err := rand.Read(clientIDbytes)
+                if err != nil {
+                        return nil, err
+                }
+                c.Debugf("mlock: generated clientId: %v", clientIDbytes)
+        }
+        return &MemcacheLock{ctx: c, key: key, clientID: clientIDbytes}, nil
+}
+
+// Obtain obtains the lock. Returns true iff we were able to obtain the lock.
+func (m *MemcacheLock) Obtain() bool {

[Vadim Sh., 2015/03/06 03:39:26]

nit: TryLock() is more canonical name, imho. TryLock()\Unlock().

Obtain()\Release() is more appropriate for mutex (you obtain a mutex, but you
lock a lock). Anyhow, 'sync' packages uses Lock()\Unlock().

[Vadim Sh., 2015/03/06 03:46:30]

For completness you can also implement "Lock()" that busy waits on "TryLock()"
until locked. Then your lock will implement sync.Locker :)

[iannucci, 2015/03/07 02:27:52]

SGTM
Yeah, ok, thought about this. I'll do it.

+        err := memcache.Add(m.ctx, &memcache.Item{
+                Key: m.key, Expiration: MemcacheLockTime, Value: m.clientID})

[Vadim Sh., 2015/03/06 03:39:26]

nit: namespace lock keys, e.g Key: "mlock:" + m.key

[iannucci, 2015/03/07 02:27:52]

done.

+        if err != nil {
+                if err != memcache.ErrNotStored {
+                        m.ctx.Warningf("Error adding: %s", err)
+                }
+                if !m.checkAndRefresh(true) {
+                        return false
+                }
+        }
+
+        atomic.StoreUint32(&m.running, 1)
+        go func() {
+                for m.Check() {
+                        time.Sleep(time.Second)

[Vadim Sh., 2015/03/06 03:39:26]

this can be a problem if all goroutines must complete before request handler
finishes. Can you rewrite it to use conditional variable or channel? So it
doesn't sit and wait for nothing 1 sec.

[iannucci, 2015/03/07 02:27:52]

OK, done. I should have done that from the beginning...

+                        if !m.checkAndRefresh(true) {
+                                m.ctx.Warningf("lost lock (%s:%s): %s", m.key, string(m.clientID), err)
+                                break
+                        }
+                }
+                m.checkAndRefresh(false)
+        }()
+
+        return true
+}
+
+// Check ensures that we still have control of the lock. Returns true if we

[Vadim Sh., 2015/03/06 03:39:26]

nit: Check returns true if we ...

"ensures" implies it actively does something

[iannucci, 2015/03/07 02:27:53]

good point

+// probably still have the lock :)
+func (m *MemcacheLock) Check() bool {
+        return atomic.LoadUint32(&m.running) == 1
+}
+
+// Release drops the lock (if we still own it, otherwise leaves it alone).
+func (m *MemcacheLock) Release() {
+        m.checkAndRefresh(false)
+        atomic.StoreUint32(&m.running, 0)
+}
+
+func (m *MemcacheLock) checkAndRefresh(refresh bool) bool {
+        itm, err := memcache.Get(m.ctx, m.key)
+        if err != nil {
+                m.ctx.Warningf("mlock: error getting: %s", err)
+                return false
+        }
+
+        if len(itm.Value) != 0 && !bytes.Equal(itm.Value, m.clientID) {
+                m.ctx.Infof("mlock: lock owned by %s", itm.Value)
+                return false
+        }
+
+        var op string
+        if refresh {
+                itm.Value = m.clientID
+                itm.Expiration = MemcacheLockTime
+                op = "refresh"

[Vadim Sh., 2015/03/06 03:39:26]

nit: move "op" to argument instead of "refresh".

checkAndRefresh("release") is more readable than checkAndRefresh(false)

[iannucci, 2015/03/07 02:27:52]

and then make it module-level constants? k :)

+        } else {
+                itm.Value = []byte{}

[Vadim Sh., 2015/03/06 03:39:26]

at this point you can compare itm.Value to empty byte buffer and return early if
it is, not attempting to swap it.

[iannucci, 2015/03/07 02:27:52]

Ah, sure. Done.

+                itm.Expiration = time.Second
+                op = "release"
+        }
+
+        err = memcache.CompareAndSwap(m.ctx, itm)

[Vadim Sh., 2015/03/06 03:39:26]

CompareAndSwap taking single argument blows my mind :)

[iannucci, 2015/03/07 02:27:52]

Yeah, it's a really weird API. AND! The Item object has hidden state, so not
only do you have to do a Get before, but you can only CAS with the SAME Item
that Get returned. :(

+        if err != nil {
+                m.ctx.Warningf("failed to %s lock (%s:%s): %s",

[Vadim Sh., 2015/03/06 03:39:26]

should a race condition on a lock (when compare and swap fails "compare" stage)
be logged as warning (it is now)?

[iannucci, 2015/03/07 02:27:52]

shrug. It's true that it failed to lock though. At this point we nominally hold
the lock.

+                        op, m.key, string(m.clientID), err)
+                return false
+        }
+
+        m.ctx.Infof("%sed lock (%s:%s)", op, m.key, string(m.clientID))

[Vadim Sh., 2015/03/06 03:39:26]

clientID is ugly byte string. Perhaps %v instead of %s?

[iannucci, 2015/03/07 02:27:52]

it's usually not though. Maybe %q, then we'll see escapes if it's really
non-utf8?

+        return true
+}