allow concurrent resolves

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)
 }

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"})