diff --git a/pkg/storage/unified/resource/bulk.go b/pkg/storage/unified/resource/bulk.go index 1ac0ff2032f..e665485dc3c 100644 --- a/pkg/storage/unified/resource/bulk.go +++ b/pkg/storage/unified/resource/bulk.go @@ -6,10 +6,14 @@ import ( "fmt" "io" "net/http" + "sync" + "time" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" "google.golang.org/grpc/metadata" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" authlib "github.com/grafana/authlib/types" @@ -328,3 +332,84 @@ func (b *batchRunner) RollbackRequested() bool { } return false } + +type bulkRV struct { + max int64 + counter int64 +} + +// Used when executing a bulk import so that we can generate snowflake RVs in the past +func newBulkRV() *bulkRV { + t := snowflakeFromTime(time.Now()) + return &bulkRV{ + max: t, + counter: 0, + } +} + +func (x *bulkRV) next(obj metav1.Object) int64 { + ts := snowflakeFromTime(obj.GetCreationTimestamp().Time) + anno := obj.GetAnnotations() + if anno != nil { + v := anno[utils.AnnoKeyUpdatedTimestamp] + t, err := time.Parse(time.RFC3339, v) + if err == nil { + ts = snowflakeFromTime(t) + } + } + if ts > x.max || ts < 0 { + ts = x.max + } + + x.counter++ + return ts + x.counter +} + +type BulkLock struct { + running map[string]bool + mu sync.Mutex +} + +func NewBulkLock() *BulkLock { + return &BulkLock{ + running: make(map[string]bool), + } +} + +func (x *BulkLock) Start(keys []*resourcepb.ResourceKey) error { + x.mu.Lock() + defer x.mu.Unlock() + + // First verify that it is not already running + ids := make([]string, len(keys)) + for i, k := range keys { + id := NSGR(k) + if x.running[id] { + return &apierrors.StatusError{ErrStatus: metav1.Status{ + Code: http.StatusPreconditionFailed, + Message: "bulk export is already running", + }} + } + ids[i] = id + } + + // Then add the keys to the lock + for _, k := range ids { + x.running[k] = true + } + return nil +} + +func (x *BulkLock) Finish(keys []*resourcepb.ResourceKey) { + x.mu.Lock() + defer x.mu.Unlock() + for _, k := range keys { + delete(x.running, NSGR(k)) + } +} + +func (x *BulkLock) Active() bool { + x.mu.Lock() + defer x.mu.Unlock() + return len(x.running) > 0 +} diff --git a/pkg/storage/unified/resource/datastore.go b/pkg/storage/unified/resource/datastore.go index 82f1e1e59d5..8a930492420 100644 --- a/pkg/storage/unified/resource/datastore.go +++ b/pkg/storage/unified/resource/datastore.go @@ -537,6 +537,27 @@ func (d *dataStore) Delete(ctx context.Context, key DataKey) error { return d.kv.Delete(ctx, dataSection, key.String()) } +func (n *dataStore) batchDelete(ctx context.Context, keys []DataKey) error { + for len(keys) > 0 { + batch := keys + if len(batch) > dataBatchSize { + batch = batch[:dataBatchSize] + } + + keys = keys[len(batch):] + stringKeys := make([]string, len(batch)) + for _, dataKey := range batch { + stringKeys = append(stringKeys, dataKey.String()) + } + + if err := n.kv.BatchDelete(ctx, dataSection, stringKeys); err != nil { + return err + } + } + + return nil +} + // ParseKey parses a string key into a DataKey struct func ParseKey(key string) (DataKey, error) { parts := strings.Split(key, "/") diff --git a/pkg/storage/unified/resource/datastore_test.go b/pkg/storage/unified/resource/datastore_test.go index fbfb828f133..8f167c2e16e 100644 --- a/pkg/storage/unified/resource/datastore_test.go +++ b/pkg/storage/unified/resource/datastore_test.go @@ -2950,6 +2950,42 @@ func TestDataStore_getGroupResources(t *testing.T) { } } +func TestDataStore_BatchDelete(t *testing.T) { + ds := setupTestDataStore(t) + ctx := context.Background() + + keys := make([]DataKey, 95) + for i := 0; i < 95; 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) + err := ds.Save(ctx, keys[i], bytes.NewReader([]byte(content))) + require.NoError(t, err) + } + + err := ds.batchDelete(ctx, keys) + require.NoError(t, err) + + // Verify all events were deleted + for i := 0; i < 95; i++ { + _, err := ds.Get(ctx, DataKey{ + Namespace: "test-namespace", + Group: "test-group", + Resource: "test-resource", + Name: fmt.Sprintf("test-name-%d", i), + }) + require.Error(t, err, "Resource should have been deleted") + } +} + func TestDataStore_BatchGet(t *testing.T) { ds := setupTestDataStore(t) ctx := context.Background() diff --git a/pkg/storage/unified/resource/storage_backend.go b/pkg/storage/unified/resource/storage_backend.go index 0f65867f351..0de97b0355e 100644 --- a/pkg/storage/unified/resource/storage_backend.go +++ b/pkg/storage/unified/resource/storage_backend.go @@ -18,6 +18,7 @@ import ( "github.com/prometheus/client_golang/prometheus" "go.opentelemetry.io/otel/trace" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "github.com/grafana/grafana/pkg/apimachinery/utils" "github.com/grafana/grafana/pkg/storage/unified/resourcepb" @@ -54,6 +55,7 @@ func convertEmptyToClusterNamespace(namespace string, withExperimentalClusterSco type kvStorageBackend struct { snowflake *snowflake.Node kv KV + bulkLock *BulkLock dataStore *dataStore eventStore *eventStore notifier *notifier @@ -102,6 +104,7 @@ func NewKVStorageBackend(opts KVBackendOptions) (StorageBackend, error) { backend := &kvStorageBackend{ kv: kv, + bulkLock: NewBulkLock(), dataStore: newDataStore(kv), eventStore: eventStore, notifier: newNotifier(eventStore, notifierOptions{}), @@ -1236,6 +1239,187 @@ func (k *kvStorageBackend) GetResourceLastImportTimes(ctx context.Context) iter. } } +func (b *kvStorageBackend) ProcessBulk(ctx context.Context, setting BulkSettings, iter BulkRequestIterator) *resourcepb.BulkResponse { + // TODO cross-node lock + err := b.bulkLock.Start(setting.Collection) + if err != nil { + return &resourcepb.BulkResponse{ + Error: AsErrorResult(err), + } + } + defer b.bulkLock.Finish(setting.Collection) + + bulkRvGenerator := newBulkRV() + summaries := make(map[string]*resourcepb.BulkResponse_Summary, len(setting.Collection)) + rsp := &resourcepb.BulkResponse{} + + if setting.RebuildCollection { + for _, key := range setting.Collection { + events := make([]string, 0) + for evtKeyStr, err := range b.eventStore.ListKeysSince(ctx, 1) { + if err != nil { + b.log.Error("failed to list event: %s", err) + return rsp + } + + evtKey, err := ParseEventKey(evtKeyStr) + if err != nil { + b.log.Error("error parsing event key: %s", err) + return rsp + } + + if evtKey.Group != key.Group || evtKey.Resource != key.Resource || evtKey.Namespace != key.Namespace { + continue + } + + events = append(events, evtKeyStr) + } + + if err := b.eventStore.batchDelete(ctx, events); err != nil { + b.log.Error("failed to delete events: %s", err) + return rsp + } + + historyKeys := make([]DataKey, 0) + + for dataKey, err := range b.dataStore.Keys(ctx, ListRequestKey{ + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + }, SortOrderAsc) { + if err != nil { + b.log.Error("failed to list collection before delete: %s", err) + return rsp + } + + historyKeys = append(historyKeys, dataKey) + } + + previousCount := int64(len(historyKeys)) + if err := b.dataStore.batchDelete(ctx, historyKeys); err != nil { + b.log.Error("failed to delete collection: %s", err) + return rsp + } + summaries[NSGR(key)] = &resourcepb.BulkResponse_Summary{ + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + PreviousCount: previousCount, + } + } + } else { + for _, key := range setting.Collection { + summaries[NSGR(key)] = &resourcepb.BulkResponse_Summary{ + Namespace: key.Namespace, + Group: key.Group, + Resource: key.Resource, + } + } + } + + obj := &unstructured.Unstructured{} + + saved := make([]DataKey, 0) + rollback := func() { + // we don't have transactions in the kv store, so we simply delete everything we created + err = b.dataStore.batchDelete(ctx, saved) + if err != nil { + b.log.Error("failed to delete during rollback: %s", err) + } + } + + for iter.Next() { + if iter.RollbackRequested() { + rollback() + break + } + + req := iter.Request() + if req == nil { + rollback() + rsp.Error = AsErrorResult(fmt.Errorf("missing request")) + break + } + + rsp.Processed++ + + var action DataAction + switch resourcepb.WatchEvent_Type(req.Action) { + case resourcepb.WatchEvent_ADDED: + action = DataActionCreated + // Check if resource already exists for create operations + _, err := b.dataStore.GetLatestResourceKey(ctx, GetRequestKey{ + Group: req.Key.Group, + Resource: req.Key.Resource, + Namespace: req.Key.Namespace, + Name: req.Key.Name, + }) + if err == nil { + rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{ + Key: req.Key, + Action: req.Action, + Error: "resource already exists", + }) + continue + } + if !errors.Is(err, ErrNotFound) { + rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{ + Key: req.Key, + Action: req.Action, + Error: fmt.Sprintf("failed to check if resource exists: %s", err), + }) + continue + } + case resourcepb.WatchEvent_MODIFIED: + action = DataActionUpdated + case resourcepb.WatchEvent_DELETED: + action = DataActionDeleted + default: + rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{ + Key: req.Key, + Action: req.Action, + Error: "invalid event type", + }) + continue + } + + err := obj.UnmarshalJSON(req.Value) + if err != nil { + rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{ + Key: req.Key, + Action: req.Action, + Error: "unable to unmarshal json", + }) + continue + } + + dataKey := DataKey{ + Group: req.Key.Group, + Resource: req.Key.Resource, + Namespace: req.Key.Namespace, + Name: req.Key.Name, + ResourceVersion: bulkRvGenerator.next(obj), + Action: action, + Folder: req.Folder, + } + err = b.dataStore.Save(ctx, dataKey, bytes.NewReader(req.Value)) + if err != nil { + rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{ + Key: req.Key, + Action: req.Action, + Error: fmt.Sprintf("failed to save resource: %s", err), + }) + continue + } + + saved = append(saved, dataKey) + } + + // TODO update last import time + + return rsp +} + // readAndClose reads all data from a ReadCloser and ensures it's closed, // combining any errors from both operations. func readAndClose(r io.ReadCloser) ([]byte, error) {