From eca379e98e636e92d60abcc195aedc09c7d63c9f Mon Sep 17 00:00:00 2001 From: Georges Chaudy Date: Tue, 21 Oct 2025 17:42:07 +0200 Subject: [PATCH] kvstore: update storage backend to use batch get (#112704) * consume batchget * remove rv * remove un-necessary limit * do not pass stop to the iterators * ensure we return the error --- pkg/storage/unified/resource/datastore.go | 59 +++++++ .../unified/resource/datastore_test.go | 145 ++++++++++++++++ .../unified/resource/storage_backend.go | 158 +++++++++--------- 3 files changed, 284 insertions(+), 78 deletions(-) diff --git a/pkg/storage/unified/resource/datastore.go b/pkg/storage/unified/resource/datastore.go index d5fc0cdf915..8ecb878b675 100644 --- a/pkg/storage/unified/resource/datastore.go +++ b/pkg/storage/unified/resource/datastore.go @@ -19,6 +19,8 @@ const ( dataSection = "unified/data" // cache groupResourcesCacheKey = "group-resources" + // batch operations + dataBatchSize = 50 // default batch size for BatchGet operations ) // dataStore is a data store that uses a KV store to store data. @@ -360,6 +362,63 @@ func (d *dataStore) Get(ctx context.Context, key DataKey) (io.ReadCloser, error) return d.kv.Get(ctx, dataSection, key.String()) } +// BatchGet retrieves multiple data objects in batches. +// It returns an iterator that yields DataObj results for the given keys. +// Keys are processed in batches (default 50) to balance between efficiency and memory usage. +// Non-existent entries will not appear in the result. +func (d *dataStore) BatchGet(ctx context.Context, keys []DataKey) iter.Seq2[DataObj, error] { + return func(yield func(DataObj, error) bool) { + // Validate all keys first + for _, key := range keys { + if err := key.Validate(); err != nil { + yield(DataObj{}, fmt.Errorf("invalid data key %s: %w", key.String(), err)) + return + } + } + + // Process keys in batches + for i := 0; i < len(keys); i += dataBatchSize { + end := i + dataBatchSize + if end > len(keys) { + end = len(keys) + } + batch := keys[i:end] + + // Convert DataKeys to string keys and create a mapping + stringKeys := make([]string, len(batch)) + keyMap := make(map[string]DataKey) // map string key back to DataKey + for j, key := range batch { + strKey := key.String() + stringKeys[j] = strKey + keyMap[strKey] = key + } + + // Call kv.BatchGet for this batch + for kv, err := range d.kv.BatchGet(ctx, dataSection, stringKeys) { + if err != nil { + yield(DataObj{}, err) + return + } + + // Look up the original DataKey + dataKey, ok := keyMap[kv.Key] + if !ok { + yield(DataObj{}, fmt.Errorf("unexpected key in batch response: %s", kv.Key)) + return + } + + // Yield the DataObj + if !yield(DataObj{ + Key: dataKey, + Value: kv.Value, + }, nil) { + return + } + } + } + } +} + func (d *dataStore) Save(ctx context.Context, key DataKey, value io.Reader) error { if err := key.Validate(); err != nil { return fmt.Errorf("invalid data key: %w", err) diff --git a/pkg/storage/unified/resource/datastore_test.go b/pkg/storage/unified/resource/datastore_test.go index 9b405847a9a..39d995b810a 100644 --- a/pkg/storage/unified/resource/datastore_test.go +++ b/pkg/storage/unified/resource/datastore_test.go @@ -2949,3 +2949,148 @@ func TestDataStore_getGroupResources(t *testing.T) { require.True(t, foundCombinations[expected], "Expected combination not found: %s", expected) } } + +func TestDataStore_BatchGet(t *testing.T) { + ds := setupTestDataStore(t) + ctx := context.Background() + + t.Run("batch get multiple existing keys", func(t *testing.T) { + // Create test data + keys := make([]DataKey, 5) + expectedContent := make(map[string]string) + + for i := 0; i < 5; i++ { + rv := node.Generate().Int64() + keys[i] = DataKey{ + Namespace: "test-namespace", + Group: "test-group", + Resource: "test-resource", + Name: fmt.Sprintf("test-name-%d", i), + ResourceVersion: rv, + Action: DataActionCreated, + Folder: "test-folder", + } + content := fmt.Sprintf("test-value-%d", i) + expectedContent[keys[i].Name] = content + err := ds.Save(ctx, keys[i], bytes.NewReader([]byte(content))) + require.NoError(t, err) + } + + // Batch get all keys + results := make([]DataObj, 0, 5) + for obj, err := range ds.BatchGet(ctx, keys) { + require.NoError(t, err) + results = append(results, obj) + } + + // Verify all keys were returned + require.Len(t, results, 5) + + // Verify content matches + for _, result := range results { + resultBytes, err := io.ReadAll(result.Value) + require.NoError(t, err) + expectedValue, ok := expectedContent[result.Key.Name] + require.True(t, ok, "Unexpected key in results: %s", result.Key.Name) + require.Equal(t, expectedValue, string(resultBytes)) + } + }) + + t.Run("batch get with some non-existent keys", func(t *testing.T) { + // Create 3 existing keys + existingKeys := make([]DataKey, 3) + for i := 0; i < 3; i++ { + rv := node.Generate().Int64() + existingKeys[i] = DataKey{ + Namespace: "test-namespace", + Group: "test-group", + Resource: "test-resource", + Name: fmt.Sprintf("existing-%d", i), + ResourceVersion: rv, + Action: DataActionCreated, + Folder: "test-folder", + } + err := ds.Save(ctx, existingKeys[i], bytes.NewReader([]byte(fmt.Sprintf("value-%d", i)))) + require.NoError(t, err) + } + + // Create 2 non-existent keys (not saved to datastore) + nonExistentKeys := make([]DataKey, 2) + for i := 0; i < 2; i++ { + rv := node.Generate().Int64() + nonExistentKeys[i] = DataKey{ + Namespace: "test-namespace", + Group: "test-group", + Resource: "test-resource", + Name: fmt.Sprintf("non-existent-%d", i), + ResourceVersion: rv, + Action: DataActionCreated, + Folder: "test-folder", + } + } + + // Combine existing and non-existent keys + allKeys := append(existingKeys, nonExistentKeys...) + + // Batch get all keys + results := make([]DataObj, 0, 3) + for obj, err := range ds.BatchGet(ctx, allKeys) { + require.NoError(t, err) + results = append(results, obj) + } + + // Should only return the 3 existing keys + require.Len(t, results, 3) + + // Verify only existing keys are returned + for _, result := range results { + require.Contains(t, result.Key.Name, "existing-") + // Verify content + resultBytes, err := io.ReadAll(result.Value) + require.NoError(t, err) + require.NotEmpty(t, resultBytes) + } + }) + + t.Run("batch get with large number of keys to test batching", func(t *testing.T) { + numKeys := 150 + keys := make([]DataKey, numKeys) + expectedContent := make(map[string]string) + + for i := 0; i < numKeys; i++ { + rv := node.Generate().Int64() + keys[i] = DataKey{ + Namespace: "batch-test", + Group: "test-group", + Resource: "test-resource", + Name: fmt.Sprintf("item-%d", i), + ResourceVersion: rv, + Action: DataActionCreated, + Folder: "test-folder", + } + content := fmt.Sprintf("content-%d", i) + expectedContent[keys[i].Name] = content + err := ds.Save(ctx, keys[i], bytes.NewReader([]byte(content))) + require.NoError(t, err) + } + + // Batch get all keys + results := make([]DataObj, 0, numKeys) + for obj, err := range ds.BatchGet(ctx, keys) { + require.NoError(t, err) + results = append(results, obj) + } + + // Verify all keys were returned + require.Len(t, results, numKeys) + + // Verify content matches for all keys + for _, result := range results { + resultBytes, err := io.ReadAll(result.Value) + require.NoError(t, err) + expectedValue, ok := expectedContent[result.Key.Name] + require.True(t, ok, "Unexpected key in results: %s", result.Key.Name) + require.Equal(t, expectedValue, string(resultBytes)) + } + }) +} diff --git a/pkg/storage/unified/resource/storage_backend.go b/pkg/storage/unified/resource/storage_backend.go index 0ec839b5a3f..d6afaf43638 100644 --- a/pkg/storage/unified/resource/storage_backend.go +++ b/pkg/storage/unified/resource/storage_backend.go @@ -364,14 +364,14 @@ func (k *kvStorageBackend) ListIterator(ctx context.Context, req *resourcepb.Lis break } } + // Create pull-style iterator from BatchGet + next, stop := iter.Pull2(k.dataStore.BatchGet(ctx, keys)) + defer stop() + iter := kvListIterator{ - keys: keys, - currentIndex: -1, - ctx: ctx, - listRV: listRV, - offset: offset, - limit: req.Limit + 1, // TODO: for now we need at least one more item. Fix the caller - dataStore: k.dataStore, + listRV: listRV, + offset: offset, + next: next, } err := cb(&iter) if err != nil { @@ -383,52 +383,44 @@ func (k *kvStorageBackend) ListIterator(ctx context.Context, req *resourcepb.Lis // kvListIterator implements ListIterator for KV storage type kvListIterator struct { - ctx context.Context - keys []DataKey - currentIndex int - dataStore *dataStore - listRV int64 - offset int64 - limit int64 + listRV int64 + offset int64 - // current - rv int64 - err error - value []byte + // pull-style iterator + next func() (DataObj, error, bool) + + // current item state + currentDataObj *DataObj + value []byte + err error } func (i *kvListIterator) Next() bool { - i.currentIndex++ - - if i.currentIndex >= len(i.keys) { + // Pull next item from the iterator + dataObj, err, ok := i.next() + if !ok { return false } - - if int64(i.currentIndex) >= i.limit { - return false - } - - i.rv, i.err = i.keys[i.currentIndex].ResourceVersion, nil - - data, err := i.dataStore.Get(i.ctx, i.keys[i.currentIndex]) if err != nil { i.err = err return false } - i.value, i.err = readAndClose(data) - if i.err != nil { + i.currentDataObj = &dataObj + + i.value, err = readAndClose(dataObj.Value) + if err != nil { + i.err = err return false } - // increment the offset i.offset++ return true } func (i *kvListIterator) Error() error { - return nil + return i.err } func (i *kvListIterator) ContinueToken() string { @@ -439,19 +431,31 @@ func (i *kvListIterator) ContinueToken() string { } func (i *kvListIterator) ResourceVersion() int64 { - return i.rv + if i.currentDataObj != nil { + return i.currentDataObj.Key.ResourceVersion + } + return 0 } func (i *kvListIterator) Namespace() string { - return i.keys[i.currentIndex].Namespace + if i.currentDataObj != nil { + return i.currentDataObj.Key.Namespace + } + return "" } func (i *kvListIterator) Name() string { - return i.keys[i.currentIndex].Name + if i.currentDataObj != nil { + return i.currentDataObj.Key.Name + } + return "" } func (i *kvListIterator) Folder() string { - return i.keys[i.currentIndex].Folder + if i.currentDataObj != nil { + return i.currentDataObj.Key.Folder + } + return "" } func (i *kvListIterator) Value() []byte { @@ -825,13 +829,14 @@ func (k *kvStorageBackend) ListHistory(ctx context.Context, req *resourcepb.List // Pagination: filter out items up to and including lastSeenRV pagedKeys := applyPagination(filteredKeys, lastSeenRV, sortAscending) + // Create pull-style iterator from BatchGet + next, stop := iter.Pull2(k.dataStore.BatchGet(ctx, pagedKeys)) + defer stop() + iter := kvHistoryIterator{ - keys: pagedKeys, - currentIndex: -1, - ctx: ctx, listRV: listRV, sortAscending: sortAscending, - dataStore: k.dataStore, + next: next, } err := fn(&iter) @@ -889,14 +894,15 @@ func (k *kvStorageBackend) processTrashEntries(ctx context.Context, req *resourc // Pagination: filter out items up to and including lastSeenRV pagedKeys := applyPagination(filteredKeys, lastSeenRV, sortAscending) + // Create pull-style iterator from BatchGet + next, stop := iter.Pull2(k.dataStore.BatchGet(ctx, pagedKeys)) + defer stop() + iter := kvHistoryIterator{ - keys: pagedKeys, - currentIndex: -1, - ctx: ctx, listRV: listRV, sortAscending: sortAscending, - dataStore: k.dataStore, skipProvisioned: true, + next: next, } err = fn(&iter) @@ -909,43 +915,36 @@ func (k *kvStorageBackend) processTrashEntries(ctx context.Context, req *resourc // kvHistoryIterator implements ListIterator for KV storage history type kvHistoryIterator struct { - ctx context.Context - keys []DataKey - currentIndex int listRV int64 sortAscending bool skipProvisioned bool - dataStore *dataStore - // current - rv int64 - err error - value []byte - folder string + // pull-style iterator + next func() (DataObj, error, bool) + + // current item state + currentDataObj *DataObj + value []byte + folder string + err error } func (i *kvHistoryIterator) Next() bool { - i.currentIndex++ - - if i.currentIndex >= len(i.keys) { + // Pull next item from the iterator + dataObj, err, ok := i.next() + if !ok { return false } - - key := i.keys[i.currentIndex] - i.rv = key.ResourceVersion - - // Read the value from the ReadCloser - data, err := i.dataStore.Get(i.ctx, key) if err != nil { i.err = err return false } - if data == nil { - i.err = fmt.Errorf("data is nil") - return false - } - i.value, i.err = readAndClose(data) - if i.err != nil { + + i.currentDataObj = &dataObj + + i.value, err = readAndClose(dataObj.Value) + if err != nil { + i.err = err return false } @@ -963,7 +962,6 @@ func (i *kvHistoryIterator) Next() bool { return false } i.folder = meta.GetFolder() - i.err = nil // if the resource is provisioned and we are skipping provisioned resources, continue onto the next one if i.skipProvisioned && meta.GetAnnotation(utils.AnnoKeyManagerKind) != "" { @@ -978,31 +976,35 @@ func (i *kvHistoryIterator) Error() error { } func (i *kvHistoryIterator) ContinueToken() string { - if i.currentIndex < 0 || i.currentIndex >= len(i.keys) { + if i.currentDataObj == nil { return "" } + rv := i.currentDataObj.Key.ResourceVersion token := ContinueToken{ - StartOffset: i.rv, - ResourceVersion: i.keys[i.currentIndex].ResourceVersion, + StartOffset: rv, + ResourceVersion: rv, SortAscending: i.sortAscending, } return token.String() } func (i *kvHistoryIterator) ResourceVersion() int64 { - return i.rv + if i.currentDataObj != nil { + return i.currentDataObj.Key.ResourceVersion + } + return 0 } func (i *kvHistoryIterator) Namespace() string { - if i.currentIndex >= 0 && i.currentIndex < len(i.keys) { - return i.keys[i.currentIndex].Namespace + if i.currentDataObj != nil { + return i.currentDataObj.Key.Namespace } return "" } func (i *kvHistoryIterator) Name() string { - if i.currentIndex >= 0 && i.currentIndex < len(i.keys) { - return i.keys[i.currentIndex].Name + if i.currentDataObj != nil { + return i.currentDataObj.Key.Name } return "" }