Merge pull request #546 from scothis/stasher

Type safe stashing of context values

README.md

@@ -915,19 +915,21 @@ To setup a Config for a test and make assertions that the expected behavior matc
 
 The stash allows passing arbitrary state between sub reconcilers within the scope of a single reconciler request. Values are stored on the context by [`StashValue`](https://pkg.go.dev/reconciler.io/runtime/reconcilers#StashValue) and accessed via [`RetrieveValue`](https://pkg.go.dev/reconciler.io/runtime/reconcilers#RetrieveValue).
 
+A [`Stasher`](https://pkg.go.dev/reconciler.io/runtime/reconcilers#Stasher) provides a convenient way to interact with typed values. Create a [`NewStasher`](https://pkg.go.dev/reconciler.io/runtime/reconcilers#NewStasher) using the type of the value being stashed and a unique stash key. All operations through a stasher, including retrieval are type safe with options for handling missing values on retrieval.
+
 For testing, given stashed values can be defined in a [SubReconcilerTests](#subreconcilertests) with [`GivenStashedValues`](https://pkg.go.dev/reconciler.io/runtime/testing#SubReconcilerTestCase.GivenStashedValues). Newly stashed or mutated values expectations are defined with [`ExpectStashedValues`](https://pkg.go.dev/reconciler.io/runtime/testing#SubReconcilerTestCase.ExpectStashedValues). An optional, custom function for asserting stashed values can be provided via [`VerifyStashedValue`](https://pkg.go.dev/reconciler.io/runtime/testing#SubReconcilerTestCase.VerifyStashedValue).
 
 **Example:**
 
 ```go
-const exampleStashKey reconcilers.StashKey = "example"
+var exampleStasher = reconcilers.NewStasher[Example]("example")
 
 func StashExampleSubReconciler(c reconcilers.Config) reconcilers.SubReconciler[*examplev1.MyExample] {
  return &reconcilers.SyncReconciler[*examplev1.MyExample]{
  Name: "StashExample",
  Sync: func(ctx context.Context, resource *examplev1.MyExample) error {
  value := Example{} // something we want to expose to a sub reconciler later in this chain
- reconcilers.StashValue(ctx, exampleStashKey, value)
+ exampleStasher.Store(ctx, value)
  return nil
  },
  }
@@ -938,9 +940,9 @@ func StashExampleSubReconciler(c reconcilers.Config) reconcilers.SubReconciler[*
  return &reconcilers.SyncReconciler[*examplev1.MyExample]{
  Name: "StashExample",
  Sync: func(ctx context.Context, resource *examplev1.MyExample) error {
- value, ok := reconcilers.RetrieveValue(ctx, exampleStashKey).(Example)
- if !ok {
- return nil, fmt.Errorf("expected stashed value for key %q", exampleStashKey)
+ value, err := exampleStasher.RetrieveOrError(ctx)
+ if err == reconcilers.ErrStashValueNotFound {
+ return nil, fmt.Errorf("%w for key %q", err, exampleStasher.Key())
  }
  ... // do something with the value
  },

reconcilers/childset.go

@@ -232,13 +232,13 @@ func (r *ChildSetReconciler[T, CT, CLT]) childReconcilerFor(desired CT, desiredE
  return desired, desiredErr
  },
  ReflectChildStatusOnParent: func(ctx context.Context, parent T, child CT, err error) {
- result := retrieveChildSetResult[CT](ctx)
+ result := childSetResultStasher[CT]().RetrieveOrEmpty(ctx)
  result.Children = append(result.Children, ChildSetPartialResult[CT]{
  Id: id,
  Child: child,
  Err: err,
  })
- stashChildSetResult(ctx, result)
+ childSetResultStasher[CT]().Store(ctx, result)
  },
  HarmonizeImmutableFields: r.HarmonizeImmutableFields,
  MergeBeforeUpdate: r.MergeBeforeUpdate,
@@ -370,7 +370,7 @@ func (r *ChildSetReconciler[T, CT, CLT]) composeChildReconcilers(ctx context.Con
 }
 
 func (r *ChildSetReconciler[T, CT, CLT]) reflectStatus(ctx context.Context, parent T) {
- result := clearChildSetResult[CT](ctx)
+ result := childSetResultStasher[CT]().Clear(ctx)
  r.ReflectChildrenStatusOnParent(ctx, parent, result)
 }
 
@@ -392,26 +392,8 @@ func (r *ChildSetResult[T]) AggregateError() error {
  return utilerrors.NewAggregate(errs)
 }
 
-const childSetResultStashKey StashKey = "reconciler.io/runtime:childSetResult"
-
-func retrieveChildSetResult[T client.Object](ctx context.Context) ChildSetResult[T] {
- value := RetrieveValue(ctx, childSetResultStashKey)
- if result, ok := value.(ChildSetResult[T]); ok {
- return result
- }
- return ChildSetResult[T]{}
-}
-
-func stashChildSetResult[T client.Object](ctx context.Context, result ChildSetResult[T]) {
- StashValue(ctx, childSetResultStashKey, result)
-}
-
-func clearChildSetResult[T client.Object](ctx context.Context) ChildSetResult[T] {
- value := ClearValue(ctx, childSetResultStashKey)
- if result, ok := value.(ChildSetResult[T]); ok {
- return result
- }
- return ChildSetResult[T]{}
+func childSetResultStasher[T client.Object]() Stasher[ChildSetResult[T]] {
+ return NewStasher[ChildSetResult[T]]("reconciler.io/runtime:childSetResult")
 }
 
 const knownChildrenStashKey StashKey = "reconciler.io/runtime:knownChildren"

reconcilers/stash.go

@@ -18,6 +18,7 @@ package reconcilers
 
 import (
  "context"
+ "errors"
  "fmt"
 )
 
@@ -55,3 +56,72 @@ func ClearValue(ctx context.Context, key StashKey) interface{} {
  delete(stash, key)
  return value
 }
+
+// Stasher stores and retrieves values from the stash context. The context which gets passed to its methods must be configured
+// with a stash via WithStash(). The stash is pre-configured for the context within a reconciler.
+type Stasher[T any] interface {
+ // Key is the stash key used to store and retrieve the value
+ Key() StashKey
+
+ // Store saves the value in the stash under the key
+ Store(ctx context.Context, value T)
+
+ // Clear removes the key from the stash returning the previous value, if any.
+ Clear(ctx context.Context) T
+
+ // RetrieveOrDie retrieves the value from the stash, or panics if the key is not in the stash
+ RetrieveOrDie(ctx context.Context) T
+
+ // RetrieveOrEmpty retrieves the value from the stash, or an error if the key is not in the stash
+ RetrieveOrError(ctx context.Context) (T, error)
+
+ // RetrieveOrEmpty retrieves the value from the stash, or the empty value if the key is not in the stash
+ RetrieveOrEmpty(ctx context.Context) T
+}
+
+// NewStasher creates a stasher for the value type
+func NewStasher[T any](key StashKey) Stasher[T] {
+ return &stasher[T]{
+ key: key,
+ }
+}
+
+type stasher[T any] struct {
+ key StashKey
+}
+
+func (s *stasher[T]) Key() StashKey {
+ return s.key
+}
+
+func (s *stasher[T]) Store(ctx context.Context, value T) {
+ StashValue(ctx, s.Key(), value)
+}
+
+func (s *stasher[T]) Clear(ctx context.Context) T {
+ previous, _ := ClearValue(ctx, s.Key()).(T)
+ return previous
+}
+
+func (s *stasher[T]) RetrieveOrDie(ctx context.Context) T {
+ value, err := s.RetrieveOrError(ctx)
+ if err != nil {
+ panic(err)
+ }
+ return value
+}
+
+var ErrStashValueNotFound = errors.New("value not found in stash")
+
+func (s *stasher[T]) RetrieveOrError(ctx context.Context) (T, error) {
+ value, ok := RetrieveValue(ctx, s.Key()).(T)
+ if !ok {
+ return value, ErrStashValueNotFound
+ }
+ return value, nil
+}
+
+func (s *stasher[T]) RetrieveOrEmpty(ctx context.Context) T {
+ value, _ := s.RetrieveOrError(ctx)
+ return value
+}

reconcilers/stash_test.go

@@ -47,7 +47,7 @@ func TestStash(t *testing.T) {
  }
 
  var key StashKey = "stash-key"
- ctx := WithStash(context.TODO())
+ ctx := WithStash(context.Background())
  for _, c := range tests {
  t.Run(c.name, func(t *testing.T) {
  StashValue(ctx, key, c.value)
@@ -59,7 +59,7 @@ func TestStash(t *testing.T) {
 }
 
 func TestStash_StashValue_UndecoratedContext(t *testing.T) {
- ctx := context.TODO()
+ ctx := context.Background()
  var key StashKey = "stash-key"
  value := "value"
 
@@ -72,7 +72,7 @@ func TestStash_StashValue_UndecoratedContext(t *testing.T) {
 }
 
 func TestStash_RetrieveValue_UndecoratedContext(t *testing.T) {
- ctx := context.TODO()
+ ctx := context.Background()
  var key StashKey = "stash-key"
 
  defer func() {
@@ -84,10 +84,87 @@ func TestStash_RetrieveValue_UndecoratedContext(t *testing.T) {
 }
 
 func TestStash_RetrieveValue_Undefined(t *testing.T) {
- ctx := WithStash(context.TODO())
+ ctx := WithStash(context.Background())
  var key StashKey = "stash-key"
 
  if value := RetrieveValue(ctx, key); value != nil {
  t.Error("expected RetrieveValue() to return nil for undefined key")
  }
 }
+
+func TestStasher(t *testing.T) {
+ ctx := WithStash(context.Background())
+ stasher := NewStasher[string]("my-key")
+
+ if key := stasher.Key(); key != StashKey("my-key") {
+ t.Errorf("expected key to be %q got %q", StashKey("my-key"), key)
+ }
+
+ t.Run("no value", func(t *testing.T) {
+ t.Run("RetrieveOrEmpty", func(t *testing.T) {
+ if value := stasher.RetrieveOrEmpty(ctx); value != "" {
+ t.Error("expected value to be empty")
+ }
+ })
+ t.Run("RetrieveOrError", func(t *testing.T) {
+ if value, err := stasher.RetrieveOrError(ctx); err == nil {
+ t.Error("expected err")
+ } else if value != "" {
+ t.Error("expected value to be empty")
+ }
+ })
+ t.Run("RetrieveOrDie", func(t *testing.T) {
+ defer func() {
+ if r := recover(); r != nil {
+ return
+ }
+ t.Errorf("expected to recover")
+ }()
+ value := stasher.RetrieveOrDie(ctx)
+ t.Errorf("expected to panic, got %q", value)
+ })
+ })
+
+ t.Run("has value", func(t *testing.T) {
+ stasher.Store(ctx, "hello world")
+ t.Run("RetrieveOrEmpty", func(t *testing.T) {
+ if value := stasher.RetrieveOrEmpty(ctx); value != "hello world" {
+ t.Errorf("expected value to be %q got %q", "hello world", value)
+ }
+ })
+ t.Run("RetrieveOrError", func(t *testing.T) {
+ if value, err := stasher.RetrieveOrError(ctx); err != nil {
+ t.Errorf("unexpected err: %s", err)
+ } else if value != "hello world" {
+ t.Errorf("expected value to be %q got %q", "hello world", value)
+ }
+ })
+ t.Run("RetrieveOrDie", func(t *testing.T) {
+ if value := stasher.RetrieveOrDie(ctx); value != "hello world" {
+ t.Errorf("expected value to be %q got %q", "hello world", value)
+ }
+ })
+ })
+
+ t.Run("context scoped", func(t *testing.T) {
+ stasher.Store(ctx, "hello world")
+ if value := stasher.RetrieveOrEmpty(ctx); value != "hello world" {
+ t.Errorf("expected value to be %q got %q", "hello world", value)
+ }
+ altCtx := WithStash(context.Background())
+ if value := stasher.RetrieveOrEmpty(altCtx); value != "" {
+ t.Error("expected value to be empty")
+ }
+ })
+
+ t.Run("Clear", func(t *testing.T) {
+ stasher.Store(ctx, "hello world")
+ if value := stasher.RetrieveOrEmpty(ctx); value != "hello world" {
+ t.Errorf("expected value to be %q got %q", "hello world", value)
+ }
+ stasher.Clear(ctx)
+ if value := stasher.RetrieveOrEmpty(ctx); value != "" {
+ t.Error("expected value to be empty")
+ }
+ })
+}