WIP: Implementing TTL for cached datasource instances

Author: iwysiu

backend/instancemgmt/instance_manager.go

@@ -2,14 +2,15 @@ package instancemgmt
 
 import (
 	"context"
+	"fmt"
 	"reflect"
-	"sync"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	gocache "github.com/patrickmn/go-cache"
 )
 
 var (
@@ -19,6 +20,9 @@ var (
 		Help:      "The number of active plugin instances",
 	})
 	disposeTTL = 5 * time.Second
+
+	instanceTTL     = 1 * time.Hour
+	instanceCleanup = 2 * time.Hour
 )
 
 // Instance is a marker interface for an instance.
@@ -76,18 +80,28 @@ func New(provider InstanceProvider) InstanceManager {
 	if provider == nil {
 		panic("provider cannot be nil")
 	}
+	cache := gocache.New(instanceTTL, instanceCleanup)
+	cache.OnEvicted(func(_ string, i interface{}) {
+		ci := i.(CachedInstance)
+		if disposer, valid := ci.instance.(InstanceDisposer); valid {
+			time.AfterFunc(disposeTTL, func() {
+				disposer.Dispose()
+			})
+		}
+		activeInstances.Dec()
+	})
 
 	return &instanceManager{
 		provider: provider,
-		cache:    sync.Map{},
+		cache:    cache,
 		locker:   newLocker(),
 	}
 }
 
 type instanceManager struct {
 	locker   *locker
 	provider InstanceProvider
-	cache    sync.Map
+	cache    *gocache.Cache
 }
 
 func (im *instanceManager) Get(ctx context.Context, pluginContext backend.PluginContext) (Instance, error) {
@@ -96,9 +110,10 @@ func (im *instanceManager) Get(ctx context.Context, pluginContext backend.Plugin
 		return nil, err
 	}
 	// Double-checked locking for update/create criteria
-	im.locker.RLock(cacheKey)
-	item, ok := im.cache.Load(cacheKey)
-	im.locker.RUnlock(cacheKey)
+	strKey := fmt.Sprintf("%v", cacheKey)
+	im.locker.RLock(strKey)
+	item, ok := im.cache.Get(strKey)
+	im.locker.RUnlock(strKey)
 
 	if ok {
 		ci := item.(CachedInstance)
@@ -109,35 +124,28 @@ func (im *instanceManager) Get(ctx context.Context, pluginContext backend.Plugin
 		}
 	}
 
-	im.locker.Lock(cacheKey)
-	defer im.locker.Unlock(cacheKey)
+	im.locker.Lock(strKey)
+	defer im.locker.Unlock(strKey)
 
-	if item, ok := im.cache.Load(cacheKey); ok {
+	if item, ok := im.cache.Get(strKey); ok {
 		ci := item.(CachedInstance)
 		needsUpdate := im.provider.NeedsUpdate(ctx, pluginContext, ci)
 
 		if !needsUpdate {
 			return ci.instance, nil
 		}
 
-		if disposer, valid := ci.instance.(InstanceDisposer); valid {
-			time.AfterFunc(disposeTTL, func() {
-				disposer.Dispose()
-				activeInstances.Dec()
-			})
-		} else {
-			activeInstances.Dec()
-		}
+		im.cache.Delete(strKey)
 	}
 
 	instance, err := im.provider.NewInstance(ctx, pluginContext)
 	if err != nil {
 		return nil, err
 	}
-	im.cache.Store(cacheKey, CachedInstance{
+	im.cache.Set(strKey, CachedInstance{
 		PluginContext: pluginContext,
 		instance:      instance,
-	})
+	}, gocache.DefaultExpiration)
 	activeInstances.Inc()
 
 	return instance, nil

backend/instancemgmt/instance_manager_test.go

@@ -70,6 +70,53 @@ func TestInstanceManager(t *testing.T) {
 	})
 }
 
+func TestInstanceManagerExpiration(t *testing.T) {
+	ctx := context.Background()
+	pCtx := backend.PluginContext{
+		OrgID: 1,
+		AppInstanceSettings: &backend.AppInstanceSettings{
+			Updated: time.Now(),
+		},
+	}
+
+	origInstanceTTL := instanceTTL
+	instanceTTL = time.Millisecond
+	origInstanceCleanup := instanceCleanup
+	instanceCleanup = 2 * time.Millisecond
+	t.Cleanup(func() {
+		instanceTTL = origInstanceTTL
+		instanceCleanup = origInstanceCleanup
+	})
+
+	tip := &testInstanceProvider{}
+	im := New(tip)
+
+	instance, err := im.Get(ctx, pCtx)
+	require.NoError(t, err)
+	require.NotNil(t, instance)
+	require.Equal(t, pCtx.OrgID, instance.(*testInstance).orgID)
+	require.Equal(t, pCtx.AppInstanceSettings.Updated, instance.(*testInstance).updated)
+
+	t.Run("After expiration", func(t *testing.T) {
+		instance.(*testInstance).wg.Wait()
+		require.True(t, instance.(*testInstance).disposed.Load())
+		require.Equal(t, int64(1), instance.(*testInstance).disposedTimes.Load())
+
+		newInstance, err := im.Get(ctx, pCtx)
+
+		t.Run("New instance should be created", func(t *testing.T) {
+			require.NoError(t, err)
+			require.NotNil(t, newInstance)
+			require.Equal(t, pCtx.OrgID, newInstance.(*testInstance).orgID)
+			require.Equal(t, pCtx.AppInstanceSettings.Updated, newInstance.(*testInstance).updated)
+		})
+
+		t.Run("New instance should not be the same as old instance", func(t *testing.T) {
+			require.NotSame(t, instance, newInstance)
+		})
+	})
+}
+
 func TestInstanceManagerConcurrency(t *testing.T) {
 	t.Run("Check possible race condition issues when initially creating instance", func(t *testing.T) {
 		ctx := context.Background()

go.mod

@@ -90,6 +90,7 @@ require (
 	github.com/oasdiff/yaml v0.0.0-20241210131133-6b86fb107d80 // indirect
 	github.com/oasdiff/yaml3 v0.0.0-20241210130736-a94c01f36349 // indirect
 	github.com/oklog/run v1.0.0 // indirect
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/perimeterx/marshmallow v1.1.5 // indirect
 	github.com/pierrec/lz4/v4 v4.1.21 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect

go.sum

@@ -170,6 +170,8 @@ github.com/oklog/run v1.0.0 h1:Ru7dDtJNOyC66gQ5dQmaCa0qIsAUFY3sFpK1Xk8igrw=
 github.com/oklog/run v1.0.0/go.mod h1:dlhp/R75TPv97u0XWUtDeV/lRKWPKSdTuV0TZvrmrQA=
 github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
 github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/perimeterx/marshmallow v1.1.5 h1:a2LALqQ1BlHM8PZblsDdidgv1mWi1DgC2UmX50IvK2s=
 github.com/perimeterx/marshmallow v1.1.5/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0Vchtcl8we9tYaXw=
 github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ=