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")
+		}
+	})
+}