Type safe stashing of values

Introducing a Stasher which is a convenience helper for interacting with typed values for a single stash key. The stasher itself is generic and handles casting retrieved values to the desired type. There are three different ways to handle the value not existing or being of the wrong type: - RetrieveOrDie panics if a value of the desired type is not found - RetrieveOrError errors if a value of the desired type is not found - RetrieveOrEmpty falls back to the empty value for the type Signed-off-by: Scott Andrews <scott@andrews.me>
reconcilerio · Sep 19, 2024 · c55d7cb · c55d7cb
1 parent 435c49f
commit c55d7cb
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 32 deletions.
diff --git a/README.md b/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
 		},

diff --git a/reconcilers/childset.go b/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"

diff --git a/reconcilers/stash.go b/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 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
+}
diff --git a/reconcilers/stash_test.go b/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")
+		}
+	})
+}