Unistore: fix DualWriter context cancelation on mode 1 (#93377)

* Unistore: fix DualWriter context cancelation on mode 1

Signed-off-by: Maicon Costa <maiconscosta@gmail.com>

pkg/apiserver/rest/dualwriter_mode1.go

@@ -57,30 +57,39 @@ func (d *DualWriterMode1) Create(ctx context.Context, original runtime.Object, c
 
 	createdCopy := created.DeepCopyObject()
 
-	go func(createdCopy runtime.Object) {
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage create timeout"))
-		defer cancel()
-
-		if err := enrichLegacyObject(original, createdCopy); err != nil {
-			cancel()
-		}
-
-		startStorage := time.Now()
-		storageObj, errObjectSt := d.Storage.Create(ctx, createdCopy, createValidation, options)
-		d.recordStorageDuration(errObjectSt != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			cancel()
-		}
-		areEqual := Compare(storageObj, createdCopy)
-		d.recordOutcome(mode1Str, getName(createdCopy), areEqual, method)
-		if !areEqual {
-			log.Info("object from legacy and storage are not equal")
-		}
-	}(createdCopy)
+	//nolint:errcheck
+	go d.createOnUnifiedStorage(ctx, original, createValidation, createdCopy, options)
 
 	return created, err
 }
 
+func (d *DualWriterMode1) createOnUnifiedStorage(ctx context.Context, original runtime.Object, createValidation rest.ValidateObjectFunc, createdCopy runtime.Object, options *metav1.CreateOptions) error {
+	var method = "create"
+	log := d.Log.WithValues("method", method)
+
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage create timeout"))
+	defer cancel()
+
+	if err := enrichLegacyObject(original, createdCopy); err != nil {
+		cancel()
+	}
+
+	startStorage := time.Now()
+	storageObj, errObjectSt := d.Storage.Create(ctx, createdCopy, createValidation, options)
+	d.recordStorageDuration(errObjectSt != nil, mode1Str, d.resource, method, startStorage)
+	if errObjectSt != nil {
+		cancel()
+	}
+	areEqual := Compare(storageObj, createdCopy)
+	d.recordOutcome(mode1Str, getName(createdCopy), areEqual, method)
+	if !areEqual {
+		log.Info("object from legacy and storage are not equal")
+	}
+
+	return errObjectSt
+}
+
 // Get overrides the behavior of the generic DualWriter and reads only from LegacyStorage.
 func (d *DualWriterMode1) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
 	var method = "get"
@@ -94,27 +103,36 @@ func (d *DualWriterMode1) Get(ctx context.Context, name string, options *metav1.
 	}
 	d.recordLegacyDuration(errLegacy != nil, mode1Str, d.resource, method, startLegacy)
 
-	go func(res runtime.Object) {
-		startStorage := time.Now()
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage get timeout"))
-		defer cancel()
-		storageObj, err := d.Storage.Get(ctx, name, options)
-		d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			log.Error(err, "unable to get object in storage")
-			cancel()
-		}
-
-		areEqual := Compare(storageObj, res)
-		d.recordOutcome(mode1Str, name, areEqual, method)
-		if !areEqual {
-			log.WithValues("name", name).Info("object from legacy and storage are not equal")
-		}
-	}(res)
+	//nolint:errcheck
+	go d.getFromUnifiedStorage(ctx, res, name, options)
 
 	return res, errLegacy
 }
 
+func (d *DualWriterMode1) getFromUnifiedStorage(ctx context.Context, res runtime.Object, name string, options *metav1.GetOptions) error {
+	var method = "get"
+	log := d.Log.WithValues("method", method, "name", name)
+
+	startStorage := time.Now()
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage get timeout"))
+	defer cancel()
+	storageObj, err := d.Storage.Get(ctx, name, options)
+	d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
+	if err != nil {
+		log.Error(err, "unable to get object in storage")
+		cancel()
+	}
+
+	areEqual := Compare(storageObj, res)
+	d.recordOutcome(mode1Str, name, areEqual, method)
+	if !areEqual {
+		log.WithValues("name", name).Info("object from legacy and storage are not equal")
+	}
+
+	return err
+}
+
 // List overrides the behavior of the generic DualWriter and reads only from LegacyStorage.
 func (d *DualWriterMode1) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
 	var method = "list"
@@ -128,25 +146,35 @@ func (d *DualWriterMode1) List(ctx context.Context, options *metainternalversion
 	}
 	d.recordLegacyDuration(errLegacy != nil, mode1Str, d.resource, method, startLegacy)
 
-	go func(res runtime.Object) {
-		startStorage := time.Now()
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage list timeout"))
-		defer cancel()
-		storageObj, err := d.Storage.List(ctx, options)
-		d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			cancel()
-		}
-		areEqual := Compare(storageObj, res)
-		d.recordOutcome(mode1Str, getName(res), areEqual, method)
-		if !areEqual {
-			log.Info("object from legacy and storage are not equal")
-		}
-	}(res)
+	//nolint:errcheck
+	go d.listFromUnifiedStorage(ctx, options, res)
 
 	return res, errLegacy
 }
 
+func (d *DualWriterMode1) listFromUnifiedStorage(ctx context.Context, options *metainternalversion.ListOptions, res runtime.Object) error {
+	var method = "list"
+	log := d.Log.WithValues("resourceVersion", options.ResourceVersion, "method", method)
+
+	startStorage := time.Now()
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage list timeout"))
+	defer cancel()
+	storageObj, err := d.Storage.List(ctx, options)
+	d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
+	if err != nil {
+		log.Error(err, "unable to list objects from unified storage")
+		cancel()
+	}
+	areEqual := Compare(storageObj, res)
+	d.recordOutcome(mode1Str, getName(res), areEqual, method)
+	if !areEqual {
+		log.Info("object from legacy and storage are not equal")
+	}
+
+	return err
+}
+
 func (d *DualWriterMode1) Delete(ctx context.Context, name string, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions) (runtime.Object, bool, error) {
 	var method = "delete"
 	log := d.Log.WithValues("name", name, "method", method, "name", name)
@@ -161,25 +189,34 @@ func (d *DualWriterMode1) Delete(ctx context.Context, name string, deleteValidat
 	}
 	d.recordLegacyDuration(false, mode1Str, name, method, startLegacy)
 
-	go func(res runtime.Object) {
-		startStorage := time.Now()
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage delete timeout"))
-		defer cancel()
-		storageObj, _, err := d.Storage.Delete(ctx, name, deleteValidation, options)
-		d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			cancel()
-		}
-		areEqual := Compare(storageObj, res)
-		d.recordOutcome(mode1Str, name, areEqual, method)
-		if !areEqual {
-			log.Info("object from legacy and storage are not equal")
-		}
-	}(res)
+	//nolint:errcheck
+	go d.deleteFromUnifiedStorage(ctx, res, name, deleteValidation, options)
 
 	return res, async, err
 }
 
+func (d *DualWriterMode1) deleteFromUnifiedStorage(ctx context.Context, res runtime.Object, name string, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions) error {
+	var method = "delete"
+	log := d.Log.WithValues("name", name, "method", method, "name", name)
+
+	startStorage := time.Now()
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage delete timeout"))
+	defer cancel()
+	storageObj, _, err := d.Storage.Delete(ctx, name, deleteValidation, options)
+	d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
+	if err != nil {
+		cancel()
+	}
+	areEqual := Compare(storageObj, res)
+	d.recordOutcome(mode1Str, name, areEqual, method)
+	if !areEqual {
+		log.Info("object from legacy and storage are not equal")
+	}
+
+	return err
+}
+
 // DeleteCollection overrides the behavior of the generic DualWriter and deletes only from LegacyStorage.
 func (d *DualWriterMode1) DeleteCollection(ctx context.Context, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions, listOptions *metainternalversion.ListOptions) (runtime.Object, error) {
 	var method = "delete-collection"
@@ -195,25 +232,34 @@ func (d *DualWriterMode1) DeleteCollection(ctx context.Context, deleteValidation
 	}
 	d.recordLegacyDuration(false, mode1Str, d.resource, method, startLegacy)
 
-	go func(res runtime.Object) {
-		startStorage := time.Now()
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage deletecollection timeout"))
-		defer cancel()
-		storageObj, err := d.Storage.DeleteCollection(ctx, deleteValidation, options, listOptions)
-		d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			cancel()
-		}
-		areEqual := Compare(storageObj, res)
-		d.recordOutcome(mode1Str, getName(res), areEqual, method)
-		if !areEqual {
-			log.Info("object from legacy and storage are not equal")
-		}
-	}(res)
+	//nolint:errcheck
+	go d.deleteCollectionFromUnifiedStorage(ctx, res, deleteValidation, options, listOptions)
 
 	return res, err
 }
 
+func (d *DualWriterMode1) deleteCollectionFromUnifiedStorage(ctx context.Context, res runtime.Object, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions, listOptions *metainternalversion.ListOptions) error {
+	var method = "delete-collection"
+	log := d.Log.WithValues("resourceVersion", listOptions.ResourceVersion, "method", method)
+
+	startStorage := time.Now()
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage deletecollection timeout"))
+	defer cancel()
+	storageObj, err := d.Storage.DeleteCollection(ctx, deleteValidation, options, listOptions)
+	d.recordStorageDuration(err != nil, mode1Str, d.resource, method, startStorage)
+	if err != nil {
+		cancel()
+	}
+	areEqual := Compare(storageObj, res)
+	d.recordOutcome(mode1Str, getName(res), areEqual, method)
+	if !areEqual {
+		log.Info("object from legacy and storage are not equal")
+	}
+
+	return err
+}
+
 func (d *DualWriterMode1) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool, options *metav1.UpdateOptions) (runtime.Object, bool, error) {
 	var method = "update"
 	log := d.Log.WithValues("name", name, "method", method, "name", name)
@@ -228,54 +274,63 @@ func (d *DualWriterMode1) Update(ctx context.Context, name string, objInfo rest.
 	}
 	d.recordLegacyDuration(false, mode1Str, d.resource, method, startLegacy)
 
-	go func(res runtime.Object) {
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage update timeout"))
-
-		resCopy := res.DeepCopyObject()
-		// get the object to be updated
-		foundObj, err := d.Storage.Get(ctx, name, &metav1.GetOptions{})
-		if err != nil {
-			if !apierrors.IsNotFound(err) {
-				log.WithValues("object", foundObj).Error(err, "could not get object to update")
-				cancel()
-			}
-			log.Info("object not found for update, creating one")
-		}
-
-		updated, err := objInfo.UpdatedObject(ctx, resCopy)
-		if err != nil {
-			log.WithValues("object", updated).Error(err, "could not update or create object")
-			cancel()
-		}
-
-		// if the object is found, create a new updateWrapper with the object found
-		if foundObj != nil {
-			if err := enrichLegacyObject(foundObj, resCopy); err != nil {
-				log.Error(err, "could not enrich object")
-				cancel()
-			}
-			objInfo = &updateWrapper{
-				upstream: objInfo,
-				updated:  resCopy,
-			}
-		}
-		startStorage := time.Now()
-		defer cancel()
-		storageObj, _, errObjectSt := d.Storage.Update(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
-		d.recordStorageDuration(errObjectSt != nil, mode1Str, d.resource, method, startStorage)
-		if err != nil {
-			cancel()
-		}
-		areEqual := Compare(storageObj, res)
-		d.recordOutcome(mode1Str, name, areEqual, method)
-		if !areEqual {
-			log.WithValues("name", name).Info("object from legacy and storage are not equal")
-		}
-	}(res)
+	//nolint:errcheck
+	go d.updateOnUnifiedStorage(ctx, res, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
 
 	return res, async, err
 }
 
+func (d *DualWriterMode1) updateOnUnifiedStorage(ctx context.Context, res runtime.Object, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool, options *metav1.UpdateOptions) error {
+	var method = "update"
+	log := d.Log.WithValues("name", name, "method", method, "name", name)
+
+	// Ignores cancellation signals from parent context. Will automatically be canceled after 10 seconds.
+	ctx, cancel := context.WithTimeoutCause(context.WithoutCancel(ctx), time.Second*10, errors.New("storage update timeout"))
+
+	resCopy := res.DeepCopyObject()
+	// get the object to be updated
+	foundObj, err := d.Storage.Get(ctx, name, &metav1.GetOptions{})
+	if err != nil {
+		if !apierrors.IsNotFound(err) {
+			log.WithValues("object", foundObj).Error(err, "could not get object to update")
+			cancel()
+		}
+		log.Info("object not found for update, creating one")
+	}
+
+	updated, err := objInfo.UpdatedObject(ctx, resCopy)
+	if err != nil {
+		log.WithValues("object", updated).Error(err, "could not update or create object")
+		cancel()
+	}
+
+	// if the object is found, create a new updateWrapper with the object found
+	if foundObj != nil {
+		if err := enrichLegacyObject(foundObj, resCopy); err != nil {
+			log.Error(err, "could not enrich object")
+			cancel()
+		}
+		objInfo = &updateWrapper{
+			upstream: objInfo,
+			updated:  resCopy,
+		}
+	}
+	startStorage := time.Now()
+	defer cancel()
+	storageObj, _, errObjectSt := d.Storage.Update(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
+	d.recordStorageDuration(errObjectSt != nil, mode1Str, d.resource, method, startStorage)
+	if err != nil {
+		cancel()
+	}
+	areEqual := Compare(storageObj, res)
+	d.recordOutcome(mode1Str, name, areEqual, method)
+	if !areEqual {
+		log.WithValues("name", name).Info("object from legacy and storage are not equal")
+	}
+
+	return errObjectSt
+}
+
 func (d *DualWriterMode1) Destroy() {
 	d.Storage.Destroy()
 	d.Legacy.Destroy()

pkg/apiserver/rest/dualwriter_mode1_test.go

@@ -90,6 +90,64 @@ func TestMode1_Create(t *testing.T) {
 	}
 }
 
+func TestMode1_CreateOnUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		name           string
+		input          runtime.Object
+		ctx            *context.Context
+		setupLegacyFn  func(m *mock.Mock)
+		setupStorageFn func(m *mock.Mock)
+	}
+	tests :=
+		[]testCase{
+			{
+				name:  "Create on unified storage",
+				input: exampleObj,
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("Create", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(exampleObjNoRV, nil)
+				},
+			},
+			{
+				name:  "Create on unified storage works even if parent context is canceled",
+				input: exampleObj,
+				ctx:   &ctxCanceled,
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("Create", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(exampleObjNoRV, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+			err := dw.(*DualWriterMode1).createOnUnifiedStorage(ctx, tt.input, func(context.Context, runtime.Object) error { return nil }, tt.input, &metav1.CreateOptions{})
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMode1_Get(t *testing.T) {
 	type testCase struct {
 		setupLegacyFn  func(m *mock.Mock, name string)
@@ -153,6 +211,64 @@ func TestMode1_Get(t *testing.T) {
 	}
 }
 
+func TestMode1_GetFromUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		setupLegacyFn  func(m *mock.Mock, name string)
+		setupStorageFn func(m *mock.Mock, name string)
+		ctx            *context.Context
+		name           string
+		input          string
+	}
+	tests :=
+		[]testCase{
+			{
+				name:  "Get from unified storage",
+				input: "foo",
+				setupStorageFn: func(m *mock.Mock, name string) {
+					m.On("Get", mock.Anything, name, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+			{
+				name:  "Get from unified storage works even if parent context is canceled",
+				input: "foo",
+				ctx:   &ctxCanceled,
+				setupStorageFn: func(m *mock.Mock, name string) {
+					m.On("Get", mock.Anything, name, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m, tt.input)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m, tt.input)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+			err := dw.(*DualWriterMode1).getFromUnifiedStorage(ctx, exampleObj, tt.input, &metav1.GetOptions{})
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMode1_List(t *testing.T) {
 	type testCase struct {
 		setupLegacyFn  func(m *mock.Mock)
@@ -199,6 +315,62 @@ func TestMode1_List(t *testing.T) {
 	}
 }
 
+func TestMode1_ListFromUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		ctx            *context.Context
+		name           string
+		setupLegacyFn  func(m *mock.Mock)
+		setupStorageFn func(m *mock.Mock)
+	}
+	tests :=
+		[]testCase{
+			{
+				name: "list from unified storage",
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("List", mock.Anything, mock.Anything).Return(anotherList, nil)
+				},
+			},
+			{
+				name: "list from unified storage works even if parent context is canceled",
+				ctx:  &ctxCanceled,
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("List", mock.Anything, mock.Anything).Return(anotherList, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+
+			err := dw.(*DualWriterMode1).listFromUnifiedStorage(ctx, &metainternalversion.ListOptions{}, anotherList)
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMode1_Delete(t *testing.T) {
 	type testCase struct {
 		setupLegacyFn  func(m *mock.Mock, name string)
@@ -258,6 +430,63 @@ func TestMode1_Delete(t *testing.T) {
 	}
 }
 
+func TestMode1_DeleteFromUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		ctx            *context.Context
+		setupLegacyFn  func(m *mock.Mock, name string)
+		setupStorageFn func(m *mock.Mock, name string)
+		name           string
+		input          string
+	}
+	tests :=
+		[]testCase{
+			{
+				name: "Delete from unified storage",
+				setupStorageFn: func(m *mock.Mock, input string) {
+					m.On("Delete", mock.Anything, input, mock.Anything, mock.Anything).Return(exampleObj, false, nil)
+				},
+			},
+			{
+				name: "Delete from unified storage works even if parent context is canceled",
+				ctx:  &ctxCanceled,
+				setupStorageFn: func(m *mock.Mock, input string) {
+					m.On("Delete", mock.Anything, input, mock.Anything, mock.Anything).Return(exampleObj, false, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m, tt.input)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m, tt.input)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+
+			err := dw.(*DualWriterMode1).deleteFromUnifiedStorage(ctx, exampleObj, tt.input, func(ctx context.Context, obj runtime.Object) error { return nil }, &metav1.DeleteOptions{})
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMode1_DeleteCollection(t *testing.T) {
 	type testCase struct {
 		input          *metav1.DeleteOptions
@@ -317,6 +546,65 @@ func TestMode1_DeleteCollection(t *testing.T) {
 	}
 }
 
+func TestMode1_DeleteCollectionFromUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		ctx            *context.Context
+		setupLegacyFn  func(m *mock.Mock)
+		setupStorageFn func(m *mock.Mock)
+		name           string
+		input          *metav1.DeleteOptions
+	}
+	tests :=
+		[]testCase{
+			{
+				name:  "Delete Collection from unified storage",
+				input: &metav1.DeleteOptions{TypeMeta: metav1.TypeMeta{Kind: "foo"}},
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("DeleteCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+			{
+				name:  "Delete Collection from unified storage works even if parent context is canceled",
+				input: &metav1.DeleteOptions{TypeMeta: metav1.TypeMeta{Kind: "foo"}},
+				ctx:   &ctxCanceled,
+				setupStorageFn: func(m *mock.Mock) {
+					m.On("DeleteCollection", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+
+			err := dw.(*DualWriterMode1).deleteCollectionFromUnifiedStorage(ctx, exampleObj, func(ctx context.Context, obj runtime.Object) error { return nil }, tt.input, &metainternalversion.ListOptions{})
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMode1_Update(t *testing.T) {
 	type testCase struct {
 		setupLegacyFn  func(m *mock.Mock, input string)
@@ -391,3 +679,73 @@ func TestMode1_Update(t *testing.T) {
 		})
 	}
 }
+
+func TestMode1_UpdateOnUnifiedStorage(t *testing.T) {
+	ctxCanceled, cancel := context.WithCancel(context.TODO())
+	cancel()
+
+	type testCase struct {
+		ctx            *context.Context
+		setupLegacyFn  func(m *mock.Mock, input string)
+		setupStorageFn func(m *mock.Mock, input string)
+		setupGetFn     func(m *mock.Mock, input string)
+		name           string
+		input          string
+	}
+	tests :=
+		[]testCase{
+			{
+				name:  "Update on unified storage",
+				input: "foo",
+				setupStorageFn: func(m *mock.Mock, input string) {
+					m.On("Update", mock.Anything, input, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(anotherObj, false, nil)
+				},
+				setupGetFn: func(m *mock.Mock, input string) {
+					m.On("Get", mock.Anything, input, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+			{
+				name:  "Update on unified storage works even if parent context is canceled",
+				ctx:   &ctxCanceled,
+				input: "foo",
+				setupStorageFn: func(m *mock.Mock, input string) {
+					m.On("Update", mock.Anything, input, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(anotherObj, false, nil)
+				},
+				setupGetFn: func(m *mock.Mock, input string) {
+					m.On("Get", mock.Anything, input, mock.Anything).Return(exampleObj, nil)
+				},
+			},
+		}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			l := (LegacyStorage)(nil)
+			s := (Storage)(nil)
+			m := &mock.Mock{}
+
+			ls := legacyStoreMock{m, l}
+			us := storageMock{m, s}
+
+			if tt.setupLegacyFn != nil {
+				tt.setupLegacyFn(m, tt.input)
+			}
+			if tt.setupStorageFn != nil {
+				tt.setupStorageFn(m, tt.input)
+			}
+
+			if tt.setupGetFn != nil {
+				tt.setupGetFn(m, tt.input)
+			}
+
+			ctx := context.TODO()
+			if tt.ctx != nil {
+				ctx = *tt.ctx
+			}
+
+			dw := NewDualWriter(Mode1, ls, us, p, kind)
+
+			err := dw.(*DualWriterMode1).updateOnUnifiedStorage(ctx, exampleObj, tt.input, updatedObjInfoObj{}, func(ctx context.Context, obj runtime.Object) error { return nil }, func(ctx context.Context, obj, old runtime.Object) error { return nil }, false, &metav1.UpdateOptions{})
+			assert.NoError(t, err)
+		})
+	}
+}

pkg/apiserver/rest/storage_mocks_test.go

@@ -2,6 +2,7 @@ package rest
 
 import (
 	"context"
+	"errors"
 
 	"github.com/stretchr/testify/mock"
 	"k8s.io/apimachinery/pkg/api/meta"
@@ -83,6 +84,12 @@ func (m legacyStoreMock) DeleteCollection(ctx context.Context, deleteValidation
 
 // Unified Store
 func (m storageMock) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
+	select {
+	case <-ctx.Done():
+		return nil, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, name, options)
 	if name == "object-fail" {
 		return nil, args.Error(1)
@@ -94,6 +101,12 @@ func (m storageMock) Get(ctx context.Context, name string, options *metav1.GetOp
 }
 
 func (m storageMock) Create(ctx context.Context, obj runtime.Object, createValidation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
+	select {
+	case <-ctx.Done():
+		return nil, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, obj, createValidation, options)
 	acc, err := meta.Accessor(obj)
 	if err != nil {
@@ -107,6 +120,12 @@ func (m storageMock) Create(ctx context.Context, obj runtime.Object, createValid
 }
 
 func (m storageMock) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
+	select {
+	case <-ctx.Done():
+		return nil, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, options)
 	if options.Kind == "fail" {
 		return nil, args.Error(1)
@@ -119,6 +138,12 @@ func (m storageMock) NewList() runtime.Object {
 }
 
 func (m storageMock) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool, options *metav1.UpdateOptions) (runtime.Object, bool, error) {
+	select {
+	case <-ctx.Done():
+		return nil, false, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
 	if name == "object-fail" {
 		return nil, false, args.Error(2)
@@ -127,6 +152,12 @@ func (m storageMock) Update(ctx context.Context, name string, objInfo rest.Updat
 }
 
 func (m storageMock) Delete(ctx context.Context, name string, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions) (runtime.Object, bool, error) {
+	select {
+	case <-ctx.Done():
+		return nil, false, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, name, deleteValidation, options)
 	if name == "object-fail" {
 		return nil, false, args.Error(2)
@@ -135,6 +166,12 @@ func (m storageMock) Delete(ctx context.Context, name string, deleteValidation r
 }
 
 func (m storageMock) DeleteCollection(ctx context.Context, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions, listOptions *metainternalversion.ListOptions) (runtime.Object, error) {
+	select {
+	case <-ctx.Done():
+		return nil, errors.New("context canceled")
+	default:
+	}
+
 	args := m.Called(ctx, deleteValidation, options, listOptions)
 	if options.Kind == "fail" {
 		return nil, args.Error(1)