Merge pull request #1082 from fabriziosestito/fix/storage-watcher-race-condition

fix(storage): avoid dropping ADD events after a concurrent store delete

Author: fabriziosestito

internal/storage/store_test.go

@@ -2,6 +2,7 @@ package storage
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -466,6 +467,115 @@ func (suite *storeTestSuite) TestWatchList() {
 	suite.NotEmpty(bookmarkObj.ResourceVersion)
 }
 
+// When a non-Deleted event arrives for an object the store no longer holds, the watcher must still broadcast the payload.
+func (suite *storeTestSuite) TestHandleMessageNotFoundStillBroadcasts() {
+	ctx, cancel := context.WithCancel(suite.T().Context())
+	defer cancel()
+
+	w, err := suite.broadcaster.Watch()
+	suite.Require().NoError(err)
+	defer w.Stop()
+
+	go suite.watcher.Start(ctx)
+	suite.Require().NoError(suite.nc.Flush())
+
+	sbom := &storagev1alpha1.SBOM{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "ghost",
+			Namespace: "default",
+		},
+	}
+	sbomBytes, err := json.Marshal(sbom)
+	suite.Require().NoError(err)
+
+	payload, err := json.Marshal(event{
+		EventType: watch.Added,
+		Object:    runtime.RawExtension{Raw: sbomBytes},
+	})
+	suite.Require().NoError(err)
+
+	suite.Require().NoError(suite.nc.Publish("watch.sboms", payload))
+	suite.Require().NoError(suite.nc.Flush())
+
+	events := mustReadEvents(suite.T(), w, 1)
+	suite.Equal(watch.Added, events[0].Type)
+
+	got, ok := events[0].Object.(*storagev1alpha1.SBOM)
+	suite.Require().True(ok)
+	suite.Equal("ghost", got.Name)
+	suite.Equal("default", got.Namespace)
+}
+
+// When the store holds a different object at the same namespace/name, the watcher must broadcast the payload, not the refetched object.
+func (suite *storeTestSuite) TestHandleMessageUIDMismatchBroadcastsPayload() {
+	ctx, cancel := context.WithCancel(suite.T().Context())
+	defer cancel()
+
+	w, err := suite.broadcaster.Watch()
+	suite.Require().NoError(err)
+	defer w.Stop()
+
+	go suite.watcher.Start(ctx)
+	suite.Require().NoError(suite.nc.Flush())
+
+	// Populate the store at default/collide. The stale event published next will carry a different UID for the same namespace/name.
+	storedUID := types.UID("stored-uid")
+	stored := &storagev1alpha1.SBOM{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "collide",
+			Namespace: "default",
+			UID:       storedUID,
+		},
+		SPDX: runtime.RawExtension{Raw: []byte(`{"stored": true}`)},
+	}
+	key := keyPrefix + "/default/collide"
+	err = suite.store.Create(ctx, key, stored, &storagev1alpha1.SBOM{}, 0)
+	suite.Require().NoError(err)
+
+	// Drain the ADDED event produced by the Create above
+	created := mustReadEvents(suite.T(), w, 1)
+	suite.Equal(watch.Added, created[0].Type)
+	createdSBOM, ok := created[0].Object.(*storagev1alpha1.SBOM)
+	suite.Require().True(ok)
+	suite.Equal(storedUID, createdSBOM.UID)
+
+	// Publish a stale ADDED carrying a different UID at the same namespace/name
+	staleUID := types.UID("stale-uid")
+	suite.Require().NotEqual(storedUID, staleUID)
+
+	stale := &storagev1alpha1.SBOM{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "collide",
+			Namespace: "default",
+			UID:       staleUID,
+		},
+	}
+	staleBytes, err := json.Marshal(stale)
+	suite.Require().NoError(err)
+
+	payload, err := json.Marshal(event{
+		EventType: watch.Added,
+		Object:    runtime.RawExtension{Raw: staleBytes},
+	})
+	suite.Require().NoError(err)
+
+	suite.Require().NoError(suite.nc.Publish("watch.sboms", payload))
+	suite.Require().NoError(suite.nc.Flush())
+
+	events := mustReadEvents(suite.T(), w, 1)
+	suite.Equal(watch.Added, events[0].Type)
+
+	got, ok := events[0].Object.(*storagev1alpha1.SBOM)
+	suite.Require().True(ok)
+	suite.Equal(staleUID, got.UID, "expected payload UID, not stored UID")
+
+	// Stored object should still be intact and retrievable with its own UID
+	fetched := &storagev1alpha1.SBOM{}
+	err = suite.store.Get(ctx, key, k8sstorage.GetOptions{}, fetched)
+	suite.Require().NoError(err)
+	suite.Equal(storedUID, fetched.UID)
+}
+
 func (suite *storeTestSuite) TestGetList() {
 	key := keyPrefix + "/default"
 	sbom1 := storagev1alpha1.SBOM{

internal/storage/watcher.go

@@ -90,32 +90,18 @@ func (w *natsWatcher) handleMessage(ctx context.Context, msg *nats.Msg) error {
 	}
 
 	obj := w.store.newFunc()
-	err := json.Unmarshal(payload.Object.Raw, obj)
-	if err != nil {
+	if err := json.Unmarshal(payload.Object.Raw, obj); err != nil {
 		return fmt.Errorf("failed to decode object: %w", err)
 	}
 
-	// NOTE: For deleted events, we broadcast the event without re-fetching the object from the store.
-	// This is because the object was deleted, and we are not able to fetch it again.
-	// The object might not be the complete object, as the transforum function may have been applied.
+	// For deleted events broadcast the payload directly since the store no longer has it.
+	// Otherwise rehydrate fields stripped by the publish-side transform (e.g. SBOM.SPDX) from the store's current state.
 	if payload.EventType != watch.Deleted {
-		metaAccessor, err := meta.Accessor(obj)
+		rehydrated, err := w.rehydrate(ctx, obj)
 		if err != nil {
-			return fmt.Errorf("failed to get meta accessor: %w", err)
-		}
-		key := fmt.Sprintf("%s/%s/%s/%s", storagev1alpha1.GroupName, w.resource, metaAccessor.GetNamespace(), metaAccessor.GetName())
-
-		if err := w.store.Get(ctx, key, storage.GetOptions{}, obj); err != nil {
-			if storage.IsNotFound(err) {
-				// Object not found, possibly deleted after the event was sent.
-				w.logger.DebugContext(ctx, "Object not found in store while handling message, skipping",
-					"key", key,
-				)
-
-				return nil
-			}
-			return fmt.Errorf("failed to get object from store while handling message: %w", err)
+			return err
 		}
+		obj = rehydrated
 	}
 
 	if err := w.watchBroadcaster.Action(payload.EventType, obj); err != nil {
@@ -129,6 +115,43 @@ func (w *natsWatcher) handleMessage(ctx context.Context, msg *nats.Msg) error {
 	return nil
 }
 
+// rehydrate returns the stored object matching the payload, or the payload itself when the store does not hold a matching object.
+// Falling back to the payload (rather than dropping the event) keeps downstream watchers like the GC dependency graph builder consistent. The following DELETED event reconciles client state.
+func (w *natsWatcher) rehydrate(ctx context.Context, payloadObj runtime.Object) (runtime.Object, error) {
+	payloadAccessor, err := meta.Accessor(payloadObj)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get meta accessor: %w", err)
+	}
+	key := fmt.Sprintf("%s/%s/%s/%s", storagev1alpha1.GroupName, w.resource, payloadAccessor.GetNamespace(), payloadAccessor.GetName())
+
+	fetched := w.store.newFunc()
+	if err := w.store.Get(ctx, key, storage.GetOptions{}, fetched); err != nil {
+		if !storage.IsNotFound(err) {
+			return nil, fmt.Errorf("failed to get object from store while handling message: %w", err)
+		}
+		w.logger.DebugContext(ctx, "Object not found in store while handling message; broadcasting payload",
+			"key", key,
+		)
+		return payloadObj, nil
+	}
+
+	fetchedAccessor, err := meta.Accessor(fetched)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get meta accessor for fetched object: %w", err)
+	}
+	if fetchedAccessor.GetUID() != payloadAccessor.GetUID() {
+		// Same namespace/name, different object: it was deleted and recreated before we got here.
+		w.logger.DebugContext(ctx, "Stored object UID does not match payload; broadcasting payload",
+			"key", key,
+			"payloadUID", payloadAccessor.GetUID(),
+			"storedUID", fetchedAccessor.GetUID(),
+		)
+		return payloadObj, nil
+	}
+
+	return fetched, nil
+}
+
 // natsBroadcaster broadcasts watch events using NATS.
 type natsBroadcaster struct {
 	nc               *nats.Conn