Add action name to context (#10)

* add action name to context * extract the const to the public package * add a key type
clear-street · Jan 6, 2025 · 93c7784 · 93c7784
1 parent 677cd7b
commit 93c7784
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/internal/generator/generator.go b/internal/generator/generator.go
@@ -13,7 +13,9 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-var fileHeader = "Code generated by reinforcer, DO NOT EDIT."
+var (
+	fileHeader = "Code generated by reinforcer, DO NOT EDIT."
+)
 
 // FileConfig holds the code generation configuration for a specific type
 type FileConfig struct {
@@ -249,6 +251,7 @@ func generateCommon(outPkg string) (string, error) {
 		jen.Id("name").Id("string"),
 		jen.Id("fn").Func().Params(jen.Id("ctx").Qual("context", "Context")).Id("error"),
 	).Id("error").Block(
+		jen.Id("ctx").Op("=").Qual("context", "WithValue").Params(jen.Id("ctx"), jen.Qual("github.com/clear-street/reinforcer/pkg/runner", "ActionNameKey").Block(), jen.Id("name")),
 		jen.Return(jen.Id("b").Dot("runnerFactory").Dot("GetRunner").Call(jen.Id("name")).Dot("Run").Call(jen.Id("ctx"), jen.Id("fn"))),
 	))
 	return renderToString(f)

diff --git a/internal/generator/generator_test.go b/internal/generator/generator_test.go
@@ -49,6 +49,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -72,6 +73,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -190,6 +192,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -213,6 +216,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -386,6 +390,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -409,6 +414,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -507,6 +513,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -530,6 +537,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -622,6 +630,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -645,6 +654,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -771,6 +781,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -794,6 +805,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,
@@ -883,6 +895,7 @@ package resilient
 
 import (
 	"context"
+	runner "github.com/clear-street/reinforcer/pkg/runner"
 	goresilience "github.com/slok/goresilience"
 )
 
@@ -906,6 +919,7 @@ func WithRetryableErrorPredicate(fn func(string, error) bool) Option {
 	}
 }
 func (b *base) run(ctx context.Context, name string, fn func(ctx context.Context) error) error {
+	ctx = context.WithValue(ctx, runner.ActionNameKey{}, name)
 	return b.runnerFactory.GetRunner(name).Run(ctx, fn)
 }
 `,

diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
@@ -1,10 +1,14 @@
 package runner
 
 import (
-	"github.com/slok/goresilience"
 	"sync"
+
+	"github.com/slok/goresilience"
 )
 
+// ActionNameKey is the key used to store the action name in the context.
+type ActionNameKey struct{}
+
 // Factory of runners
 type Factory struct {
 	mu          sync.RWMutex