Unified Storage Indexer: Build tenant indexes concurrently (#95795)

* WIP - build tenant indexes concurrently

* adds local dev seeders and readme for indexer

* update logging and adds locking in getShard()

* update logs

* Adds Namespaces func. Initializes index after ResourceServer is created.

* fixes Count() and updates test lint issues

* Initialize index separately. Don't do it when creating the resource server. Makes testing really awkward.

* fix lint error

* handles error when getting namespaces

* updates test and index helper funcs
This commit is contained in:
owensmallwood authored and GitHub committed 2024-11-06 12:58:07 -06:00
1 parent 3877537dc0
commit b6596db75e
13 files changed
+373 -70

No files matched your search

@@ -108,6 +108,10 @@ func (s *cdkBackend) getPath(key *ResourceKey, rv int64) string {
return buffer.String()
}
func (s *cdkBackend) Namespaces(ctx context.Context) ([]string, error) {
return nil, fmt.Errorf("not implemented")
}
func (s *cdkBackend) WriteEvent(ctx context.Context, event WriteEvent) (rv int64, err error) {
// Scope the lock
{
+113 -46
View File
@@ -4,6 +4,7 @@ import (
"context"
golog "log"
"path/filepath"
"sync"
"time"
"github.com/blevesearch/bleve/v2"
@@ -31,20 +32,22 @@ type Opts struct {
}
type Index struct {
shards map[string]Shard
opts Opts
s *server
log log.Logger
tracer tracing.Tracer
shardMutex sync.RWMutex
shards map[string]*Shard
opts Opts
s *server
log log.Logger
tracer tracing.Tracer
}
func NewIndex(s *server, opts Opts, tracer tracing.Tracer) *Index {
idx := &Index{
s: s,
opts: opts,
shards: make(map[string]Shard),
log: log.New("unifiedstorage.search.index"),
tracer: tracer,
shardMutex: sync.RWMutex{},
s: s,
opts: opts,
shards: make(map[string]*Shard),
log: log.New("unifiedstorage.search.index"),
tracer: tracer,
}
return idx
@@ -128,32 +131,86 @@ func (i *Index) AddToBatches(ctx context.Context, list *ListResponse) ([]string,
}
func (i *Index) Init(ctx context.Context) error {
logger := i.log.FromContext(ctx)
ctx, span := i.tracer.Start(ctx, tracingPrexfixIndex+"Init")
defer span.End()
start := time.Now().Unix()
group := errgroup.Group{}
group.SetLimit(i.opts.Workers)
totalObjects := 0
// Get all tenants currently in Unified Storage
tenants, err := i.s.backend.Namespaces(ctx)
if err != nil {
return err
}
for _, tenant := range tenants {
group.Go(func() error {
logger.Info("initializing index for tenant", "tenant", tenant)
objs, err := i.InitForTenant(ctx, tenant)
if err != nil {
return err
}
totalObjects += objs
return nil
})
}
err = group.Wait()
if err != nil {
return err
}
//index all remaining batches for all tenants
logger.Info("indexing remaining batches", "shards", len(i.shards))
err = i.IndexBatches(ctx, 1, i.allTenants())
if err != nil {
return err
}
end := time.Now().Unix()
totalDocCount := getTotalDocCount(i)
logger.Info("Initial indexing finished", "seconds", float64(end-start), "objs_fetched", totalObjects, "objs_indexed", totalDocCount)
span.AddEvent(
"indexing finished",
trace.WithAttributes(attribute.Int64("objects_indexed", int64(totalDocCount))),
trace.WithAttributes(attribute.Int64("objects_fetched", int64(totalObjects))),
)
if IndexServerMetrics != nil {
IndexServerMetrics.IndexCreationTime.WithLabelValues().Observe(float64(end - start))
}
return nil
}
func (i *Index) InitForTenant(ctx context.Context, namespace string) (int, error) {
ctx, span := i.tracer.Start(ctx, tracingPrexfixIndex+"InitForTenant")
defer span.End()
logger := i.log.FromContext(ctx)
start := time.Now().Unix()
resourceTypes := fetchResourceTypes()
totalObjectsFetched := 0
for _, rt := range resourceTypes {
logger.Info("indexing resource", "kind", rt.Key.Resource, "list_limit", i.opts.ListLimit, "batch_size", i.opts.BatchSize, "workers", i.opts.Workers)
logger.Debug("indexing resource", "kind", rt.Key.Resource, "list_limit", i.opts.ListLimit, "batch_size", i.opts.BatchSize, "workers", i.opts.Workers, "namespace", namespace)
r := &ListRequest{Options: rt, Limit: int64(i.opts.ListLimit)}
r.Options.Key.Namespace = namespace // scope the list to a tenant or this will take forever when US has 1M+ resources
// Paginate through the list of resources and index each page
for {
logger.Info("fetching resource list", "kind", rt.Key.Resource)
logger.Debug("fetching resource list", "kind", rt.Key.Resource, "namespace", namespace)
list, err := i.s.List(ctx, r)
if err != nil {
return err
return totalObjectsFetched, err
}
totalObjectsFetched += len(list.Items)
logger.Info("indexing batch", "kind", rt.Key.Resource, "count", len(list.Items))
logger.Debug("indexing batch", "kind", rt.Key.Resource, "count", len(list.Items), "namespace", namespace)
//add changes to batches for shards with changes in the List
err = i.writeBatch(ctx, list)
if err != nil {
return err
return totalObjectsFetched, err
}
if list.NextPageToken == "" {
@@ -164,21 +221,13 @@ func (i *Index) Init(ctx context.Context) error {
}
}
//index all remaining batches
logger.Info("indexing remaining batches", "shards", len(i.shards))
err := i.IndexBatches(ctx, 1, i.allTenants())
if err != nil {
return err
}
span.AddEvent(
"indexing finished for tenant",
trace.WithAttributes(attribute.Int64("objects_indexed", int64(totalObjectsFetched))),
trace.WithAttributes(attribute.String("tenant", namespace)),
)
span.AddEvent("indexing finished", trace.WithAttributes(attribute.Int64("objects_indexed", int64(totalObjectsFetched))))
end := time.Now().Unix()
logger.Info("Initial indexing finished", "seconds", float64(end-start))
if IndexServerMetrics != nil {
IndexServerMetrics.IndexCreationTime.WithLabelValues().Observe(float64(end - start))
}
return nil
return totalObjectsFetched, nil
}
func (i *Index) writeBatch(ctx context.Context, list *ListResponse) error {
@@ -207,6 +256,17 @@ func (i *Index) Index(ctx context.Context, data *Data) error {
}
tenant := res.Namespace
logger.Debug("indexing resource for tenant", "res", string(data.Value.Value), "tenant", tenant)
// if tenant doesn't exist, they may have been created during initial indexing
_, ok := i.shards[tenant]
if !ok {
i.log.Info("tenant not found, initializing their index", "tenant", tenant)
_, err = i.InitForTenant(ctx, tenant)
if err != nil {
return err
}
}
shard, err := i.getShard(tenant)
if err != nil {
return err
@@ -320,18 +380,20 @@ func (i *Index) Search(ctx context.Context, request *SearchRequest) (*IndexResul
return &IndexResults{Values: results, Groups: groups}, nil
}
func (i *Index) Count() (uint64, error) {
var total uint64
// Count returns the total doc count
func (i *Index) Count() (int, error) {
total := 0
for _, shard := range i.shards {
count, err := shard.index.DocCount()
if err != nil {
i.log.Error("failed to get doc count", "error", err)
}
total += count
total += int(count)
}
return total, nil
}
// allTenants returns a list of all tenants in the index
func (i *Index) allTenants() []string {
tenants := make([]string, 0, len(i.shards))
for tenant := range i.shards {
@@ -340,7 +402,10 @@ func (i *Index) allTenants() []string {
return tenants
}
func (i *Index) getShard(tenant string) (Shard, error) {
func (i *Index) getShard(tenant string) (*Shard, error) {
i.shardMutex.Lock()
defer i.shardMutex.Unlock()
shard, ok := i.shards[tenant]
if ok {
return shard, nil
@@ -348,16 +413,16 @@ func (i *Index) getShard(tenant string) (Shard, error) {
index, path, err := i.createIndex()
if err != nil {
return Shard{}, err
return &Shard{}, err
}
shard = Shard{
shard = &Shard{
index: index,
path: path,
batch: index.NewBatch(),
}
// TODO: do we need to lock this?
i.shards[tenant] = shard
return shard, nil
}
@@ -389,17 +454,19 @@ func createInMemoryIndex() (bleve.Index, string, error) {
// TODO - fetch from api
func fetchResourceTypes() []*ListOptions {
items := []*ListOptions{}
items = append(items, &ListOptions{
Key: &ResourceKey{
Group: "playlist.grafana.app",
Resource: "playlists",
items = append(items,
&ListOptions{
Key: &ResourceKey{
Group: "playlist.grafana.app",
Resource: "playlists",
},
},
}, &ListOptions{
Key: &ResourceKey{
Group: "folder.grafana.app",
Resource: "folders",
&ListOptions{
Key: &ResourceKey{
Group: "folder.grafana.app",
Resource: "folders",
},
},
},
&ListOptions{
Key: &ResourceKey{
Group: "dashboard.grafana.app",
@@ -87,23 +87,23 @@ func (s *IndexMetrics) Describe(ch chan<- *prometheus.Desc) {
s.IndexLatency.Describe(ch)
}
// getTotalDocCount returns the total number of documents in the index
func getTotalDocCount(index *Index) float64 {
var totalCount float64
totalCount = 0
if index == nil {
return totalCount
}
for _, shard := range index.shards {
docCount, err := shard.index.DocCount()
if err != nil {
continue
}
totalCount += float64(docCount)
count, _ := shard.index.DocCount()
totalCount += float64(count)
}
return totalCount
}
// getTotalIndexSize returns the total size of the index directory when using a file-based index
func getTotalIndexSize(dir string) (int64, error) {
var totalSize int64
+4 -3
View File
@@ -29,6 +29,7 @@ func TestIndexDashboard(t *testing.T) {
require.NoError(t, err)
assertCountEquals(t, index, 1)
require.Equal(t, 1, len(index.allTenants()))
assertSearchCountEquals(t, index, "*", nil, 1)
}
@@ -83,7 +84,7 @@ func TestLookupNames(t *testing.T) {
err := index.writeBatch(testContext, list)
require.NoError(t, err)
assertCountEquals(t, index, uint64(records))
assertCountEquals(t, index, records)
query := ""
chunk := ids[:100] // query for n folders by id
for _, id := range chunk {
@@ -185,7 +186,7 @@ func newTestIndex(t *testing.T, batchSize int) *Index {
return &Index{
tracer: trace,
shards: make(map[string]Shard),
shards: make(map[string]*Shard),
log: log.New("unifiedstorage.search.index"),
opts: Opts{
ListLimit: 5000,
@@ -195,7 +196,7 @@ func newTestIndex(t *testing.T, batchSize int) *Index {
}
}
func assertCountEquals(t *testing.T, index *Index, expected uint64) {
func assertCountEquals(t *testing.T, index *Index, expected int) {
total, err := index.Count()
require.NoError(t, err)
assert.Equal(t, expected, total)
+10 -2
View File
@@ -75,6 +75,8 @@ type StorageBackend interface {
// Get all events from the store
// For HA setups, this will be more events than the local WriteEvent above!
WatchWriteEvents(ctx context.Context) (<-chan *WrittenEvent, error)
Namespaces(ctx context.Context) ([]string, error)
}
// This interface is not exposed to end users directly
@@ -182,7 +184,7 @@ func NewResourceServer(opts ResourceServerOptions) (ResourceServer, error) {
UserID: 1,
IsGrafanaAdmin: true,
}))
return &server{
s := &server{
tracer: opts.Tracer,
log: logger,
backend: opts.Backend,
@@ -194,7 +196,9 @@ func NewResourceServer(opts ResourceServerOptions) (ResourceServer, error) {
now: opts.Now,
ctx: ctx,
cancel: cancel,
}, nil
}
return s, nil
}
var _ ResourceServer = &server{}
@@ -765,6 +769,10 @@ func (s *server) Origin(ctx context.Context, req *OriginRequest) (*OriginRespons
// Index returns the search index. If the index is not initialized, it will be initialized.
func (s *server) Index(ctx context.Context) (*Index, error) {
if err := s.Init(ctx); err != nil {
return nil, err
}
index := s.index.(*IndexServer)
if index.index == nil {
err := index.Init(ctx, s)