go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/impl/prod/memcache.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package prod
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"google.golang.org/appengine"
    22  	"google.golang.org/appengine/memcache"
    23  
    24  	mc "go.chromium.org/luci/gae/service/memcache"
    25  )
    26  
    27  // useMC adds a gae.Memcache implementation to context, accessible
    28  // by gae.GetMC(c)
    29  func useMC(c context.Context) context.Context {
    30  	return mc.SetRawFactory(c, func(ci context.Context) mc.RawInterface {
    31  		return mcImpl{getAEContext(ci)}
    32  	})
    33  }
    34  
    35  type mcImpl struct {
    36  	aeCtx context.Context
    37  }
    38  
    39  type mcItem struct {
    40  	i *memcache.Item
    41  }
    42  
    43  var _ mc.Item = mcItem{}
    44  
    45  func (i mcItem) Key() string               { return i.i.Key }
    46  func (i mcItem) Value() []byte             { return i.i.Value }
    47  func (i mcItem) Flags() uint32             { return i.i.Flags }
    48  func (i mcItem) Expiration() time.Duration { return i.i.Expiration }
    49  
    50  func (i mcItem) SetKey(k string) mc.Item {
    51  	i.i.Key = k
    52  	return i
    53  }
    54  func (i mcItem) SetValue(v []byte) mc.Item {
    55  	i.i.Value = v
    56  	return i
    57  }
    58  func (i mcItem) SetFlags(f uint32) mc.Item {
    59  	i.i.Flags = f
    60  	return i
    61  }
    62  func (i mcItem) SetExpiration(d time.Duration) mc.Item {
    63  	i.i.Expiration = d
    64  	return i
    65  }
    66  
    67  func (i mcItem) SetAll(other mc.Item) {
    68  	if other == nil {
    69  		i.i = &memcache.Item{Key: i.i.Key}
    70  	} else {
    71  		k := i.i.Key
    72  		*i.i = *other.(mcItem).i
    73  		i.i.Key = k
    74  	}
    75  }
    76  
    77  // mcF2R (MC fake-to-real) converts a mc.Item. i must originate from inside
    78  // this package for this function to work (see the panic message for why).
    79  //
    80  // If the item's Value == nil, it will be copied and replaced with []byte{}.
    81  func mcF2R(i mc.Item) *memcache.Item {
    82  	if mci, ok := i.(mcItem); ok {
    83  		if mci.i.Value == nil {
    84  			ret := *mci.i
    85  			ret.Value = []byte{}
    86  			return &ret
    87  		}
    88  		return mci.i
    89  	}
    90  	panic(
    91  		"you may not use other mc.Item implementations with this " +
    92  			"implementation of gae.Memcache, since it will cause all CompareAndSwap " +
    93  			"operations to fail. Please use the NewItem api instead.")
    94  }
    95  
    96  // mcMF2R (MC multi-fake-to-real) converts a slice of mc.Item to a slice of
    97  // *memcache.Item.
    98  func mcMF2R(items []mc.Item) []*memcache.Item {
    99  	realItems := make([]*memcache.Item, len(items))
   100  	for i, itm := range items {
   101  		realItems[i] = mcF2R(itm)
   102  	}
   103  	return realItems
   104  }
   105  
   106  func (m mcImpl) NewItem(key string) mc.Item {
   107  	return mcItem{&memcache.Item{Key: key}}
   108  }
   109  
   110  func doCB(err error, cb mc.RawCB) error {
   111  	if me, ok := err.(appengine.MultiError); ok {
   112  		for _, err := range me {
   113  			cb(err)
   114  		}
   115  		err = nil
   116  	}
   117  	return err
   118  }
   119  
   120  func (m mcImpl) DeleteMulti(keys []string, cb mc.RawCB) error {
   121  	return doCB(memcache.DeleteMulti(m.aeCtx, keys), cb)
   122  }
   123  
   124  func (m mcImpl) AddMulti(items []mc.Item, cb mc.RawCB) error {
   125  	return doCB(memcache.AddMulti(m.aeCtx, mcMF2R(items)), cb)
   126  }
   127  
   128  func (m mcImpl) SetMulti(items []mc.Item, cb mc.RawCB) error {
   129  	return doCB(memcache.SetMulti(m.aeCtx, mcMF2R(items)), cb)
   130  }
   131  
   132  func (m mcImpl) GetMulti(keys []string, cb mc.RawItemCB) error {
   133  	realItems, err := memcache.GetMulti(m.aeCtx, keys)
   134  	if err != nil {
   135  		return err
   136  	}
   137  	for _, k := range keys {
   138  		itm := realItems[k]
   139  		if itm == nil {
   140  			cb(nil, memcache.ErrCacheMiss)
   141  		} else {
   142  			cb(mcItem{itm}, nil)
   143  		}
   144  	}
   145  	return nil
   146  }
   147  
   148  func (m mcImpl) CompareAndSwapMulti(items []mc.Item, cb mc.RawCB) error {
   149  	return doCB(memcache.CompareAndSwapMulti(m.aeCtx, mcMF2R(items)), cb)
   150  }
   151  
   152  func (m mcImpl) Increment(key string, delta int64, initialValue *uint64) (uint64, error) {
   153  	if initialValue == nil {
   154  		return memcache.IncrementExisting(m.aeCtx, key, delta)
   155  	}
   156  	return memcache.Increment(m.aeCtx, key, delta, *initialValue)
   157  }
   158  
   159  func (m mcImpl) Flush() error {
   160  	return memcache.Flush(m.aeCtx)
   161  }
   162  
   163  func (m mcImpl) Stats() (*mc.Statistics, error) {
   164  	stats, err := memcache.Stats(m.aeCtx)
   165  	if err != nil {
   166  		return nil, err
   167  	}
   168  	return (*mc.Statistics)(stats), nil
   169  }