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
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 com
iannucci
2015/03/07 02:27:52
haha :P. It used to be `time.Duration(numSecs) * t
|
+ |
+// 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. TryLo
Vadim Sh.
2015/03/06 03:46:30
For completness you can also implement "Lock()" th
iannucci
2015/03/07 02:27:52
SGTM
|
+ 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.ke
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 compl
iannucci
2015/03/07 02:27:52
OK, done. I should have done that from the beginni
|
+ 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" impli
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".
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 b
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 min
iannucci
2015/03/07 02:27:52
Yeah, it's a really weird API. AND! The Item objec
|
+ 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 an
iannucci
2015/03/07 02:27:52
shrug. It's true that it failed to lock though. At
|
+ 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 o
iannucci
2015/03/07 02:27:52
it's usually not though. Maybe %q, then we'll see
|
+ return true |
+} |