From 5e5068b0f2bbcd35c0e165bf5707db9d22894920 Mon Sep 17 00:00:00 2001 From: Andrey Smirnov Date: Thu, 15 May 2025 15:47:49 +0400 Subject: [PATCH] feat: move `Modify` implementation so State This moved `Modify` from controller adapter to state, allowing Modify with plain State. Fixes https://github.com/siderolabs/omni/issues/1136 Signed-off-by: Andrey Smirnov --- .../internal/controllerstate/adapter.go | 20 +--- pkg/safe/state.go | 26 +++++ pkg/state/conformance/state.go | 98 +++++++++++++++++++ pkg/state/options.go | 7 ++ pkg/state/state.go | 10 ++ pkg/state/wrap.go | 49 ++++++++++ 6 files changed, 191 insertions(+), 19 deletions(-) diff --git a/pkg/controller/runtime/internal/controllerstate/adapter.go b/pkg/controller/runtime/internal/controllerstate/adapter.go index 951f846..1f7f1aa 100644 --- a/pkg/controller/runtime/internal/controllerstate/adapter.go +++ b/pkg/controller/runtime/internal/controllerstate/adapter.go @@ -223,24 +223,6 @@ func (adapter *StateAdapter) modify( emptyResource.Metadata().Namespace(), emptyResource.Metadata().Type(), adapter.Name, emptyResource.Metadata().ID()) } - _, err := adapter.State.Get(ctx, emptyResource.Metadata()) - if err != nil { - if state.IsNotFoundError(err) { - err = updateFunc(emptyResource) - if err != nil { - return nil, err - } - - if err = adapter.State.Create(ctx, emptyResource, state.WithCreateOwner(adapter.Name)); err != nil { - return nil, err - } - - return emptyResource, nil - } - - return nil, fmt.Errorf("error querying current object state: %w", err) - } - updateOptions := []state.UpdateOption{state.WithUpdateOwner(adapter.Name)} modifyOptions := controller.ToModifyOptions(options...) @@ -250,7 +232,7 @@ func (adapter *StateAdapter) modify( updateOptions = append(updateOptions, state.WithExpectedPhaseAny()) } - return adapter.State.UpdateWithConflicts(ctx, emptyResource.Metadata(), updateFunc, updateOptions...) + return adapter.State.ModifyWithResult(ctx, emptyResource, updateFunc, updateOptions...) } // AddFinalizer implements controller.Runtime interface. diff --git a/pkg/safe/state.go b/pkg/safe/state.go index 3bc76b5..0317563 100644 --- a/pkg/safe/state.go +++ b/pkg/safe/state.go @@ -200,6 +200,32 @@ func StateWatchKind[T resource.Resource](ctx context.Context, st state.CoreState return nil } +// StateModify is a type safe wrapper around state.Modify. +func StateModify[T resource.Resource](ctx context.Context, st state.State, r T, fn func(T) error, options ...state.UpdateOption) error { + return st.Modify(ctx, r, func(r resource.Resource) error { + arg, ok := r.(T) + if !ok { + return fmt.Errorf("type mismatch: expected %T, got %T", arg, r) + } + + return fn(arg) + }, options...) +} + +// StateModifyWithResult is a type safe wrapper around state.ModifyWithResult. +func StateModifyWithResult[T resource.Resource](ctx context.Context, st state.State, r T, fn func(T) error, options ...state.UpdateOption) (T, error) { + got, err := st.ModifyWithResult(ctx, r, func(r resource.Resource) error { + arg, ok := r.(T) + if !ok { + return fmt.Errorf("type mismatch: expected %T, got %T", arg, r) + } + + return fn(arg) + }, options...) + + return typeAssertOrZero[T](got, err) +} + // List is a type safe wrapper around resource.List. type List[T resource.Resource] struct { list resource.List diff --git a/pkg/state/conformance/state.go b/pkg/state/conformance/state.go index 1de2f5b..f620283 100644 --- a/pkg/state/conformance/state.go +++ b/pkg/state/conformance/state.go @@ -6,6 +6,7 @@ package conformance import ( "context" + "errors" "fmt" "math/rand" "regexp" @@ -1459,6 +1460,103 @@ func (suite *StateSuite) TestTeardownAndDestroy() { suite.Require().NoError(eg.Wait()) } +// TestModify verifies Modify. +func (suite *StateSuite) TestModify() { + path1 := NewPathResource(suite.getNamespace(), "var/run/modify") + + ctx := context.Background() + + p1, err := safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Set("foo", "bar") + + return nil + }) + suite.Require().NoError(err) + + suite.Assert().Equal(resource.String(path1), resource.String(p1)) + suite.Assert().Equal("bar", p1.Metadata().Labels().Raw()["foo"]) + suite.Assert().Empty(p1.Metadata().Owner()) + + p2, err := safe.StateGet[*PathResource](ctx, suite.State, path1.Metadata()) + suite.Require().NoError(err) + + suite.Assert().Equal(resource.String(path1), resource.String(p2)) + suite.Assert().Equal("bar", p2.Metadata().Labels().Raw()["foo"]) + + p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Delete("foo") + + return nil + }) + suite.Require().NoError(err) + + suite.Assert().True(p1.Metadata().Labels().Empty()) + + p2, err = safe.StateGet[*PathResource](ctx, suite.State, path1.Metadata()) + suite.Require().NoError(err) + + suite.Assert().True(p2.Metadata().Labels().Empty()) + + _, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(*PathResource) error { + return errors.New("modify error") + }) + suite.Require().EqualError(err, "modify error") + + _, err = suite.State.Teardown(ctx, path1.Metadata()) + suite.Require().NoError(err) + + _, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Set("foo", "bar") + + return nil + }) + suite.Require().Error(err) + suite.Assert().True(state.IsPhaseConflictError(err)) + + p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Set("foo", "bar2") + + return nil + }, state.WithExpectedPhaseAny()) + suite.Require().NoError(err) + suite.Assert().Equal(resource.PhaseTearingDown, p1.Metadata().Phase()) + suite.Assert().Equal("bar2", p1.Metadata().Labels().Raw()["foo"]) +} + +// TestModifyWithOwner verifies Modify with Owner. +func (suite *StateSuite) TestModifyWithOwner() { + path1 := NewPathResource(suite.getNamespace(), "var/run/modify/owned") + + ctx := context.Background() + + p1, err := safe.StateModifyWithResult(ctx, suite.State, path1, func(*PathResource) error { + return nil + }, state.WithUpdateOwner("owner")) + suite.Require().NoError(err) + + suite.Assert().Equal(resource.String(path1), resource.String(p1)) + suite.Assert().Equal("owner", p1.Metadata().Owner()) + + p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Set("foo", "bar") + + return nil + }, state.WithUpdateOwner("owner")) + suite.Require().NoError(err) + + suite.Assert().Equal(resource.String(path1), resource.String(p1)) + suite.Assert().Equal("owner", p1.Metadata().Owner()) + suite.Assert().Equal("bar", p1.Metadata().Labels().Raw()["foo"]) + + _, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error { + r.Metadata().Labels().Set("foo", "baz") + + return nil + }) + suite.Require().Error(err) + suite.Require().True(state.IsOwnerConflictError(err)) +} + func assertContextIsCanceled(t *testing.T, ctx context.Context) { //nolint:revive t.Helper() diff --git a/pkg/state/options.go b/pkg/state/options.go index 7bdb043..e059117 100644 --- a/pkg/state/options.go +++ b/pkg/state/options.go @@ -122,6 +122,13 @@ func WithExpectedPhaseAny() UpdateOption { } } +// WithUpdateOptions sets update options for the update request. +func WithUpdateOptions(opts UpdateOptions) UpdateOption { + return func(options *UpdateOptions) { + *options = opts + } +} + // TeardownOptions for the CoreState.Teardown function. type TeardownOptions struct { Owner string diff --git a/pkg/state/state.go b/pkg/state/state.go index ef1f832..746d167 100644 --- a/pkg/state/state.go +++ b/pkg/state/state.go @@ -144,4 +144,14 @@ type State interface { // It's not an error to tear down a resource which is already being torn down. // The call blocks until all resource finalizers are empty. TeardownAndDestroy(context.Context, resource.Pointer, ...TeardownAndDestroyOption) error + + // Modify modifies an existing resource or creates a new one. + // + // It is a shorthand for Get+UpdateWithConflicts+Create. + Modify(ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption) error + + // ModifyWithResult modifies an existing resource or creates a new one. + // + // It is a shorthand for Get+UpdateWithConflicts+Create. + ModifyWithResult(ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption) (resource.Resource, error) } diff --git a/pkg/state/wrap.go b/pkg/state/wrap.go index 02fa321..1d54213 100644 --- a/pkg/state/wrap.go +++ b/pkg/state/wrap.go @@ -6,6 +6,9 @@ package state import ( "context" + "fmt" + + "github.com/siderolabs/go-pointer" "github.com/cosi-project/runtime/pkg/resource" ) @@ -237,3 +240,49 @@ func (state coreWrapper) TeardownAndDestroy(ctx context.Context, resourcePointer return state.Destroy(ctx, resourcePointer, WithDestroyOwner(options.Owner)) } + +// Modify modifies an existing resource or creates a new one. +// +// It is a shorthand for Get+UpdateWithConflicts+Create. +func (state coreWrapper) Modify( + ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption, +) error { + _, err := state.ModifyWithResult(ctx, emptyResource, updateFunc, options...) + + return err +} + +// ModifyWithResult modifies an existing resource or creates a new one. +// +// It is a shorthand for Get+UpdateWithConflicts+Create. +func (state coreWrapper) ModifyWithResult( + ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption, +) (resource.Resource, error) { + opts := UpdateOptions{ + ExpectedPhase: pointer.To(resource.PhaseRunning), + } + + for _, opt := range options { + opt(&opts) + } + + _, err := state.Get(ctx, emptyResource.Metadata()) + if err != nil { + if IsNotFoundError(err) { + err = updateFunc(emptyResource) + if err != nil { + return nil, err + } + + if err = state.Create(ctx, emptyResource, WithCreateOwner(opts.Owner)); err != nil { + return nil, err + } + + return emptyResource, nil + } + + return nil, fmt.Errorf("error querying current object state: %w", err) + } + + return state.UpdateWithConflicts(ctx, emptyResource.Metadata(), updateFunc, WithUpdateOptions(opts)) +}