Skip to content

Commit 5e5068b

Browse files
smirasmira
committed
feat: move Modify implementation so State
This moved `Modify` from controller adapter to state, allowing Modify with plain State. Fixes siderolabs/omni#1136 Signed-off-by: Andrey Smirnov <[email protected]>
1 parent 1148efe commit 5e5068b

File tree

6 files changed

+191
-19
lines changed

6 files changed

+191
-19
lines changed

‎pkg/controller/runtime/internal/controllerstate/adapter.go‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -223,24 +223,6 @@ func (adapter *StateAdapter) modify(
223223
emptyResource.Metadata().Namespace(), emptyResource.Metadata().Type(), adapter.Name, emptyResource.Metadata().ID())
224224
}
225225

226-
_, err := adapter.State.Get(ctx, emptyResource.Metadata())
227-
if err != nil {
228-
if state.IsNotFoundError(err) {
229-
err = updateFunc(emptyResource)
230-
if err != nil {
231-
return nil, err
232-
}
233-
234-
if err = adapter.State.Create(ctx, emptyResource, state.WithCreateOwner(adapter.Name)); err != nil {
235-
return nil, err
236-
}
237-
238-
return emptyResource, nil
239-
}
240-
241-
return nil, fmt.Errorf("error querying current object state: %w", err)
242-
}
243-
244226
updateOptions := []state.UpdateOption{state.WithUpdateOwner(adapter.Name)}
245227

246228
modifyOptions := controller.ToModifyOptions(options...)
@@ -250,7 +232,7 @@ func (adapter *StateAdapter) modify(
250232
updateOptions = append(updateOptions, state.WithExpectedPhaseAny())
251233
}
252234

253-
return adapter.State.UpdateWithConflicts(ctx, emptyResource.Metadata(), updateFunc, updateOptions...)
235+
return adapter.State.ModifyWithResult(ctx, emptyResource, updateFunc, updateOptions...)
254236
}
255237

256238
// AddFinalizer implements controller.Runtime interface.

‎pkg/safe/state.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,32 @@ func StateWatchKind[T resource.Resource](ctx context.Context, st state.CoreState
200200
return nil
201201
}
202202

203+
// StateModify is a type safe wrapper around state.Modify.
204+
func StateModify[T resource.Resource](ctx context.Context, st state.State, r T, fn func(T) error, options ...state.UpdateOption) error {
205+
return st.Modify(ctx, r, func(r resource.Resource) error {
206+
arg, ok := r.(T)
207+
if !ok {
208+
return fmt.Errorf("type mismatch: expected %T, got %T", arg, r)
209+
}
210+
211+
return fn(arg)
212+
}, options...)
213+
}
214+
215+
// StateModifyWithResult is a type safe wrapper around state.ModifyWithResult.
216+
func StateModifyWithResult[T resource.Resource](ctx context.Context, st state.State, r T, fn func(T) error, options ...state.UpdateOption) (T, error) {
217+
got, err := st.ModifyWithResult(ctx, r, func(r resource.Resource) error {
218+
arg, ok := r.(T)
219+
if !ok {
220+
return fmt.Errorf("type mismatch: expected %T, got %T", arg, r)
221+
}
222+
223+
return fn(arg)
224+
}, options...)
225+
226+
return typeAssertOrZero[T](got, err)
227+
}
228+
203229
// List is a type safe wrapper around resource.List.
204230
type List[T resource.Resource] struct {
205231
list resource.List

‎pkg/state/conformance/state.go‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package conformance
66

77
import (
88
"context"
9+
"errors"
910
"fmt"
1011
"math/rand"
1112
"regexp"
@@ -1459,6 +1460,103 @@ func (suite *StateSuite) TestTeardownAndDestroy() {
14591460
suite.Require().NoError(eg.Wait())
14601461
}
14611462

1463+
// TestModify verifies Modify.
1464+
func (suite *StateSuite) TestModify() {
1465+
path1 := NewPathResource(suite.getNamespace(), "var/run/modify")
1466+
1467+
ctx := context.Background()
1468+
1469+
p1, err := safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1470+
r.Metadata().Labels().Set("foo", "bar")
1471+
1472+
return nil
1473+
})
1474+
suite.Require().NoError(err)
1475+
1476+
suite.Assert().Equal(resource.String(path1), resource.String(p1))
1477+
suite.Assert().Equal("bar", p1.Metadata().Labels().Raw()["foo"])
1478+
suite.Assert().Empty(p1.Metadata().Owner())
1479+
1480+
p2, err := safe.StateGet[*PathResource](ctx, suite.State, path1.Metadata())
1481+
suite.Require().NoError(err)
1482+
1483+
suite.Assert().Equal(resource.String(path1), resource.String(p2))
1484+
suite.Assert().Equal("bar", p2.Metadata().Labels().Raw()["foo"])
1485+
1486+
p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1487+
r.Metadata().Labels().Delete("foo")
1488+
1489+
return nil
1490+
})
1491+
suite.Require().NoError(err)
1492+
1493+
suite.Assert().True(p1.Metadata().Labels().Empty())
1494+
1495+
p2, err = safe.StateGet[*PathResource](ctx, suite.State, path1.Metadata())
1496+
suite.Require().NoError(err)
1497+
1498+
suite.Assert().True(p2.Metadata().Labels().Empty())
1499+
1500+
_, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(*PathResource) error {
1501+
return errors.New("modify error")
1502+
})
1503+
suite.Require().EqualError(err, "modify error")
1504+
1505+
_, err = suite.State.Teardown(ctx, path1.Metadata())
1506+
suite.Require().NoError(err)
1507+
1508+
_, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1509+
r.Metadata().Labels().Set("foo", "bar")
1510+
1511+
return nil
1512+
})
1513+
suite.Require().Error(err)
1514+
suite.Assert().True(state.IsPhaseConflictError(err))
1515+
1516+
p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1517+
r.Metadata().Labels().Set("foo", "bar2")
1518+
1519+
return nil
1520+
}, state.WithExpectedPhaseAny())
1521+
suite.Require().NoError(err)
1522+
suite.Assert().Equal(resource.PhaseTearingDown, p1.Metadata().Phase())
1523+
suite.Assert().Equal("bar2", p1.Metadata().Labels().Raw()["foo"])
1524+
}
1525+
1526+
// TestModifyWithOwner verifies Modify with Owner.
1527+
func (suite *StateSuite) TestModifyWithOwner() {
1528+
path1 := NewPathResource(suite.getNamespace(), "var/run/modify/owned")
1529+
1530+
ctx := context.Background()
1531+
1532+
p1, err := safe.StateModifyWithResult(ctx, suite.State, path1, func(*PathResource) error {
1533+
return nil
1534+
}, state.WithUpdateOwner("owner"))
1535+
suite.Require().NoError(err)
1536+
1537+
suite.Assert().Equal(resource.String(path1), resource.String(p1))
1538+
suite.Assert().Equal("owner", p1.Metadata().Owner())
1539+
1540+
p1, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1541+
r.Metadata().Labels().Set("foo", "bar")
1542+
1543+
return nil
1544+
}, state.WithUpdateOwner("owner"))
1545+
suite.Require().NoError(err)
1546+
1547+
suite.Assert().Equal(resource.String(path1), resource.String(p1))
1548+
suite.Assert().Equal("owner", p1.Metadata().Owner())
1549+
suite.Assert().Equal("bar", p1.Metadata().Labels().Raw()["foo"])
1550+
1551+
_, err = safe.StateModifyWithResult(ctx, suite.State, path1, func(r *PathResource) error {
1552+
r.Metadata().Labels().Set("foo", "baz")
1553+
1554+
return nil
1555+
})
1556+
suite.Require().Error(err)
1557+
suite.Require().True(state.IsOwnerConflictError(err))
1558+
}
1559+
14621560
func assertContextIsCanceled(t *testing.T, ctx context.Context) { //nolint:revive
14631561
t.Helper()
14641562

‎pkg/state/options.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,13 @@ func WithExpectedPhaseAny() UpdateOption {
122122
}
123123
}
124124

125+
// WithUpdateOptions sets update options for the update request.
126+
func WithUpdateOptions(opts UpdateOptions) UpdateOption {
127+
return func(options *UpdateOptions) {
128+
*options = opts
129+
}
130+
}
131+
125132
// TeardownOptions for the CoreState.Teardown function.
126133
type TeardownOptions struct {
127134
Owner string

‎pkg/state/state.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,4 +144,14 @@ type State interface {
144144
// It's not an error to tear down a resource which is already being torn down.
145145
// The call blocks until all resource finalizers are empty.
146146
TeardownAndDestroy(context.Context, resource.Pointer, ...TeardownAndDestroyOption) error
147+
148+
// Modify modifies an existing resource or creates a new one.
149+
//
150+
// It is a shorthand for Get+UpdateWithConflicts+Create.
151+
Modify(ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption) error
152+
153+
// ModifyWithResult modifies an existing resource or creates a new one.
154+
//
155+
// It is a shorthand for Get+UpdateWithConflicts+Create.
156+
ModifyWithResult(ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption) (resource.Resource, error)
147157
}

‎pkg/state/wrap.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ package state
66

77
import (
88
"context"
9+
"fmt"
10+
11+
"github.com/siderolabs/go-pointer"
912

1013
"github.com/cosi-project/runtime/pkg/resource"
1114
)
@@ -237,3 +240,49 @@ func (state coreWrapper) TeardownAndDestroy(ctx context.Context, resourcePointer
237240

238241
return state.Destroy(ctx, resourcePointer, WithDestroyOwner(options.Owner))
239242
}
243+
244+
// Modify modifies an existing resource or creates a new one.
245+
//
246+
// It is a shorthand for Get+UpdateWithConflicts+Create.
247+
func (state coreWrapper) Modify(
248+
ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption,
249+
) error {
250+
_, err := state.ModifyWithResult(ctx, emptyResource, updateFunc, options...)
251+
252+
return err
253+
}
254+
255+
// ModifyWithResult modifies an existing resource or creates a new one.
256+
//
257+
// It is a shorthand for Get+UpdateWithConflicts+Create.
258+
func (state coreWrapper) ModifyWithResult(
259+
ctx context.Context, emptyResource resource.Resource, updateFunc func(resource.Resource) error, options ...UpdateOption,
260+
) (resource.Resource, error) {
261+
opts := UpdateOptions{
262+
ExpectedPhase: pointer.To(resource.PhaseRunning),
263+
}
264+
265+
for _, opt := range options {
266+
opt(&opts)
267+
}
268+
269+
_, err := state.Get(ctx, emptyResource.Metadata())
270+
if err != nil {
271+
if IsNotFoundError(err) {
272+
err = updateFunc(emptyResource)
273+
if err != nil {
274+
return nil, err
275+
}
276+
277+
if err = state.Create(ctx, emptyResource, WithCreateOwner(opts.Owner)); err != nil {
278+
return nil, err
279+
}
280+
281+
return emptyResource, nil
282+
}
283+
284+
return nil, fmt.Errorf("error querying current object state: %w", err)
285+
}
286+
287+
return state.UpdateWithConflicts(ctx, emptyResource.Metadata(), updateFunc, WithUpdateOptions(opts))
288+
}

0 commit comments

Comments
 (0)