package main import ( "sync" "sync/atomic" "time" "git.curoverse.com/arvados.git/sdk/go/arvados" "git.curoverse.com/arvados.git/sdk/go/arvadosclient" "github.com/hashicorp/golang-lru" ) type cache struct { TTL arvados.Duration MaxCollectionEntries int MaxCollectionBytes int64 MaxPermissionEntries int MaxUUIDEntries int stats cacheStats pdhs *lru.TwoQueueCache collections *lru.TwoQueueCache permissions *lru.TwoQueueCache setupOnce sync.Once } type cacheStats struct { Requests uint64 `json:"Cache.Requests"` CollectionBytes uint64 `json:"Cache.CollectionBytes"` CollectionEntries int `json:"Cache.CollectionEntries"` CollectionHits uint64 `json:"Cache.CollectionHits"` PDHHits uint64 `json:"Cache.UUIDHits"` PermissionHits uint64 `json:"Cache.PermissionHits"` APICalls uint64 `json:"Cache.APICalls"` } type cachedPDH struct { expire time.Time pdh string } type cachedCollection struct { expire time.Time collection *arvados.Collection } type cachedPermission struct { expire time.Time } func (c *cache) setup() { var err error c.pdhs, err = lru.New2Q(c.MaxUUIDEntries) if err != nil { panic(err) } c.collections, err = lru.New2Q(c.MaxCollectionEntries) if err != nil { panic(err) } c.permissions, err = lru.New2Q(c.MaxPermissionEntries) if err != nil { panic(err) } } var selectPDH = map[string]interface{}{ "select": []string{"portable_data_hash"}, } func (c *cache) Stats() cacheStats { c.setupOnce.Do(c.setup) return cacheStats{ Requests: atomic.LoadUint64(&c.stats.Requests), CollectionBytes: c.collectionBytes(), CollectionEntries: c.collections.Len(), CollectionHits: atomic.LoadUint64(&c.stats.CollectionHits), PDHHits: atomic.LoadUint64(&c.stats.PDHHits), PermissionHits: atomic.LoadUint64(&c.stats.PermissionHits), APICalls: atomic.LoadUint64(&c.stats.APICalls), } } func (c *cache) Get(arv *arvadosclient.ArvadosClient, targetID string, forceReload bool) (*arvados.Collection, error) { c.setupOnce.Do(c.setup) atomic.AddUint64(&c.stats.Requests, 1) permOK := false permKey := arv.ApiToken + "\000" + targetID if forceReload { } else if ent, cached := c.permissions.Get(permKey); cached { ent := ent.(*cachedPermission) if ent.expire.Before(time.Now()) { c.permissions.Remove(permKey) } else { permOK = true atomic.AddUint64(&c.stats.PermissionHits, 1) } } var pdh string if arvadosclient.PDHMatch(targetID) { pdh = targetID } else if forceReload { } else if ent, cached := c.pdhs.Get(targetID); cached { ent := ent.(*cachedPDH) if ent.expire.Before(time.Now()) { c.pdhs.Remove(targetID) } else { pdh = ent.pdh atomic.AddUint64(&c.stats.PDHHits, 1) } } var collection *arvados.Collection if pdh != "" { collection = c.lookupCollection(pdh) } if collection != nil && permOK { return collection, nil } else if collection != nil { // Ask API for current PDH for this targetID. Most // likely, the cached PDH is still correct; if so, // _and_ the current token has permission, we can // use our cached manifest. atomic.AddUint64(&c.stats.APICalls, 1) var current arvados.Collection err := arv.Get("collections", targetID, selectPDH, ¤t) if err != nil { return nil, err } if current.PortableDataHash == pdh { exp := time.Now().Add(time.Duration(c.TTL)) c.permissions.Add(permKey, &cachedPermission{ expire: exp, }) if pdh != targetID { c.pdhs.Add(targetID, &cachedPDH{ expire: exp, pdh: pdh, }) } return collection, err } else { // PDH changed, but now we know we have // permission -- and maybe we already have the // new PDH in the cache. if coll := c.lookupCollection(current.PortableDataHash); coll != nil { return coll, nil } } } // Collection manifest is not cached. atomic.AddUint64(&c.stats.APICalls, 1) err := arv.Get("collections", targetID, nil, &collection) if err != nil { return nil, err } exp := time.Now().Add(time.Duration(c.TTL)) c.permissions.Add(permKey, &cachedPermission{ expire: exp, }) c.pdhs.Add(targetID, &cachedPDH{ expire: exp, pdh: collection.PortableDataHash, }) c.collections.Add(collection.PortableDataHash, &cachedCollection{ expire: exp, collection: collection, }) if int64(len(collection.ManifestText)) > c.MaxCollectionBytes/int64(c.MaxCollectionEntries) { go c.pruneCollections() } return collection, nil } // pruneCollections checks the total bytes occupied by manifest_text // in the collection cache and removes old entries as needed to bring // the total size down to CollectionBytes. It also deletes all expired // entries. // // pruneCollections does not aim to be perfectly correct when there is // concurrent cache activity. func (c *cache) pruneCollections() { var size int64 now := time.Now() keys := c.collections.Keys() entsize := make([]int, len(keys)) expired := make([]bool, len(keys)) for i, k := range keys { v, ok := c.collections.Peek(k) if !ok { continue } ent := v.(*cachedCollection) n := len(ent.collection.ManifestText) size += int64(n) entsize[i] = n expired[i] = ent.expire.Before(now) } for i, k := range keys { if expired[i] { c.collections.Remove(k) size -= int64(entsize[i]) } } for i, k := range keys { if size <= c.MaxCollectionBytes { break } if expired[i] { // already removed this entry in the previous loop continue } c.collections.Remove(k) size -= int64(entsize[i]) } } // collectionBytes returns the approximate memory size of the // collection cache. func (c *cache) collectionBytes() uint64 { var size uint64 for _, k := range c.collections.Keys() { v, ok := c.collections.Peek(k) if !ok { continue } size += uint64(len(v.(*cachedCollection).collection.ManifestText)) } return size } func (c *cache) lookupCollection(pdh string) *arvados.Collection { if pdh == "" { return nil } else if ent, cached := c.collections.Get(pdh); !cached { return nil } else { ent := ent.(*cachedCollection) if ent.expire.Before(time.Now()) { c.collections.Remove(pdh) return nil } else { atomic.AddUint64(&c.stats.CollectionHits, 1) return ent.collection } } }