feat: add caching to reduce db calls

DanielleMaywood · DanielleMaywood · commit de6e6af41f33 · 2024-11-07T13:07:27.000Z
diff --git a/coderd/autobuild/cache.go b/coderd/autobuild/cache.go
@@ -0,0 +1,41 @@
+package autobuild
+
+import (
+	"errors"
+	"sync"
+)
+
+var errCacheLoaderNil = errors.New("loader is nil")
+
+type cacheOf[K comparable, V any] struct {
+	mu sync.Mutex
+	m  map[K]V
+}
+
+func newCacheOf[K comparable, V any]() *cacheOf[K, V] {
+	return &cacheOf[K, V]{
+		m: make(map[K]V),
+	}
+}
+
+func (c *cacheOf[K, V]) LoadOrStore(key K, loader func() (V, error)) (V, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	value, found := c.m[key]
+	if !found {
+		if loader == nil {
+			return *new(V), errCacheLoaderNil
+		}
+
+		loaded, err := loader()
+		if err != nil {
+			return *new(V), err
+		}
+
+		c.m[key] = loaded
+		return loaded, nil
+	}
+
+	return value, nil
+}
diff --git a/coderd/autobuild/cache_test.go b/coderd/autobuild/cache_test.go
@@ -0,0 +1,84 @@
+package autobuild
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCache(t *testing.T) {
+	t.Parallel()
+
+	t.Run("CallsLoaderOnce", func(t *testing.T) {
+		callCount := 0
+		cache := newCacheOf[uuid.UUID, int]()
+		key := uuid.New()
+
+		// Call LoadOrStore for key `key` for the first time.
+		// We expect this to call our loader function.
+		value, err := cache.LoadOrStore(key, func() (int, error) {
+			callCount += 1
+			return 1, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, 1, value)
+		require.Equal(t, 1, callCount)
+
+		// Call LoadOrStore for key `key` for the second time.
+		// We expect this to return data from the previous load.
+		value, err = cache.LoadOrStore(key, func() (int, error) {
+			callCount += 1
+
+			// We return a different value to further check
+			// that this function isn't called.
+			return 2, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, 1, value)
+		require.Equal(t, 1, callCount)
+	})
+
+	t.Run("ReturnsErrOnLoaderErr", func(t *testing.T) {
+		exampleErr := errors.New("example error")
+		cache := newCacheOf[uuid.UUID, int]()
+		key := uuid.New()
+
+		_, err := cache.LoadOrStore(key, func() (int, error) {
+			return 0, exampleErr
+		})
+		require.Error(t, err)
+		require.Equal(t, exampleErr, err)
+	})
+
+	t.Run("CanCacheWithMultipleKeys", func(t *testing.T) {
+		cache := newCacheOf[uuid.UUID, int]()
+		keyA := uuid.New()
+		keyB := uuid.New()
+
+		// We first insert data with our first key
+		value, err := cache.LoadOrStore(keyA, func() (int, error) {
+			return 10, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, 10, value)
+
+		// Next we insert data with a different key
+		value, err = cache.LoadOrStore(keyB, func() (int, error) {
+			return 20, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, 20, value)
+
+		// Now we check the data is still available for the first key
+		value, err = cache.LoadOrStore(keyA, nil)
+		require.NoError(t, err)
+		require.Equal(t, 10, value)
+
+		// And that the data is also still available for the second key
+		value, err = cache.LoadOrStore(keyB, nil)
+		require.NoError(t, err)
+		require.Equal(t, 20, value)
+	})
+}
diff --git a/coderd/autobuild/lifecycle_executor.go b/coderd/autobuild/lifecycle_executor.go
@@ -154,6 +154,16 @@ func (e *Executor) runOnce(t time.Time) Stats {
 	// Limit the concurrency to avoid overloading the database.
 	eg.SetLimit(10)
 
+	// We cache these values to help reduce load on the database.
+	// These could be outdated during our execution, but this is
+	// unlikely to be noticed or cause any unwanted behaviour.
+	var (
+		userCache             = newCacheOf[uuid.UUID, database.User]()
+		templateCache         = newCacheOf[uuid.UUID, database.Template]()
+		templateVersionCache  = newCacheOf[uuid.UUID, database.TemplateVersion]()
+		templateScheduleCache = newCacheOf[uuid.UUID, schedule.TemplateScheduleOptions]()
+	)
+
 	for _, ws := range workspaces {
 		wsID := ws.ID
 		wsName := ws.Name
@@ -184,7 +194,9 @@ func (e *Executor) runOnce(t time.Time) Stats {
 						return xerrors.Errorf("get workspace by id: %w", err)
 					}
 
-					user, err := tx.GetUserByID(e.ctx, ws.OwnerID)
+					user, err := userCache.LoadOrStore(ws.OwnerID, func() (database.User, error) {
+						return tx.GetUserByID(e.ctx, ws.OwnerID)
+					})
 					if err != nil {
 						return xerrors.Errorf("get user by id: %w", err)
 					}
@@ -200,17 +212,23 @@ func (e *Executor) runOnce(t time.Time) Stats {
 						return xerrors.Errorf("get latest provisioner job: %w", err)
 					}
 
-					templateSchedule, err := (*(e.templateScheduleStore.Load())).Get(e.ctx, tx, ws.TemplateID)
+					templateSchedule, err := templateScheduleCache.LoadOrStore(ws.TemplateID, func() (schedule.TemplateScheduleOptions, error) {
+						return (*(e.templateScheduleStore.Load())).Get(e.ctx, tx, ws.TemplateID)
+					})
 					if err != nil {
 						return xerrors.Errorf("get template scheduling options: %w", err)
 					}
 
-					tmpl, err = tx.GetTemplateByID(e.ctx, ws.TemplateID)
+					tmpl, err := templateCache.LoadOrStore(ws.TemplateID, func() (database.Template, error) {
+						return tx.GetTemplateByID(e.ctx, ws.TemplateID)
+					})
 					if err != nil {
 						return xerrors.Errorf("get template by ID: %w", err)
 					}
 
-					activeTemplateVersion, err = tx.GetTemplateVersionByID(e.ctx, tmpl.ActiveVersionID)
+					activeTemplateVersion, err = templateVersionCache.LoadOrStore(tmpl.ActiveVersionID, func() (database.TemplateVersion, error) {
+						return tx.GetTemplateVersionByID(e.ctx, tmpl.ActiveVersionID)
+					})
 					if err != nil {
 						return xerrors.Errorf("get active template version by ID: %w", err)
 					}
