pkg/services/libraryelements/cache_test.go - main - githubmirror/grafana

pkg/services/libraryelements/cache_test.go
368 строк12 KB
maicon
Library Elements: Speed up folder-tree resolution for GET /api/library-elements on large instances (#127762)
28 июл 2026, 17:16
28 июл 2026, 17:16d46296b
package libraryelements

import (
	"context"
	"encoding/json"
	"errors"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/grafana/grafana/pkg/services/folder"
	"github.com/grafana/grafana/pkg/services/folder/foldertest"
	"github.com/grafana/grafana/pkg/services/libraryelements/model"
	"github.com/grafana/grafana/pkg/services/org"
	searchmodel "github.com/grafana/grafana/pkg/services/search/model"
	"github.com/grafana/grafana/pkg/services/user"
	"github.com/grafana/grafana/pkg/util/testutil"
)

// trackingFolderService wraps a folder.Service and tracks how many times GetFolders is called.
// searchHits/searchErr are honored independently of FakeService.ExpectedHitList/
// ExpectedError so tests can drive the search path into failure without also
// breaking the GetFolders fallback that shares those fields on FakeService.
type trackingFolderService struct {
	*foldertest.FakeService
	getFoldersCallCount atomic.Int32
	searchHits          searchmodel.HitList
	searchErr           error
	overrideSearch      bool
}

func newTrackingFolderService() *trackingFolderService {
	return &trackingFolderService{
		FakeService: foldertest.NewFakeService(),
	}
}

func (s *trackingFolderService) GetFolders(ctx context.Context, q folder.GetFoldersQuery) ([]*folder.Folder, error) {
	s.getFoldersCallCount.Add(1)
	return s.FakeService.GetFolders(ctx, q)
}

func (s *trackingFolderService) SearchFolders(ctx context.Context, q folder.SearchFoldersQuery) (searchmodel.HitList, error) {
	if s.overrideSearch {
		return s.searchHits, s.searchErr
	}
	return s.FakeService.SearchFolders(ctx, q)
}

func (s *trackingFolderService) GetCallCount() int {
	return int(s.getFoldersCallCount.Load())
}

func TestIntegration_FolderTreeCache(t *testing.T) {
	testutil.SkipIntegrationTestInShortMode(t)

	t.Run("GetAll calls GetFolders once and caches across requests", func(t *testing.T) {
		sc := setupTestScenario(t)

		// Create multiple library panels
		for i := 0; i < 5; i++ {
			// nolint:staticcheck
			command := getCreatePanelCommand(sc.folder.ID, sc.folder.UID, "Panel "+string(rune('A'+i)))
			sc.reqContext.Req.Body = mockRequestBody(command)
			resp := sc.service.createHandler(sc.reqContext)
			require.Equal(t, 200, resp.Status())
		}

		// Replace folder service with tracking one
		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{sc.folder}
		trackingSvc.AddFolder(sc.folder)

		originalFolderSvc := sc.service.folderService
		sc.service.folderService = trackingSvc
		sc.service.treeCache = newFolderTreeCache(trackingSvc, false)
		defer func() { sc.service.folderService = originalFolderSvc }()

		// First request
		resp := sc.service.getAllHandler(sc.reqContext)
		require.Equal(t, 200, resp.Status())

		var result libraryElementsSearch
		err := json.Unmarshal(resp.Body(), &result)
		require.NoError(t, err)
		require.Equal(t, int64(5), result.Result.TotalCount)
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "GetFolders should be called once for tree build")

		// Second request reuses global cache
		resp = sc.service.getAllHandler(sc.reqContext)
		require.Equal(t, 200, resp.Status())
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "Second request should reuse cached tree")
	})
}

func TestFolderTreeCache_Unit(t *testing.T) {
	t.Run("builds tree from GetFolders result", func(t *testing.T) {
		sc := setupTestScenario(t)

		fakeSvc := foldertest.NewFakeService()
		fakeSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
			{UID: "folder-b", Title: "Folder B", OrgID: 1, ParentUID: "folder-a"},
		}

		cache := newFolderTreeCache(fakeSvc, false)
		tree, err := cache.get(context.Background(), sc.reqContext.SignedInUser)
		require.NoError(t, err)
		require.NotNil(t, tree)

		assert.True(t, tree.Contains("folder-a"))
		assert.True(t, tree.Contains("folder-b"))
		assert.Equal(t, "Folder A", tree.GetTitle("folder-a"))
		assert.Equal(t, "Folder B", tree.GetTitle("folder-b"))
	})

	t.Run("builds tree from SearchFolders result when useSearch is set", func(t *testing.T) {
		sc := setupTestScenario(t)

		fakeSvc := foldertest.NewFakeService()
		// SearchFolders returns lightweight hits (UID + parent via FolderUID).
		fakeSvc.ExpectedHitList = searchmodel.HitList{
			{UID: "folder-a", Title: "Folder A"},
			{UID: "folder-b", Title: "Folder B", FolderUID: "folder-a"},
		}

		cache := newFolderTreeCache(fakeSvc, true)
		tree, err := cache.get(context.Background(), sc.reqContext.SignedInUser)
		require.NoError(t, err)
		require.NotNil(t, tree)

		// Same tree as the GetFolders path, but sourced from the search index.
		assert.True(t, tree.Contains("folder-a"))
		assert.True(t, tree.Contains("folder-b"))
		assert.Equal(t, "Folder A", tree.GetTitle("folder-a"))
		assert.Equal(t, "Folder B", tree.GetTitle("folder-b"))
		// folder-b's parent (folder-a) is reflected in the tree.
		parents := make([]string, 0)
		for ancestor := range tree.Ancestors("folder-b") {
			parents = append(parents, ancestor.UID)
		}
		assert.Contains(t, parents, "folder-a")
	})

	t.Run("falls back to GetFolders when SearchFolders errors", func(t *testing.T) {
		sc := setupTestScenario(t)

		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
		}
		trackingSvc.overrideSearch = true
		trackingSvc.searchErr = errors.New("search index unavailable")

		cache := newFolderTreeCache(trackingSvc, true)
		tree, err := cache.get(context.Background(), sc.reqContext.SignedInUser)
		require.NoError(t, err)
		require.NotNil(t, tree)

		assert.True(t, tree.Contains("folder-a"))
		assert.Equal(t, 1, trackingSvc.GetCallCount())
	})

	t.Run("falls back to GetFolders when SearchFolders returns no hits", func(t *testing.T) {
		sc := setupTestScenario(t)

		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
		}
		// SearchFolders returns an empty hit list (index still warming up).
		trackingSvc.overrideSearch = true
		trackingSvc.searchHits = searchmodel.HitList{}

		cache := newFolderTreeCache(trackingSvc, true)
		tree, err := cache.get(context.Background(), sc.reqContext.SignedInUser)
		require.NoError(t, err)
		require.NotNil(t, tree)

		assert.True(t, tree.Contains("folder-a"))
		assert.Equal(t, 1, trackingSvc.GetCallCount())
	})

	t.Run("falls back to GetFolders when useSearch is set but the requester has no ID token", func(t *testing.T) {
		sc := setupTestScenario(t)
		sc.reqContext.IDToken = ""

		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
		}
		// Populate a hit list too — if the search path were taken, folder-b would
		// leak into the tree.
		trackingSvc.ExpectedHitList = searchmodel.HitList{
			{UID: "folder-b", Title: "Folder B"},
		}

		cache := newFolderTreeCache(trackingSvc, true)
		tree, err := cache.get(context.Background(), sc.reqContext.SignedInUser)
		require.NoError(t, err)
		require.NotNil(t, tree)

		assert.True(t, tree.Contains("folder-a"))
		// folder-b would only appear if the search path had been taken.
		assert.False(t, tree.Contains("folder-b"))
		assert.Equal(t, 1, trackingSvc.GetCallCount())
	})

	t.Run("caches tree and returns same instance on repeated calls", func(t *testing.T) {
		sc := setupTestScenario(t)

		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
		}

		cache := newFolderTreeCache(trackingSvc, false)
		ctx := context.Background()

		tree1, err := cache.get(ctx, sc.reqContext.SignedInUser)
		require.NoError(t, err)
		assert.Equal(t, 1, trackingSvc.GetCallCount())

		tree2, err := cache.get(ctx, sc.reqContext.SignedInUser)
		require.NoError(t, err)
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "Should not call GetFolders again within TTL")
		assert.Same(t, tree1, tree2, "Should return same cached tree instance")
	})

	t.Run("different user gets separate cache entry", func(t *testing.T) {
		sc := setupTestScenario(t)

		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{
			{UID: "folder-a", Title: "Folder A", OrgID: 1},
		}

		cache := newFolderTreeCache(trackingSvc, false)
		ctx := context.Background()

		// First user fetches and caches tree
		_, err := cache.get(ctx, sc.reqContext.SignedInUser)
		require.NoError(t, err)
		assert.Equal(t, 1, trackingSvc.GetCallCount())

		// Same user reuses cache
		_, err = cache.get(ctx, sc.reqContext.SignedInUser)
		require.NoError(t, err)
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "Same user should reuse cache")

		// Different user triggers a new GetFolders call
		user2 := &user.SignedInUser{
			UserID:  999,
			UserUID: "different-user-uid",
			OrgID:   sc.reqContext.GetOrgID(),
			OrgRole: org.RoleViewer,
		}
		_, err = cache.get(ctx, user2)
		require.NoError(t, err)
		assert.Equal(t, 2, trackingSvc.GetCallCount(), "Different user should trigger a new GetFolders call")
	})
}

func TestIntegration_SkipFolderTreeForAdmin(t *testing.T) {
	testutil.SkipIntegrationTestInShortMode(t)

	createPanels := func(t *testing.T, sc scenarioContext, count int) {
		t.Helper()
		for i := 0; i < count; i++ {
			// nolint:staticcheck
			command := getCreatePanelCommand(sc.folder.ID, sc.folder.UID, "Panel "+string(rune('A'+i)))
			sc.reqContext.Req.Body = mockRequestBody(command)
			resp := sc.service.createHandler(sc.reqContext)
			require.Equal(t, 200, resp.Status())
		}
	}

	replaceWithTrackingFolderSvc := func(sc scenarioContext) *trackingFolderService {
		trackingSvc := newTrackingFolderService()
		trackingSvc.ExpectedFolders = []*folder.Folder{sc.folder}
		trackingSvc.AddFolder(sc.folder)
		sc.service.folderService = trackingSvc
		sc.service.treeCache = newFolderTreeCache(trackingSvc, false)
		return trackingSvc
	}

	t.Run("admin with SkipFolderTreeForAdmin skips GetFolders", func(t *testing.T) {
		sc := setupTestScenario(t)
		createPanels(t, sc, 3)

		trackingSvc := replaceWithTrackingFolderSvc(sc)

		result, err := sc.service.getAllLibraryElements(context.Background(), sc.reqContext.SignedInUser, model.SearchLibraryElementsQuery{
			PerPage:                100,
			Page:                   1,
			SkipFolderTreeForAdmin: true,
		})
		require.NoError(t, err)
		assert.Equal(t, 3, len(result.Elements))
		assert.Equal(t, 0, trackingSvc.GetCallCount(), "GetFolders should not be called for admin with SkipFolderTreeForAdmin")

		// FolderName should be empty since folder tree was not fetched
		for _, elem := range result.Elements {
			assert.Empty(t, elem.Meta.FolderName, "FolderName should be empty when SkipFolderTreeForAdmin is set")
		}
	})

	t.Run("admin without SkipFolderTreeForAdmin fetches folder tree", func(t *testing.T) {
		sc := setupTestScenario(t)
		createPanels(t, sc, 3)

		trackingSvc := replaceWithTrackingFolderSvc(sc)

		result, err := sc.service.getAllLibraryElements(context.Background(), sc.reqContext.SignedInUser, model.SearchLibraryElementsQuery{
			PerPage:                100,
			Page:                   1,
			SkipFolderTreeForAdmin: false,
		})
		require.NoError(t, err)
		assert.Equal(t, 3, len(result.Elements))
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "GetFolders should be called once")

		for _, elem := range result.Elements {
			assert.Equal(t, sc.folder.Title, elem.Meta.FolderName, "FolderName should be populated")
		}
	})

	t.Run("non-admin with SkipFolderTreeForAdmin still fetches folder tree", func(t *testing.T) {
		sc := setupTestScenario(t)
		createPanels(t, sc, 3)

		trackingSvc := replaceWithTrackingFolderSvc(sc)

		viewer := &user.SignedInUser{
			UserID:  sc.user.UserID,
			UserUID: sc.user.UserUID,
			OrgID:   sc.user.OrgID,
			OrgRole: org.RoleViewer,
			Permissions: map[int64]map[string][]string{
				sc.user.OrgID: {
					folder.ActionFoldersRead: {folder.ScopeFoldersAll},
				},
			},
		}

		result, err := sc.service.getAllLibraryElements(context.Background(), viewer, model.SearchLibraryElementsQuery{
			PerPage:                100,
			Page:                   1,
			SkipFolderTreeForAdmin: true,
		})
		require.NoError(t, err)
		assert.Equal(t, 3, len(result.Elements))
		assert.Equal(t, 1, trackingSvc.GetCallCount(), "GetFolders should still be called for non-admin")
	})
}

func TestCacheKey(t *testing.T) {
	t.Run("withCache adds cache key to context", func(t *testing.T) {
		ctx := withCache(context.Background())
		assert.True(t, hasCache(ctx))
	})

	t.Run("hasCache returns false when not set", func(t *testing.T) {
		assert.False(t, hasCache(context.Background()))
	})
}