From 115eb31f57eaa219f3e1b0c80c011b89c06f416e Mon Sep 17 00:00:00 2001
From: Michael Bolot <michael.bolot@suse.com>
Date: Fri, 26 Aug 2022 09:27:57 -0500
Subject: [PATCH] Adding logic to limit number of cached schemas

Adds logic to ensure that only one schema/access set
is cached for each user. This should improve memory
consumption
---
 pkg/accesscontrol/access_store.go |  5 +++++
 pkg/schema/collection.go          |  2 ++
 pkg/schema/factory.go             | 28 ++++++++++++++++++++++++++--
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/pkg/accesscontrol/access_store.go b/pkg/accesscontrol/access_store.go
index c14ee6e..5e51f97 100644
--- a/pkg/accesscontrol/access_store.go
+++ b/pkg/accesscontrol/access_store.go
@@ -14,6 +14,7 @@ import (
 
 type AccessSetLookup interface {
 	AccessFor(user user.Info) *AccessSet
+	PurgeUserData(id string)
 }
 
 type AccessStore struct {
@@ -63,6 +64,10 @@ func (l *AccessStore) AccessFor(user user.Info) *AccessSet {
 	return result
 }
 
+func (l *AccessStore) PurgeUserData(id string) {
+	l.cache.Remove(id)
+}
+
 func (l *AccessStore) CacheKey(user user.Info) string {
 	d := sha256.New()
 
diff --git a/pkg/schema/collection.go b/pkg/schema/collection.go
index 916e78f..48f23ea 100644
--- a/pkg/schema/collection.go
+++ b/pkg/schema/collection.go
@@ -36,6 +36,7 @@ type Collection struct {
 	byGVR      map[schema.GroupVersionResource]string
 	byGVK      map[schema.GroupVersionKind]string
 	cache      *cache.LRUExpireCache
+	userCache  *cache.LRUExpireCache
 	lock       sync.RWMutex
 
 	ctx     context.Context
@@ -84,6 +85,7 @@ func NewCollection(ctx context.Context, baseSchema *types.APISchemas, access acc
 		byGVR:      map[schema.GroupVersionResource]string{},
 		byGVK:      map[schema.GroupVersionKind]string{},
 		cache:      cache.NewLRUExpireCache(1000),
+		userCache:  cache.NewLRUExpireCache(1000),
 		notifiers:  map[int]func(){},
 		ctx:        ctx,
 		as:         access,
diff --git a/pkg/schema/factory.go b/pkg/schema/factory.go
index b16af18..5f7ebd9 100644
--- a/pkg/schema/factory.go
+++ b/pkg/schema/factory.go
@@ -23,6 +23,7 @@ func newSchemas() (*types.APISchemas, error) {
 
 func (c *Collection) Schemas(user user.Info) (*types.APISchemas, error) {
 	access := c.as.AccessFor(user)
+	c.removeOldRecords(access, user)
 	val, ok := c.cache.Get(access.ID)
 	if ok {
 		schemas, _ := val.(*types.APISchemas)
@@ -33,11 +34,34 @@ func (c *Collection) Schemas(user user.Info) (*types.APISchemas, error) {
 	if err != nil {
 		return nil, err
 	}
-
-	c.cache.Add(access.ID, schemas, 24*time.Hour)
+	c.addToCache(access, user, schemas)
 	return schemas, nil
 }
 
+func (c *Collection) removeOldRecords(access *accesscontrol.AccessSet, user user.Info) {
+	current, ok := c.userCache.Get(user.GetName())
+	if ok {
+		currentId, cOk := current.(string)
+		if cOk && currentId != access.ID {
+			// we only want to keep around one record per user. If our current access record is invalid, purge the
+			//record of it from the cache, so we don't keep duplicates
+			c.purgeUserRecords(currentId)
+			c.userCache.Remove(user.GetName())
+		}
+	}
+}
+
+func (c *Collection) addToCache(access *accesscontrol.AccessSet, user user.Info, schemas *types.APISchemas) {
+	c.cache.Add(access.ID, schemas, 24*time.Hour)
+	c.userCache.Add(user.GetName(), access.ID, 24*time.Hour)
+}
+
+// PurgeUserRecords removes a record from the backing LRU cache before expiry
+func (c *Collection) purgeUserRecords(id string) {
+	c.cache.Remove(id)
+	c.as.PurgeUserData(id)
+}
+
 func (c *Collection) schemasForSubject(access *accesscontrol.AccessSet) (*types.APISchemas, error) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()