services/keep-web/cache.go

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