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         MaxCollectionEntries int
  20         MaxCollectionBytes   int64
  21         MaxPermissionEntries int
  22         MaxUUIDEntries       int
  23
  24         stats       cacheStats
  25         pdhs        *lru.TwoQueueCache
  26         collections *lru.TwoQueueCache
  27         permissions *lru.TwoQueueCache
  28         setupOnce   sync.Once
  29 }
  30
  31 type cacheStats struct {
  32         Requests          uint64 `json:"Cache.Requests"`
  33         CollectionBytes   uint64 `json:"Cache.CollectionBytes"`
  34         CollectionEntries int    `json:"Cache.CollectionEntries"`
  35         CollectionHits    uint64 `json:"Cache.CollectionHits"`
  36         PDHHits           uint64 `json:"Cache.UUIDHits"`
  37         PermissionHits    uint64 `json:"Cache.PermissionHits"`
  38         APICalls          uint64 `json:"Cache.APICalls"`
  39 }
  40
  41 type cachedPDH struct {
  42         expire time.Time
  43         pdh    string
  44 }
  45
  46 type cachedCollection struct {
  47         expire     time.Time
  48         collection *arvados.Collection
  49 }
  50
  51 type cachedPermission struct {
  52         expire time.Time
  53 }
  54
  55 func (c *cache) setup() {
  56         var err error
  57         c.pdhs, err = lru.New2Q(c.MaxUUIDEntries)
  58         if err != nil {
  59                 panic(err)
  60         }
  61         c.collections, err = lru.New2Q(c.MaxCollectionEntries)
  62         if err != nil {
  63                 panic(err)
  64         }
  65         c.permissions, err = lru.New2Q(c.MaxPermissionEntries)
  66         if err != nil {
  67                 panic(err)
  68         }
  69 }
  70
  71 var selectPDH = map[string]interface{}{
  72         "select": []string{"portable_data_hash"},
  73 }
  74
  75 func (c *cache) Stats() cacheStats {
  76         c.setupOnce.Do(c.setup)
  77         return cacheStats{
  78                 Requests:          atomic.LoadUint64(&c.stats.Requests),
  79                 CollectionBytes:   c.collectionBytes(),
  80                 CollectionEntries: c.collections.Len(),
  81                 CollectionHits:    atomic.LoadUint64(&c.stats.CollectionHits),
  82                 PDHHits:           atomic.LoadUint64(&c.stats.PDHHits),
  83                 PermissionHits:    atomic.LoadUint64(&c.stats.PermissionHits),
  84                 APICalls:          atomic.LoadUint64(&c.stats.APICalls),
  85         }
  86 }
  87
  88 func (c *cache) Get(arv *arvadosclient.ArvadosClient, targetID string, forceReload bool) (*arvados.Collection, error) {
  89         c.setupOnce.Do(c.setup)
  90
  91         atomic.AddUint64(&c.stats.Requests, 1)
  92
  93         permOK := false
  94         permKey := arv.ApiToken + "\000" + targetID
  95         if forceReload {
  96         } else if ent, cached := c.permissions.Get(permKey); cached {
  97                 ent := ent.(*cachedPermission)
  98                 if ent.expire.Before(time.Now()) {
  99                         c.permissions.Remove(permKey)
 100                 } else {
 101                         permOK = true
 102                         atomic.AddUint64(&c.stats.PermissionHits, 1)
 103                 }
 104         }
 105
 106         var pdh string
 107         if arvadosclient.PDHMatch(targetID) {
 108                 pdh = targetID
 109         } else if ent, cached := c.pdhs.Get(targetID); cached {
 110                 ent := ent.(*cachedPDH)
 111                 if ent.expire.Before(time.Now()) {
 112                         c.pdhs.Remove(targetID)
 113                 } else {
 114                         pdh = ent.pdh
 115                         atomic.AddUint64(&c.stats.PDHHits, 1)
 116                 }
 117         }
 118
 119         var collection *arvados.Collection
 120         if pdh != "" {
 121                 collection = c.lookupCollection(pdh)
 122         }
 123
 124         if collection != nil && permOK {
 125                 return collection, nil
 126         } else if collection != nil {
 127                 // Ask API for current PDH for this targetID. Most
 128                 // likely, the cached PDH is still correct; if so,
 129                 // _and_ the current token has permission, we can
 130                 // use our cached manifest.
 131                 atomic.AddUint64(&c.stats.APICalls, 1)
 132                 var current arvados.Collection
 133                 err := arv.Get("collections", targetID, selectPDH, &current)
 134                 if err != nil {
 135                         return nil, err
 136                 }
 137                 if current.PortableDataHash == pdh {
 138                         exp := time.Now().Add(time.Duration(c.TTL))
 139                         c.permissions.Add(permKey, &cachedPermission{
 140                                 expire: exp,
 141                         })
 142                         if pdh != targetID {
 143                                 c.pdhs.Add(targetID, &cachedPDH{
 144                                         expire: exp,
 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(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: exp,
 171                 pdh:    collection.PortableDataHash,
 172         })
 173         c.collections.Add(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(pdh string) *arvados.Collection {
 241         if pdh == "" {
 242                 return nil
 243         } else if ent, cached := c.collections.Get(pdh); !cached {
 244                 return nil
 245         } else {
 246                 ent := ent.(*cachedCollection)
 247                 if ent.expire.Before(time.Now()) {
 248                         c.collections.Remove(pdh)
 249                         return nil
 250                 } else {
 251                         atomic.AddUint64(&c.stats.CollectionHits, 1)
 252                         return ent.collection
 253                 }
 254         }
 255 }