From 93c7784ce93ec59eb9cfff081b7e6d93fa98e408 Mon Sep 17 00:00:00 2001
From: Tzach Shabtay <tshabtay@clearstreet.io>
Date: Mon, 6 Jan 2025 12:40:03 -0500
Subject: [PATCH] Add action name to context (#10)

* add action name to context

* extract the const to the public package

* add a key type
---
 internal/generator/generator.go      |  5 ++++-
 internal/generator/generator_test.go | 14 ++++++++++++++
 pkg/runner/runner.go                 |  6 +++++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/internal/generator/generator.go b/internal/generator/generator.go
index 1eabbea..16b17e0 100644
--- 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
index 6b5d08d..b60eb78 100644
--- 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
index 088e83c..bbf80fb 100644
--- 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