services/keep-web/cache.go

   1 package main
   2
   3 import (
   4         "sync"
   5         "sync/atomic"
   6         "time"
   7
   8         "git.curoverse.com/arvados.git/sdk/go/arvados"
   9         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
  10         "github.com/hashicorp/golang-lru"
  11 )
  12
  13 type cache struct {
  14         TTL                  arvados.Duration
  15         MaxCollectionEntries int
  16         MaxCollectionBytes   int64
  17         MaxPermissionEntries int
  18         MaxUUIDEntries       int
  19
  20         stats       cacheStats
  21         pdhs        *lru.TwoQueueCache
  22         collections *lru.TwoQueueCache
  23         permissions *lru.TwoQueueCache
  24         setupOnce   sync.Once
  25 }
  26
  27 type cacheStats struct {
  28         Requests          uint64 `json:"Cache.Requests"`
  29         CollectionBytes   uint64 `json:"Cache.CollectionBytes"`
  30         CollectionEntries int    `json:"Cache.CollectionEntries"`
  31         CollectionHits    uint64 `json:"Cache.CollectionHits"`
  32         PDHHits           uint64 `json:"Cache.UUIDHits"`
  33         PermissionHits    uint64 `json:"Cache.PermissionHits"`
  34         APICalls          uint64 `json:"Cache.APICalls"`
  35 }
  36
  37 type cachedPDH struct {
  38         expire time.Time
  39         pdh    string
  40 }
  41
  42 type cachedCollection struct {
  43         expire     time.Time
  44         collection *arvados.Collection
  45 }
  46
  47 type cachedPermission struct {
  48         expire time.Time
  49 }
  50
  51 func (c *cache) setup() {
  52         var err error
  53         c.pdhs, err = lru.New2Q(c.MaxUUIDEntries)
  54         if err != nil {
  55                 panic(err)
  56         }
  57         c.collections, err = lru.New2Q(c.MaxCollectionEntries)
  58         if err != nil {
  59                 panic(err)
  60         }
  61         c.permissions, err = lru.New2Q(c.MaxPermissionEntries)
  62         if err != nil {
  63                 panic(err)
  64         }
  65 }
  66
  67 var selectPDH = map[string]interface{}{
  68         "select": []string{"portable_data_hash"},
  69 }
  70
  71 func (c *cache) Stats() cacheStats {
  72         c.setupOnce.Do(c.setup)
  73         return cacheStats{
  74                 Requests:          atomic.LoadUint64(&c.stats.Requests),
  75                 CollectionBytes:   c.collectionBytes(),
  76                 CollectionEntries: c.collections.Len(),
  77                 CollectionHits:    atomic.LoadUint64(&c.stats.CollectionHits),
  78                 PDHHits:           atomic.LoadUint64(&c.stats.PDHHits),
  79                 PermissionHits:    atomic.LoadUint64(&c.stats.PermissionHits),
  80                 APICalls:          atomic.LoadUint64(&c.stats.APICalls),
  81         }
  82 }
  83
  84 func (c *cache) Get(arv *arvadosclient.ArvadosClient, targetID string, forceReload bool) (*arvados.Collection, error) {
  85         c.setupOnce.Do(c.setup)
  86
  87         atomic.AddUint64(&c.stats.Requests, 1)
  88
  89         permOK := false
  90         permKey := arv.ApiToken + "\000" + targetID
  91         if forceReload {
  92         } else if ent, cached := c.permissions.Get(permKey); cached {
  93                 ent := ent.(*cachedPermission)
  94                 if ent.expire.Before(time.Now()) {
  95                         c.permissions.Remove(permKey)
  96                 } else {
  97                         permOK = true
  98                         atomic.AddUint64(&c.stats.PermissionHits, 1)
  99                 }
 100         }
 101
 102         var pdh string
 103         if arvadosclient.PDHMatch(targetID) {
 104                 pdh = targetID
 105         } else if ent, cached := c.pdhs.Get(targetID); cached {
 106                 ent := ent.(*cachedPDH)
 107                 if ent.expire.Before(time.Now()) {
 108                         c.pdhs.Remove(targetID)
 109                 } else {
 110                         pdh = ent.pdh
 111                         atomic.AddUint64(&c.stats.PDHHits, 1)
 112                 }
 113         }
 114
 115         var collection *arvados.Collection
 116         if pdh != "" {
 117                 collection = c.lookupCollection(pdh)
 118         }
 119
 120         if collection != nil && permOK {
 121                 return collection, nil
 122         } else if collection != nil {
 123                 // Ask API for current PDH for this targetID. Most
 124                 // likely, the cached PDH is still correct; if so,
 125                 // _and_ the current token has permission, we can
 126                 // use our cached manifest.
 127                 atomic.AddUint64(&c.stats.APICalls, 1)
 128                 var current arvados.Collection
 129                 err := arv.Get("collections", targetID, selectPDH, &current)
 130                 if err != nil {
 131                         return nil, err
 132                 }
 133                 if current.PortableDataHash == pdh {
 134                         exp := time.Now().Add(time.Duration(c.TTL))
 135                         c.permissions.Add(permKey, &cachedPermission{
 136                                 expire: exp,
 137                         })
 138                         if pdh != targetID {
 139                                 c.pdhs.Add(targetID, &cachedPDH{
 140                                         expire: exp,
 141                                         pdh:    pdh,
 142                                 })
 143                         }
 144                         return collection, err
 145                 } else {
 146                         // PDH changed, but now we know we have
 147                         // permission -- and maybe we already have the
 148                         // new PDH in the cache.
 149                         if coll := c.lookupCollection(current.PortableDataHash); coll != nil {
 150                                 return coll, nil
 151                         }
 152                 }
 153         }
 154
 155         // Collection manifest is not cached.
 156         atomic.AddUint64(&c.stats.APICalls, 1)
 157         err := arv.Get("collections", targetID, nil, &collection)
 158         if err != nil {
 159                 return nil, err
 160         }
 161         exp := time.Now().Add(time.Duration(c.TTL))
 162         c.permissions.Add(permKey, &cachedPermission{
 163                 expire: exp,
 164         })
 165         c.pdhs.Add(targetID, &cachedPDH{
 166                 expire: exp,
 167                 pdh:    collection.PortableDataHash,
 168         })
 169         c.collections.Add(collection.PortableDataHash, &cachedCollection{
 170                 expire:     exp,
 171                 collection: collection,
 172         })
 173         if int64(len(collection.ManifestText)) > c.MaxCollectionBytes/int64(c.MaxCollectionEntries) {
 174                 go c.pruneCollections()
 175         }
 176         return collection, nil
 177 }
 178
 179 // pruneCollections checks the total bytes occupied by manifest_text
 180 // in the collection cache and removes old entries as needed to bring
 181 // the total size down to CollectionBytes. It also deletes all expired
 182 // entries.
 183 //
 184 // pruneCollections does not aim to be perfectly correct when there is
 185 // concurrent cache activity.
 186 func (c *cache) pruneCollections() {
 187         var size int64
 188         now := time.Now()
 189         keys := c.collections.Keys()
 190         entsize := make([]int, len(keys))
 191         expired := make([]bool, len(keys))
 192         for i, k := range keys {
 193                 v, ok := c.collections.Peek(k)
 194                 if !ok {
 195                         continue
 196                 }
 197                 ent := v.(*cachedCollection)
 198                 n := len(ent.collection.ManifestText)
 199                 size += int64(n)
 200                 entsize[i] = n
 201                 expired[i] = ent.expire.Before(now)
 202         }
 203         for i, k := range keys {
 204                 if expired[i] {
 205                         c.collections.Remove(k)
 206                         size -= int64(entsize[i])
 207                 }
 208         }
 209         for i, k := range keys {
 210                 if size <= c.MaxCollectionBytes {
 211                         break
 212                 }
 213                 if expired[i] {
 214                         // already removed this entry in the previous loop
 215                         continue
 216                 }
 217                 c.collections.Remove(k)
 218                 size -= int64(entsize[i])
 219         }
 220 }
 221
 222 // collectionBytes returns the approximate memory size of the
 223 // collection cache.
 224 func (c *cache) collectionBytes() uint64 {
 225         var size uint64
 226         for _, k := range c.collections.Keys() {
 227                 v, ok := c.collections.Peek(k)
 228                 if !ok {
 229                         continue
 230                 }
 231                 size += uint64(len(v.(*cachedCollection).collection.ManifestText))
 232         }
 233         return size
 234 }
 235
 236 func (c *cache) lookupCollection(pdh string) *arvados.Collection {
 237         if pdh == "" {
 238                 return nil
 239         } else if ent, cached := c.collections.Get(pdh); !cached {
 240                 return nil
 241         } else {
 242                 ent := ent.(*cachedCollection)
 243                 if ent.expire.Before(time.Now()) {
 244                         c.collections.Remove(pdh)
 245                         return nil
 246                 } else {
 247                         atomic.AddUint64(&c.stats.CollectionHits, 1)
 248                         return ent.collection
 249                 }
 250         }
 251 }