public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

kvstore: update storage backend to use batch get (#112704)

Browse Source 
* consume batchget

* remove rv

* remove un-necessary limit

* do not pass stop to the iterators

* ensure we return the error
This commit is contained in:
Georges Chaudy
2025-10-21 15:42:07 +00:00
committed by GitHub
parent 0a2564c4d5
commit eca379e98e
3 changed files with 284 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/storage/unified/resource/datastore.go
+59
View File
@@ -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)
pkg/storage/unified/resource/datastore_test.go
+145
View File
@@ -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))
}
})
}
pkg/storage/unified/resource/storage_backend.go
+80 -78
View File
@@ -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 ""
}