A transparent cache for datastore, backed by memcache.

filter/dscache/support.go

Index: filter/dscache/support.go
diff --git a/filter/dscache/support.go b/filter/dscache/support.go
new file mode 100644
index 0000000000000000000000000000000000000000..f57a9e9df2408c3dfca4353daffdcd043652a0d2
--- /dev/null
+++ b/filter/dscache/support.go
@@ -0,0 +1,158 @@
+// 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 dscache
+
+import (
+	"fmt"
+	"math/rand"
+	"time"
+
+	ds "github.com/luci/gae/service/datastore"
+	"github.com/luci/gae/service/memcache"
+	log "github.com/luci/luci-go/common/logging"
+	"golang.org/x/net/context"
+)
+
+type supportContext struct {
+	aid string
+	ns  string
+
+	c            context.Context
+	mc           memcache.Interface
+	mr           *rand.Rand
+	shardsForKey func(ds.Key) int
+}
+
+func (s *supportContext) numShards(k ds.Key) int {
+	ret := DefaultShards
+	if s.shardsForKey != nil {
+		ret = s.shardsForKey(k)
+	}
+	if ret < 1 {
+		return 0 // disable caching entirely
+	}
+	if ret > MaxShards {
+		ret = MaxShards
+	}
+	return ret
+}
+
+func (s *supportContext) mkRandKeys(keys []ds.Key, metas ds.MultiMetaGetter) []string {
+	ret := []string(nil)
+	for i, key := range keys {
+		if !metas.GetMetaDefault(i, CacheEnableMeta, true).(bool) {
+			continue
+		}
+		shards := s.numShards(key)
+		if shards == 0 {
+			continue
+		}
+		if ret == nil {
+			ret = make([]string, len(keys))
+		}
+		ret[i] = MakeMemcacheKey(s.mr.Intn(shards), key)
+	}
+	return ret
+}
+
+func (s *supportContext) mkAllKeys(keys []ds.Key) []string {
+	size := 0
+	nums := make([]int, len(keys))
+	for i, k := range keys {
+		if !ds.KeyIncomplete(k) {
+			shards := s.numShards(k)
+			nums[i] = shards
+			size += shards
+		}
+	}
+	if size == 0 {
+		return nil
+	}
+	ret := make([]string, 0, size)
+	for i, key := range keys {
+		if !ds.KeyIncomplete(key) {
+			keySuffix := HashKey(key)
+			for shard := 0; shard < nums[i]; shard++ {
+				ret = append(ret, fmt.Sprintf(KeyFormat, shard, keySuffix))
+			}
+		}
+	}
+	return ret
+}
+
+// crappyNonce creates a really crappy nonce using math/rand. This is generally
+// unacceptable for cryptographic purposes, but since mathrand is the only
+// mocked randomness source, we use that.
+//
+// The random values here are controlled entriely by the application, will never
+// be shown to, or provided by, the user, so this should be fine.
+//
+// Do not use this function for anything other than mkRandLockItems or your hair
+// will fall out. You've been warned.
+func (s *supportContext) crappyNonce() []byte {
+	ret := make([]byte, NonceUint32s*4)
+	for w := uint(0); w < NonceUint32s; w++ {
+		word := s.mr.Uint32()
+		for i := uint(0); i < 4; i++ {
+			ret[(w*4)+i] = byte(word >> (8 * i))
+		}
+	}
+	return ret
+}
+
+func (s *supportContext) mutation(keys []ds.Key, f func() error) error {
+	lockItems, lockKeys := s.mkAllLockItems(keys)
+	if lockItems == nil {
+		return f()
+	}
+	if err := s.mc.SetMulti(lockItems); err != nil {
+		// this is a hard failure. No mutation can occur if we're unable to set
+		// locks out. See "DANGER ZONE" in the docs.
+		(log.Fields{log.ErrorKey: err}).Errorf(
+			s.c, "dscache: HARD FAILURE: supportContext.mutation(): mc.SetMulti")
+		return err
+	}
+	err := f()
+	if err == nil {
+		if err := s.mc.DeleteMulti(lockKeys); err != nil {
+			(log.Fields{log.ErrorKey: err}).Warningf(
+				s.c, "dscache: mc.DeleteMulti")
+		}
+	}
+	return err
+}
+
+func (s *supportContext) mkRandLockItems(keys []ds.Key, metas ds.MultiMetaGetter) ([]memcache.Item, []byte) {
+	mcKeys := s.mkRandKeys(keys, metas)
+	if len(mcKeys) == 0 {
+		return nil, nil
+	}
+	nonce := s.crappyNonce()
+	ret := make([]memcache.Item, len(mcKeys))
+	for i, k := range mcKeys {
+		if k == "" {
+			continue
+		}
+		ret[i] = (s.mc.NewItem(k).
+			SetFlags(uint32(ItemHasLock)).
+			SetExpiration(time.Second * time.Duration(LockTimeSeconds)).
+			SetValue(nonce))
+	}
+	return ret, nonce
+}
+
+func (s *supportContext) mkAllLockItems(keys []ds.Key) ([]memcache.Item, []string) {
+	mcKeys := s.mkAllKeys(keys)
+	if mcKeys == nil {
+		return nil, nil
+	}
+	ret := make([]memcache.Item, len(mcKeys))
+	for i := range ret {
+		ret[i] = (s.mc.NewItem(mcKeys[i]).
+			SetFlags(uint32(ItemHasLock)).
+			SetExpiration(time.Second * time.Duration(LockTimeSeconds)))
+	}
+	return ret, mcKeys
+}