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 // Update saves a modified version (fs) to an existing collection
  90 // (coll) and, if successful, updates the relevant cache entries so
  91 // subsequent calls to Get() reflect the modifications.
  92 func (c *cache) Update(client *arvados.Client, coll arvados.Collection, fs arvados.CollectionFileSystem) error {
  93         c.setupOnce.Do(c.setup)
  94
  95         if m, err := fs.MarshalManifest("."); err != nil || m == coll.ManifestText {
  96                 return err
  97         } else {
  98                 coll.ManifestText = m
  99         }
 100         var updated arvados.Collection
 101         defer c.pdhs.Remove(coll.UUID)
 102         err := client.RequestAndDecode(&updated, "PATCH", "arvados/v1/collections/"+coll.UUID, client.UpdateBody(coll), nil)
 103         if err == nil {
 104                 c.collections.Add(client.AuthToken+"\000"+coll.PortableDataHash, &cachedCollection{
 105                         expire:     time.Now().Add(time.Duration(c.TTL)),
 106                         collection: &updated,
 107                 })
 108         }
 109         return err
 110 }
 111
 112 func (c *cache) Get(arv *arvadosclient.ArvadosClient, targetID string, forceReload bool) (*arvados.Collection, error) {
 113         c.setupOnce.Do(c.setup)
 114
 115         atomic.AddUint64(&c.stats.Requests, 1)
 116
 117         permOK := false
 118         permKey := arv.ApiToken + "\000" + targetID
 119         if forceReload {
 120         } else if ent, cached := c.permissions.Get(permKey); cached {
 121                 ent := ent.(*cachedPermission)
 122                 if ent.expire.Before(time.Now()) {
 123                         c.permissions.Remove(permKey)
 124                 } else {
 125                         permOK = true
 126                         atomic.AddUint64(&c.stats.PermissionHits, 1)
 127                 }
 128         }
 129
 130         var pdh string
 131         if arvadosclient.PDHMatch(targetID) {
 132                 pdh = targetID
 133         } else if ent, cached := c.pdhs.Get(targetID); cached {
 134                 ent := ent.(*cachedPDH)
 135                 if ent.expire.Before(time.Now()) {
 136                         c.pdhs.Remove(targetID)
 137                 } else {
 138                         pdh = ent.pdh
 139                         atomic.AddUint64(&c.stats.PDHHits, 1)
 140                 }
 141         }
 142
 143         var collection *arvados.Collection
 144         if pdh != "" {
 145                 collection = c.lookupCollection(arv.ApiToken + "\000" + pdh)
 146         }
 147
 148         if collection != nil && permOK {
 149                 return collection, nil
 150         } else if collection != nil {
 151                 // Ask API for current PDH for this targetID. Most
 152                 // likely, the cached PDH is still correct; if so,
 153                 // _and_ the current token has permission, we can
 154                 // use our cached manifest.
 155                 atomic.AddUint64(&c.stats.APICalls, 1)
 156                 var current arvados.Collection
 157                 err := arv.Get("collections", targetID, selectPDH, &current)
 158                 if err != nil {
 159                         return nil, err
 160                 }
 161                 if current.PortableDataHash == pdh {
 162                         c.permissions.Add(permKey, &cachedPermission{
 163                                 expire: time.Now().Add(time.Duration(c.TTL)),
 164                         })
 165                         if pdh != targetID {
 166                                 c.pdhs.Add(targetID, &cachedPDH{
 167                                         expire: time.Now().Add(time.Duration(c.UUIDTTL)),
 168                                         pdh:    pdh,
 169                                 })
 170                         }
 171                         return collection, err
 172                 } else {
 173                         // PDH changed, but now we know we have
 174                         // permission -- and maybe we already have the
 175                         // new PDH in the cache.
 176                         if coll := c.lookupCollection(arv.ApiToken + "\000" + current.PortableDataHash); coll != nil {
 177                                 return coll, nil
 178                         }
 179                 }
 180         }
 181
 182         // Collection manifest is not cached.
 183         atomic.AddUint64(&c.stats.APICalls, 1)
 184         err := arv.Get("collections", targetID, nil, &collection)
 185         if err != nil {
 186                 return nil, err
 187         }
 188         exp := time.Now().Add(time.Duration(c.TTL))
 189         c.permissions.Add(permKey, &cachedPermission{
 190                 expire: exp,
 191         })
 192         c.pdhs.Add(targetID, &cachedPDH{
 193                 expire: time.Now().Add(time.Duration(c.UUIDTTL)),
 194                 pdh:    collection.PortableDataHash,
 195         })
 196         c.collections.Add(arv.ApiToken+"\000"+collection.PortableDataHash, &cachedCollection{
 197                 expire:     exp,
 198                 collection: collection,
 199         })
 200         if int64(len(collection.ManifestText)) > c.MaxCollectionBytes/int64(c.MaxCollectionEntries) {
 201                 go c.pruneCollections()
 202         }
 203         return collection, nil
 204 }
 205
 206 // pruneCollections checks the total bytes occupied by manifest_text
 207 // in the collection cache and removes old entries as needed to bring
 208 // the total size down to CollectionBytes. It also deletes all expired
 209 // entries.
 210 //
 211 // pruneCollections does not aim to be perfectly correct when there is
 212 // concurrent cache activity.
 213 func (c *cache) pruneCollections() {
 214         var size int64
 215         now := time.Now()
 216         keys := c.collections.Keys()
 217         entsize := make([]int, len(keys))
 218         expired := make([]bool, len(keys))
 219         for i, k := range keys {
 220                 v, ok := c.collections.Peek(k)
 221                 if !ok {
 222                         continue
 223                 }
 224                 ent := v.(*cachedCollection)
 225                 n := len(ent.collection.ManifestText)
 226                 size += int64(n)
 227                 entsize[i] = n
 228                 expired[i] = ent.expire.Before(now)
 229         }
 230         for i, k := range keys {
 231                 if expired[i] {
 232                         c.collections.Remove(k)
 233                         size -= int64(entsize[i])
 234                 }
 235         }
 236         for i, k := range keys {
 237                 if size <= c.MaxCollectionBytes {
 238                         break
 239                 }
 240                 if expired[i] {
 241                         // already removed this entry in the previous loop
 242                         continue
 243                 }
 244                 c.collections.Remove(k)
 245                 size -= int64(entsize[i])
 246         }
 247 }
 248
 249 // collectionBytes returns the approximate memory size of the
 250 // collection cache.
 251 func (c *cache) collectionBytes() uint64 {
 252         var size uint64
 253         for _, k := range c.collections.Keys() {
 254                 v, ok := c.collections.Peek(k)
 255                 if !ok {
 256                         continue
 257                 }
 258                 size += uint64(len(v.(*cachedCollection).collection.ManifestText))
 259         }
 260         return size
 261 }
 262
 263 func (c *cache) lookupCollection(key string) *arvados.Collection {
 264         if ent, cached := c.collections.Get(key); !cached {
 265                 return nil
 266         } else {
 267                 ent := ent.(*cachedCollection)
 268                 if ent.expire.Before(time.Now()) {
 269                         c.collections.Remove(key)
 270                         return nil
 271                 } else {
 272                         atomic.AddUint64(&c.stats.CollectionHits, 1)
 273                         return ent.collection
 274                 }
 275         }
 276 }