feat(contextx): generic context injector and loader

contextx/context.go

@@ -1,6 +1,10 @@
 package contextx
 
-import "context"
+import (
+	"context"
+
+	"github.com/pkg/errors"
+)
 
 type WithContext = func(ctx context.Context) context.Context
 
@@ -12,3 +16,69 @@ func WithContextCompose(withs ...WithContext) WithContext {
 		return ctx
 	}
 }
+
+type (
+	Option[T any] func(*ctx[T])
+	Valuer[T any] func() T
+)
+
+func With[T any](v T) Option[T] {
+	return func(c *ctx[T]) {
+		c.valuer = func() T {
+			return v
+		}
+	}
+}
+
+func WithValuer[T any](valuer Valuer[T]) Option[T] {
+	return func(c *ctx[T]) {
+		c.valuer = valuer
+	}
+}
+
+func New[T any](options ...Option[T]) Context[T] {
+	c := &ctx[T]{}
+	for _, option := range options {
+		if option != nil {
+			option(c)
+		}
+	}
+	return c
+}
+
+func NewValue[T any](v T) Context[T] {
+	c := &ctx[T]{}
+	With(v)(c)
+	return c
+}
+
+type Context[T any] interface {
+	With(context.Context, T) context.Context
+	From(context.Context) (T, bool)
+	MustFrom(context.Context) T
+}
+
+type ctx[T any] struct {
+	valuer Valuer[T]
+}
+
+func (c *ctx[T]) With(ctx context.Context, v T) context.Context {
+	return WithValue(ctx, c, v)
+}
+
+func (c *ctx[T]) MustFrom(ctx context.Context) T {
+	if v, ok := ctx.Value(c).(T); ok {
+		return v
+	}
+	if c.valuer != nil {
+		return c.valuer()
+	}
+	panic(errors.Errorf("%T not found in context", c))
+}
+
+func (c *ctx[T]) From(ctx context.Context) (T, bool) {
+	if v, ok := ctx.Value(c).(T); ok {
+		return v, ok
+	}
+	return *new(T), false
+}

contextx/context_test.go

@@ -0,0 +1,115 @@
+package contextx_test
+
+import (
+	"context"
+	"testing"
+
+	. "github.com/onsi/gomega"
+
+	"github.com/xoctopus/x/contextx"
+)
+
+type Value struct {
+	Int int
+}
+
+var ValueContext = contextx.NewValue(&Value{1})
+
+func TestContext(t *testing.T) {
+	empty := context.Background()
+
+	t.Run("FailedToExtract", func(t *testing.T) {
+		t.Run("MustFrom", func(t *testing.T) {
+			defer func() {
+				err := recover().(error)
+				NewWithT(t).Expect(err.Error()).To(ContainSubstring("not found in context"))
+			}()
+
+			ctx := contextx.New[*Value](nil)
+			_ = ctx.MustFrom(empty)
+		})
+
+		t.Run("From", func(t *testing.T) {
+			ctx := contextx.New[*Value](contextx.With[*Value](&Value{2}))
+			v, ok := ctx.From(empty)
+			NewWithT(t).Expect(ok).To(BeFalse())
+			NewWithT(t).Expect(v).To(BeNil())
+		})
+	})
+
+	t.Run("FromValuer", func(t *testing.T) {
+		t.Run("New", func(t *testing.T) {
+			t.Run("With", func(t *testing.T) {
+				ctx := contextx.New(contextx.With(&Value{2}))
+				v := ctx.MustFrom(empty)
+				NewWithT(t).Expect(v).NotTo(BeNil())
+				NewWithT(t).Expect(v.Int).To(Equal(2))
+			})
+			t.Run("WithValuer", func(t *testing.T) {
+				ctx := contextx.New(contextx.WithValuer(func() *Value { return &Value{3} }))
+				v := ctx.MustFrom(empty)
+				NewWithT(t).Expect(v).NotTo(BeNil())
+				NewWithT(t).Expect(v.Int).To(Equal(3))
+			})
+		})
+		t.Run("NewValue", func(t *testing.T) {
+			ctx := contextx.NewValue(&Value{4})
+			v := ctx.MustFrom(empty)
+			NewWithT(t).Expect(v).NotTo(BeNil())
+			NewWithT(t).Expect(v.Int).To(Equal(4))
+		})
+	})
+
+	t.Run("OverwriteByInjection", func(t *testing.T) {
+		ctx := contextx.NewValue(&Value{5})
+		root := ctx.With(empty, &Value{6})
+
+		v, ok := ctx.From(root)
+		NewWithT(t).Expect(ok).To(BeTrue())
+		NewWithT(t).Expect(v.Int).To(Equal(6))
+
+		v = ctx.MustFrom(root)
+		NewWithT(t).Expect(v.Int).To(Equal(6))
+	})
+}
+
+func BenchmarkCtx_MustFrom(b *testing.B) {
+	empty := context.Background()
+	injected := ValueContext.With(empty, &Value{1})
+
+	b.Run("FromValuer", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ValueContext.MustFrom(empty)
+		}
+	})
+
+	b.Run("ExtractFromContext", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ValueContext.MustFrom(injected)
+		}
+	})
+
+	overinjected := context.Background()
+	for i := 0; i < 1000; i++ {
+		overinjected = context.WithValue(overinjected, i, i)
+	}
+	overinjected = ValueContext.With(overinjected, &Value{1})
+
+	b.Run("OverInjected", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ValueContext.MustFrom(overinjected)
+		}
+	})
+
+	overinjected2 := context.Background()
+	for i := 0; i < 1000; i++ {
+		overinjected = contextx.WithValue(overinjected, i, i)
+	}
+	overinjected2 = ValueContext.With(overinjected, &Value{1})
+
+	b.Run("OverInjected2", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ValueContext.MustFrom(overinjected2)
+		}
+	})
+}

go.mod

@@ -5,6 +5,7 @@ go 1.22.0
 require (
 	github.com/onsi/gomega v1.33.1
 	github.com/pkg/errors v0.9.1
+	golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7
 	golang.org/x/tools v0.23.0
 )
 

go.sum

@@ -12,6 +12,8 @@ github.com/onsi/gomega v1.33.1 h1:dsYjIxxSR755MDmKVsaFQTE22ChNBcuuTWgkUDSubOk=
 github.com/onsi/gomega v1.33.1/go.mod h1:U4R44UsT+9eLIaYRB2a5qajjtQYn0hauxvRm16AVYg0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7 h1:wDLEX9a7YQoKdKNQt88rtydkqDxeGaBUTnIYc3iG/mA=
+golang.org/x/exp v0.0.0-20240716175740-e3f259677ff7/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY=
 golang.org/x/mod v0.19.0 h1:fEdghXQSo20giMthA7cd28ZC+jts4amQ3YMXiP5oMQ8=
 golang.org/x/mod v0.19.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
 golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys=