impl/cloud: Add support for "memcached" memcache.

impl/cloud/memcache.go

+// Copyright 2016 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package cloud
+
+import (
+        "crypto/sha256"
+        "encoding/hex"
+        "strconv"
+        "time"
+
+        "github.com/luci/gae/service/info"
+        mc "github.com/luci/gae/service/memcache"
+
+        "github.com/bradfitz/gomemcache/memcache"
+        "golang.org/x/net/context"
+)
+
+const (
+        // memcacheKeyPrefix is the common prefix prepended to memcached keys created
+        // by this package. It is intended to ensure that keys do not conflict with
+        // other users of the service.
+        memcacheKeyPrefix = "github.com/luci/gae/impl/cloud:"
+
+        // keyHashSizeThreshold is a threshold for key hashing. If the key's length
+        // exceeds this threshold, the key will be hashed and the hash used in its
+        // place with the memcache service.
+        //
+        // This is an implementation detail, but will be visible to the user in the
+        // Key field of the memcache entry on retrieval.
+        keyHashSizeThreshold = 250
+)
+
+// memcacheClient is a "service/memcache" implementation built on top of a
+// "memcached" client connection.
+//
+// Because "memcached" has no concept of a namespace, we differentiate memcache
+// entries by prepending "memcacheKeyPrefix:SHA256(namespace):" to each key.
+type memcacheClient struct {
+        client *memcache.Client
+}
+
+func (m *memcacheClient) use(c context.Context) context.Context {
+        return mc.SetRawFactory(c, func(ic context.Context) mc.RawInterface {
+                return bindMemcacheClient(m, info.GetNamespace(ic))
+        })
+}
+
+type memcacheItem struct {
+        native *memcache.Item
+}
+
+func (it *memcacheItem) Key() string   { return it.native.Key }
+func (it *memcacheItem) Value() []byte { return it.native.Value }
+func (it *memcacheItem) Flags() uint32 { return it.native.Flags }
+func (it *memcacheItem) Expiration() time.Duration {
+        return time.Duration(it.native.Expiration) * time.Second
+}
+
+func (it *memcacheItem) SetKey(v string) mc.Item {
+        it.native.Key = v
+        return it
+}
+
+func (it *memcacheItem) SetValue(v []byte) mc.Item {
+        it.native.Value = v
+        return it
+}
+
+func (it *memcacheItem) SetFlags(v uint32) mc.Item {
+        it.native.Flags = v
+        return it
+}
+
+func (it *memcacheItem) SetExpiration(v time.Duration) mc.Item {
+        it.native.Expiration = int32(v.Seconds())
+        return it
+}
+
+func (it *memcacheItem) SetAll(other mc.Item) {
+        origKey := it.native.Key
+
+        var on memcache.Item
+        if other != nil {
+                on = *(other.(*memcacheItem).native)
+        }
+        it.native = &on
+        it.native.Key = origKey
+}
+
+func hashBytes(b []byte) string {
+        hash := sha256.Sum256(b)
+        return hex.EncodeToString(hash[:])
+}
+
+type boundMemcacheClient struct {
+        *memcacheClient
+        keyPrefix string
+}
+
+func bindMemcacheClient(mc *memcacheClient, ns string) *boundMemcacheClient {
+        nsPrefix := hashBytes([]byte(ns))
+        return &boundMemcacheClient{
+                memcacheClient: mc,
+                keyPrefix:      memcacheKeyPrefix + nsPrefix + ":",
+        }
+}
+
+func (*boundMemcacheClient) newMemcacheItem(nativeKey string) *memcacheItem {
+        return &memcacheItem{
+                native: &memcache.Item{
+                        Key: nativeKey,
+                },
+        }
+}
+
+// makeKey constructs the actual key used with the memcache service. This
+// includes a service-specific prefix, the key's namespace, and the key itself.
+// If the key's length exceeds the keyHashSizeThreshold, the key will be stored
+// as a hash.
+func (bmc *boundMemcacheClient) makeKey(base string) string {
+        if len(base) > keyHashSizeThreshold {
+                base = hashBytes([]byte(base))
+        }
+        return bmc.keyPrefix + base
+}
+
+func (bmc *boundMemcacheClient) userKey(key string) string { return key[len(bmc.keyPrefix):] }
+
+func (bmc *boundMemcacheClient) nativeItem(itm mc.Item) *memcache.Item {
+        ni := *(itm.(*memcacheItem).native)
+        ni.Key = bmc.makeKey(ni.Key)
+        return &ni
+}
+
+func (bmc *boundMemcacheClient) NewItem(key string) mc.Item { return bmc.newMemcacheItem(key) }
+
+func (bmc *boundMemcacheClient) AddMulti(items []mc.Item, cb mc.RawCB) error {
+        for _, itm := range items {
+                err := bmc.client.Add(bmc.nativeItem(itm))
+                cb(bmc.translateErr(err))
+        }
+        return nil
+}
+
+func (bmc *boundMemcacheClient) SetMulti(items []mc.Item, cb mc.RawCB) error {
+        for _, itm := range items {
+                err := bmc.client.Set(bmc.nativeItem(itm))
+                cb(bmc.translateErr(err))
+        }
+        return nil
+}
+
+func (bmc *boundMemcacheClient) GetMulti(keys []string, cb mc.RawItemCB) error {
+        nativeKeys := make([]string, len(keys))
+        for i, key := range keys {
+                nativeKeys[i] = bmc.makeKey(key)
+        }
+
+        itemMap, err := bmc.client.GetMulti(nativeKeys)
+        if err != nil {
+                return bmc.translateErr(err)
+        }
+
+        // Translate the item keys back to user keys.
+        for _, v := range itemMap {
+                v.Key = bmc.userKey(v.Key)
+        }
+
+        for _, k := range nativeKeys {
+                if it := itemMap[k]; it != nil {
+                        cb(&memcacheItem{native: it}, nil)
+                } else {
+                        cb(nil, mc.ErrCacheMiss)
+                }
+        }
+        return nil
+}
+
+func (bmc *boundMemcacheClient) DeleteMulti(keys []string, cb mc.RawCB) error {
+        for _, k := range keys {
+                err := bmc.client.Delete(bmc.makeKey(k))
+                cb(bmc.translateErr(err))
+        }
+        return nil
+}
+
+func (bmc *boundMemcacheClient) CompareAndSwapMulti(items []mc.Item, cb mc.RawCB) error {
+        for _, itm := range items {
+                err := bmc.client.CompareAndSwap(bmc.nativeItem(itm))
+                cb(bmc.translateErr(err))
+        }
+        return nil
+}
+
+func (bmc *boundMemcacheClient) Increment(key string, delta int64, initialValue *uint64) (uint64, error) {
+        // key is now the native key (namespaced).
+        key = bmc.makeKey(key)
+
+        op := func() (newValue uint64, err error) {
+                switch {
+                case delta > 0:
+                        newValue, err = bmc.client.Increment(key, uint64(delta))
+                case delta < 0:
+                        newValue, err = bmc.client.Decrement(key, uint64(-delta))
+                default:
+                        // We don't want to change the value, but we want to return ErrNotStored
+                        // if the value doesn't exist. Use Get.
+                        _, err = bmc.client.Get(key)
+                }
+                err = bmc.translateErr(err)
+                return
+        }
+
+        if initialValue == nil {
+                return op()
+        }
+
+        // The Memcache service doesn't have an "IncrementExisting" equivalent. We
+        // will emulate this with other memcache operations, using Add to set the
+        // initial value if appropriate.
+        var (
+                itm *memcacheItem
+                iv  = *initialValue
+        )
+        for {
+                // Perform compare-and-swap.
+                nv, err := op()
+                if err != mc.ErrCacheMiss {
+                        return nv, err
+                }
+
+                // The value doesn't exist. Use "Add" to set the initial value. We will
+                // calculate the "initial value" as if delta were applied so we can do this
+                // in one operation.
+                //
+                // We only need to do this once per invocation, so we will use "itm == nil"
+                // as a sentinel for uninitialized.
+                if itm == nil {
+                        // Overflow wraps around (to zero), and underflow is capped at 0.
+                        if delta < 0 {
+                                udelta := uint64(-delta)
+                                if udelta >= iv {
+                                        // Would underflow, cap at 0.
+                                        iv = 0
+                                } else {
+                                        iv -= udelta
+                                }
+                        } else {
+                                // Apply delta. This will automatically wrap on overflow.
+                                iv += uint64(delta)
+                        }
+
+                        itm = bmc.newMemcacheItem(key)
+                        itm.SetValue([]byte(strconv.FormatUint(iv, 10)))
+                }
+                switch err := bmc.client.Add(itm.native); err {
+                case nil:
+                        // Item was successfully set.
+                        return iv, nil
+
+                case mc.ErrNotStored:
+                        // Something else set it in between "op" and "Add". Try "op" again.
+                        break
+
+                default:
+                        return 0, err
+                }
+        }
+}
+
+func (bmc *boundMemcacheClient) Flush() error {
+        // Unfortunately there's not really a good way to flush just a single
+        // namespace, so Flush will flush all memcache.
+        return bmc.translateErr(bmc.client.FlushAll())
+}
+
+func (bmc *boundMemcacheClient) Stats() (*mc.Statistics, error) { return nil, mc.ErrNoStats }
+
+func (*boundMemcacheClient) translateErr(err error) error {
+        switch err {
+        case memcache.ErrCacheMiss:
+                return mc.ErrCacheMiss
+        case memcache.ErrCASConflict:
+                return mc.ErrCASConflict
+        case memcache.ErrNotStored:
+                return mc.ErrNotStored
+        case memcache.ErrServerError:
+                return mc.ErrServerError
+        case memcache.ErrNoStats:
+                return mc.ErrNoStats
+        default:
+                return err
+        }
+}