A simple memcache lock for appengine.

go/src/infra/gae/libs/mlock/mlock.go

Index: go/src/infra/gae/libs/mlock/mlock.go
diff --git a/go/src/infra/gae/libs/mlock/mlock.go b/go/src/infra/gae/libs/mlock/mlock.go
new file mode 100644
index 0000000000000000000000000000000000000000..5e0425a1da306dae5069d15c0260747d14ad792f
--- /dev/null
+++ b/go/src/infra/gae/libs/mlock/mlock.go
@@ -0,0 +1,135 @@
+// 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
+}