diff --git a/container.go b/container.go index 61a592f..dbae029 100644 --- a/container.go +++ b/container.go @@ -9,15 +9,15 @@ import ( type ( Container[C any] struct { - config C + config C + // immutable once container is built. No lock needed. Each factory is a singleflight and thus + // holds its own lock for service construction. factories map[string]func(context.Context, *Container[C]) (any, error) - services map[string]any - sync.Mutex - frozen atomic.Bool } ContainerBuilder[C any] struct { container *Container[C] + mu sync.Mutex frozen atomic.Bool } @@ -26,12 +26,12 @@ type ( ServiceID[T any] string ) +// New initializes a new ContainerBuilder func New[C any](config C) *ContainerBuilder[C] { c := &Container[C]{ config: config, factories: make(map[string]func(context.Context, *Container[C]) (any, error)), - services: make(map[string]any), } return &ContainerBuilder[C]{ @@ -43,6 +43,8 @@ func (c *ContainerBuilder[C]) isFrozen() bool { return c.frozen.Load() } +// Build wil build the Dependency Container +// If the ContainerBuilder is already build, it will return an error, otherwhise it'll return the container func (c *ContainerBuilder[C]) Build() (*Container[C], error) { if c.isFrozen() { return nil, ErrContainerSealed @@ -56,21 +58,55 @@ func (c *Container[C]) Config() C { return c.config } +// Inject injects a final type into the container. +// This is a convenient function as it just wraps the final type into a factory function +func Inject[T, C any](cc *ContainerBuilder[C], id ServiceID[T], obj T) error { + return Register(cc, id, func(ctx context.Context, c *Container[C]) (T, error) { + return obj, nil + }) +} + func Register[T, C any](cc *ContainerBuilder[C], id ServiceID[T], fac Factory[T, C]) error { if cc.isFrozen() { var v T return fmt.Errorf("cannot register %s[%T]: %w", id, v, ErrContainerSealed) } - cc.container.Lock() - defer cc.container.Unlock() + cc.mu.Lock() + defer cc.mu.Unlock() - cc.container.factories[string(id)] = func(ctx context.Context, c *Container[C]) (any, error) { + cc.container.factories[string(id)] = singleFlight(func(ctx context.Context, c *Container[C]) (any, error) { return fac(ctx, c) - } + }) return nil } +// singleFlight wraps a factory so it is invoked at most once successfully +// +// it can fail multiple times, so service construction recovery is possible. +// +// Concurrent calls to the factory will block until the first call returns. +func singleFlight[C any](fac func(context.Context, *Container[C]) (any, error)) func(context.Context, *Container[C]) (any, error) { + var ( + mu sync.Mutex + done bool + val any + ) + return func(ctx context.Context, c *Container[C]) (any, error) { + mu.Lock() + defer mu.Unlock() + if done { + return val, nil + } + v, err := fac(ctx, c) + if err != nil { + return nil, err + } + val, done = v, true + return val, nil + } +} + func Get[T, C any](ctx context.Context, cc *Container[C], service ServiceID[T]) (T, error) { return get(ctx, cc, service) } @@ -91,22 +127,11 @@ func get[T, C any](ctx context.Context, c *Container[C], service ServiceID[T]) ( } func fromContainer[T, C any](ctx context.Context, c *Container[C], id ServiceID[T]) (any, error) { - s, ok := c.services[string(id)] - if ok { - return s, nil - } - f, ok := c.factories[string(id)] if !ok { return nil, fmt.Errorf("no factory is registered for %+v", id) } - srv, err := f(ctx, c) - if err != nil { - return nil, err - } - - c.services[string(id)] = srv - return srv, nil + return f(ctx, c) } diff --git a/container_test.go b/container_test.go index 06a4d7b..f25a12f 100644 --- a/container_test.go +++ b/container_test.go @@ -3,6 +3,8 @@ package di import ( "context" "errors" + "sync" + "sync/atomic" "testing" "github.com/stretchr/testify/assert" @@ -21,6 +23,10 @@ type ( B struct { BName string } + CD struct { + *A + *B + } ) func TestContainerRegister(t *testing.T) { @@ -54,6 +60,163 @@ func TestContainerRegister(t *testing.T) { assert.Equal(t, c.Config().B, b.BName) } +func TestConcurrentResolve(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + cb := New(&Config{A: "a", B: "b"}) + + var err error + + const ( + sidA ServiceID[*A] = "service.a" + sidB ServiceID[*B] = "service.b" + sidC ServiceID[*CD] = "service.c" + ) + + err = Register(cb, sidA, func(ctx context.Context, cc *Container[*Config]) (*A, error) { + return &A{AName: cc.Config().A}, nil + }) + err = Register(cb, sidB, func(ctx context.Context, cc *Container[*Config]) (*B, error) { + return &B{BName: cc.Config().B}, nil + }) + + err = Register(cb, sidC, func(ctx context.Context, cc *Container[*Config]) (*CD, error) { + a, err := Get(ctx, cc, sidA) + if err != nil { + return nil, err + } + b, err := Get(ctx, cc, sidB) + if err != nil { + return nil, err + } + return &CD{A: a, B: b}, nil + }) + + assert.NoError(t, err) + + assert.NoError(t, err) + + c, err := cb.Build() + assert.NoError(t, err) + + wg := &sync.WaitGroup{} + + resultc := make(chan *CD, 100) + errc := make(chan error, 100) + + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + res, err := Get(ctx, c, sidC) + if err != nil { + errc <- err + return + } + + resultc <- res + }() + } + + go func() { + wg.Wait() + close(resultc) + close(errc) + }() + + res, err := Get(ctx, c, sidC) + assert.NoError(t, err) + + for a := range resultc { + assert.Same(t, res, a) + } + + for err := range errc { + t.Error(err) + } +} + +// TestConcurrentNestedResolve exercises the case the plain TestConcurrentResolve +// misses: a factory that resolves a dependency via Get while being resolved +// concurrently. It guards against the self-deadlock that occurs when the +// container lock is held across factory execution, and asserts each service is +// constructed exactly once (single-flight) and shared. +func TestConcurrentNestedResolve(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + cb := New(&Config{A: "a", B: "b"}) + + var aBuilds, bBuilds atomic.Int64 + sidA := ServiceID[*A]("service.a") + sidB := ServiceID[*B]("service.b") + + assert.NoError(t, Register(cb, sidB, func(ctx context.Context, cc *Container[*Config]) (*B, error) { + bBuilds.Add(1) + return &B{BName: cc.Config().B}, nil + })) + // A depends on B: nested Get inside A's factory. + assert.NoError(t, Register(cb, sidA, func(ctx context.Context, cc *Container[*Config]) (*A, error) { + aBuilds.Add(1) + b, err := Get(ctx, cc, sidB) + if err != nil { + return nil, err + } + return &A{AName: cc.Config().A + b.BName}, nil + })) + + c, err := cb.Build() + assert.NoError(t, err) + + const n = 50 + wg := &sync.WaitGroup{} + results := make([]*A, n) + errs := make([]error, n) + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + results[i], errs[i] = Get(ctx, c, sidA) + }(i) + } + wg.Wait() + + first := results[0] + for i := 0; i < n; i++ { + assert.NoError(t, errs[i]) + assert.Same(t, first, results[i]) // all share the one singleton + } + assert.Equal(t, int64(1), aBuilds.Load(), "A built exactly once") + assert.Equal(t, int64(1), bBuilds.Load(), "B built exactly once") +} + +// TestResolveErrorNotCached asserts a factory error is not memoized: a later Get +// retries the factory. xstream's startup Retry relies on this behavior. +func TestResolveErrorNotCached(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + cb := New(&Config{A: "a", B: "b"}) + + var calls atomic.Int64 + sidA := ServiceID[*A]("service.a") + assert.NoError(t, Register(cb, sidA, func(ctx context.Context, cc *Container[*Config]) (*A, error) { + if calls.Add(1) == 1 { + return nil, errors.New("transient failure") + } + return &A{AName: "ok"}, nil + })) + + c, err := cb.Build() + assert.NoError(t, err) + + _, err = Get(ctx, c, sidA) + assert.Error(t, err) // first attempt fails and is NOT cached + + a, err := Get(ctx, c, sidA) + assert.NoError(t, err) // retry succeeds + assert.Equal(t, "ok", a.AName) + assert.Equal(t, int64(2), calls.Load()) +} + func TestContainerSeal(t *testing.T) { cb := New(&Config{A: "a", B: "b"})