A transparent cache for datastore, backed by memcache.

filter/dscache/dscache.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 dscache
+
+import (
+        "bytes"
+        "compress/zlib"
+        "crypto/sha1"
+        "encoding/base64"
+        "fmt"
+        "io/ioutil"
+        "math/rand"
+        "reflect"
+        "time"
+
+        "github.com/luci/gae/service/datastore"
+        "github.com/luci/gae/service/memcache"
+)
+
+var (
+        // LockTimeSeconds is the number of seconds that a "lock" memcache entry will
+        // have its expiration set to. It's set to just over half of the frontend
+        // request handler timeout (currently 60 seconds).
+        LockTimeSeconds = 31
+
+        // CacheTimeSeconds is the default number of seconds that a positively cached
+        // entity will be retained (memcache contention notwithstanding). A value of
+        // 0 is infinite.
+        CacheTimeSeconds int64 = 0
+
+        // CompressionThreshold is the number of bytes of entity value after which
+        // compression kicks in.
+        CompressionThreshold = 860
+
+        // DefaultShards is the default number of key sharding to do.
+        DefaultShards int = 1
+
+        // DefaultEnable indicates whether or not caching is globally enabled or
+        // disabled by default. Can still be overridden by CacheEnableMeta.
+        DefaultEnabled = true
+
+        // ShardsForKey is a user-controllable global function which calculates
+        // the number of shards to use for a certain datastore key. The provided
+        // key will always be non-nil, complete and valid.
+        //
+        // The # shards returned may be between 1 and 256. Values above this range
+        // will be clamped into that range. A return value of 0 means that NO cache
+        // operations should be done for this key, regardless of the dscache.enable
+        // setting.
+        //
+        // If ShardsForKey is nil, the value of DefaultShards is used.
+        ShardsForKey func(datastore.Key) int
+)
+
+const (
+        MemcacheVersion = "1"
+
+        // KeyFormat is the format string used to generate memcache keys. It's
+        //   gae:<version>:<shard#>:<base64_std_nopad(sha1(datastore.Key))>
+        KeyFormat      = "gae:" + MemcacheVersion + ":%x:%s"
+        Sha1B64Padding = 1
+        Sha1B64Size    = 28 - Sha1B64Padding
+
+        MaxShards          = 256
+        MaxShardsLen       = len("ff")
+        InternalGAEPadding = 96
+        ValueSizeLimit     = (1000 * 1000) - InternalGAEPadding - MaxShardsLen
+
+        CacheEnableMeta     = "dscache.enable"
+        CacheExpirationMeta = "dscache.expiration"
+
+        // NonceBytes is the number of bytes to use in the 'lock' nonce. It must be
+        // a multiple of 4.
+        NonceBytes = 8
+)
+
+type CompressionType byte
+
+const (
+        NoCompression CompressionType = iota
+        ZlibCompression
+)
+
+func (c CompressionType) String() string {
+        switch c {
+        case NoCompression:
+                return "NoCompression"
+        case ZlibCompression:
+                return "ZlibCompression"
+        default:
+                return fmt.Sprintf("UNKNOWN_CompressionType(%d)", c)
+        }
+}
+
+// FlagValue is used to indicate if a memcache entry currently contains an
+// item or a lock.
+type FlagValue uint32
+
+const (
+        ItemHasData FlagValue = iota
+        ItemHasLock
+)
+
+func meta(val datastore.PropertyMap, key string, dflt interface{}) interface{} {
+        ret, err := val.GetMeta(key)
+        if err == datastore.ErrMetaFieldUnset {
+                return dflt
+        }
+        if reflect.TypeOf(dflt) != reflect.TypeOf(ret) {
+                // TODO(riannucci): Panic? Log a warning?
+                return dflt
+        }
+        return ret
+}
+
+func numShards(k datastore.Key) int {
+        if datastore.KeyIncomplete(k) {
+                return 0
+        }
+        ret := DefaultShards
+        if ShardsForKey != nil {
+                ret = ShardsForKey(k)
+        }
+        if ret < 1 {
+                return 0 // disable caching entirely
+        }
+        if ret > MaxShards {
+                ret = MaxShards
+        }
+        return ret
+}
+
+func mkKeySuffix(k datastore.Key) string {
+        buf := bytes.Buffer{}
+        if err := datastore.WriteKey(&buf, datastore.WithoutContext, k); err != nil {
+                panic(err) // can't happen, since we're using a byte buffer.
+        }
+        dgst := sha1.Sum(buf.Bytes())
+        buf.Reset()
+        enc := base64.NewEncoder(base64.StdEncoding, &buf)
+        _, err := enc.Write(dgst[:])
+        enc.Close()
+        if err != nil {
+                panic(err) // can't happen, since we're using a byte buffer.
+        }
+        return buf.String()[:buf.Len()-Sha1B64Padding]
+}
+
+func mkRandKeys(mr *rand.Rand, keys []datastore.Key) []string {
+        ret := make([]string, len(keys))
+        for i, key := range keys {
+                shards := numShards(key)
+                if shards == 0 {
+                        continue
+                }
+                ret[i] = fmt.Sprintf(KeyFormat, mr.Intn(shards), mkKeySuffix(key))
+        }
+        return ret
+}
+
+func mkAllKeys(keys []datastore.Key) []string {
+        size := 0
+        nums := make([]int, len(keys))
+        for i, k := range keys {
+                shards := numShards(k)
+                nums[i] = shards
+                size += shards
+        }
+        if size == 0 {
+                return nil
+        }
+        ret := make([]string, 0, size)
+        for i, key := range keys {
+                keySuffix := mkKeySuffix(key)
+                for shard := 0; shard < nums[i]; shard++ {
+                        ret = append(ret, fmt.Sprintf(KeyFormat, shard, keySuffix))
+                }
+        }
+        return ret
+}
+
+func init() {
+        // TODO(riannucci): remove this when using crypto/rand instead of math/rand.
+        if NonceBytes%4 != 0 {
+                panic("NonceBytes must be a multiple of 4")
+        }
+}
+
+// crappyNonce creates a really crapyp nonce using math/rand. This is generally
+// unacceptable for cryptographic purposes, but since mathrand is the only
+// mocked randomness source, we use that.
+//
+// Do not use this function for anything other than mkRandLockItems or your hair
+// will fall out. You've been warned.
+//
+// TODO(riannucci): make this use crypto/rand instead
+func crappyNonce(mr *rand.Rand) []byte {
+        ret := make([]byte, NonceBytes)
+        for w := uint(0); w < NonceBytes/4; w++ {
+                word := mr.Uint32()
+                for i := uint(0); i < 4; i++ {
+                        ret[(w*4)+i] = byte(word >> (8 * i))
+                }
+        }
+        return ret
+}
+
+func mkRandLockItems(mc memcache.Interface, mr *rand.Rand, keys []datastore.Key) []memcache.Item {
+        mcKeys := mkRandKeys(mr, keys)
+        if mcKeys == nil {
+                return nil
+        }
+        nonce := crappyNonce(mr)
+        ret := make([]memcache.Item, len(mcKeys))
+        for i := range ret {
+                ret[i] = (mc.NewItem(mcKeys[i]).
+                        SetFlags(uint32(ItemHasLock)).
+                        SetExpiration(time.Second * time.Duration(LockTimeSeconds)).
+                        SetValue(nonce))
+        }
+        return ret
+}
+
+func mkAllLockItems(mc memcache.Interface, keys []datastore.Key) ([]memcache.Item, []string) {
+        mcKeys := mkAllKeys(keys)
+        if mcKeys == nil {
+                return nil, nil
+        }
+        ret := make([]memcache.Item, len(mcKeys))
+        for i := range ret {
+                ret[i] = (mc.NewItem(mcKeys[i]).
+                        SetFlags(uint32(ItemHasLock)).
+                        SetExpiration(time.Second * time.Duration(LockTimeSeconds)))
+        }
+        return ret, mcKeys
+}
+
+func mkItemData(pm datastore.PropertyMap) ([]byte, error) {
+        pm, _ = pm.Save(false)
+
+        data := []byte(nil)
+
+        if pm.DataLen() > 0 {
+                buf := bytes.Buffer{}
+                if err := buf.WriteByte(byte(NoCompression)); err != nil {
+                        panic(err) // can't happen, since we're using a byte buffer.
+                }
+                if err := pm.Write(&buf, datastore.WithoutContext); err != nil {
+                        panic(err) // can't happen, since we're using a byte buffer.
+                }
+
+                data = buf.Bytes()
+                if buf.Len() > CompressionThreshold {
+                        buf2 := bytes.NewBuffer(make([]byte, 0, len(data)))
+                        if err := buf2.WriteByte(byte(ZlibCompression)); err != nil {
+                                panic(err) // can't happen, since we're using a byte buffer.
+                        }
+                        writer := zlib.NewWriter(buf2)
+                        _, err := writer.Write(data[1:]) // skip the NoCompression byte
+                        writer.Close()
+                        if err != nil {
+                                panic(err) // can't happen, since we're using a byte buffer.
+                        }
+                        data = buf2.Bytes()
+                }
+                if len(data) > ValueSizeLimit {
+                        // TODO(riannucci): should we cache the fact that it's too big?
+                        data = nil // disables memcaching
+                }
+        }
+
+        return data, nil
+}
+
+func decodeValue(val []byte, ns, aid string) (datastore.PropertyMap, error) {
+        if len(val) == 0 {
+                return nil, datastore.ErrNoSuchEntity
+        }
+        buf := bytes.NewBuffer(val)
+        compTypeByte, err := buf.ReadByte()
+        if err != nil {
+                panic(err) // can't happen, since we're using a byte buffer.
+        }
+        if CompressionType(compTypeByte) == ZlibCompression {
+                reader, err := zlib.NewReader(buf)
+                if err != nil {
+                        return nil, err
+                }
+                defer reader.Close()
+                data, err := ioutil.ReadAll(reader)
+                if err != nil {
+                        return nil, err
+                }
+                buf = bytes.NewBuffer(data)
+        }
+        ret := datastore.PropertyMap{}
+        err = ret.Read(buf, datastore.WithoutContext, ns, aid)
+        return ret, err
+}
+
+// TODO(riannucci): Should there be a way to purge the cache entries for a range
+// of datastore keys?