private/mud: support registering multiple implementations

Change-Id: I5df24fbc799bd11311606d8b285fd64fbad9d93c

private/mud/implementation.go

@@ -0,0 +1,43 @@
+// Copyright (C) 2024 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package mud
+
+import (
+	"context"
+)
+
+// Implementation registers a new []T component, which will be filled with any registered instances.
+// Instances will be marked with "Optional{}" tag, and will be injected only, if they are initialized.
+// It's the responsibility of the Init code to exclude / include them during initialization.
+func Implementation[L ~[]T, Instance any, T any](ball *Ball) {
+	if lookup[L](ball) == nil {
+		RegisterManual[L](ball, func(ctx context.Context) (L, error) {
+			var instances L
+			component := lookup[L](ball)
+			for _, req := range component.requirements {
+				c, _ := lookupByType(ball, req)
+				// only initialized instances are inject to the implementation list
+				if c.instance != nil {
+					instances = append(instances, c.instance.(T))
+				}
+			}
+			return instances, nil
+		})
+	}
+	lookup[L](ball).requirements = append(lookup[L](ball).requirements, typeOf[Instance]())
+	Tag[Instance](ball, Optional{})
+}
+
+// ImplementationOf is a ForEach filter to get all the dependency of an implementation.
+func ImplementationOf[L ~[]T, T any](ball *Ball) ComponentSelector {
+	component := MustLookupComponent[L](ball)
+	return func(c *Component) bool {
+		for _, dep := range component.requirements {
+			if dep == c.target {
+				return true
+			}
+		}
+		return false
+	}
+}

private/mud/implementation_test.go

@@ -0,0 +1,121 @@
+// Copyright (C) 2024 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package mud
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"storj.io/common/testcontext"
+)
+
+func TestImplementationAllInit(t *testing.T) {
+	ctx := testcontext.New(t)
+
+	ball := NewBall()
+	Provide[PostgresDB](ball, func() PostgresDB {
+		return PostgresDB{}
+	})
+	Provide[CockroachDB](ball, func() CockroachDB {
+		return CockroachDB{}
+	})
+
+	Implementation[[]DBAdapter, PostgresDB](ball)
+	Implementation[[]DBAdapter, CockroachDB](ball)
+
+	err := ForEach(ball, Initialize(ctx), All)
+	require.NoError(t, err)
+
+	// by default all dependencies are initialized, and usable
+	err = Execute0(ctx, ball, func(impl []DBAdapter) {
+		require.Len(t, impl, 2)
+	})
+	require.NoError(t, err)
+}
+
+func TestImplementationOneInit(t *testing.T) {
+	ctx := testcontext.New(t)
+
+	ball := NewBall()
+	Provide[PostgresDB](ball, func() PostgresDB {
+		return PostgresDB{}
+	})
+	Provide[CockroachDB](ball, func() CockroachDB {
+		return CockroachDB{}
+	})
+
+	Implementation[[]DBAdapter, PostgresDB](ball)
+	Implementation[[]DBAdapter, CockroachDB](ball)
+
+	pg := MustLookupComponent[PostgresDB](ball)
+	err := pg.Init(ctx)
+	require.NoError(t, err)
+
+	adapters := MustLookupComponent[[]DBAdapter](ball)
+
+	// this init will use all the initialized dependencies (in our case, postgres only)
+	err = adapters.Init(ctx)
+	require.NoError(t, err)
+
+	// Cockroach is not initialized, therefore it's an optional dependency.
+	err = Execute0(ctx, ball, func(impl []DBAdapter) {
+		require.Len(t, impl, 1)
+		require.Equal(t, "postgres", impl[0].Name())
+	})
+	require.NoError(t, err)
+
+}
+
+func TestImplementationMarkedOptional(t *testing.T) {
+	ctx := testcontext.New(t)
+
+	ball := NewBall()
+	Provide[PostgresDB](ball, func() PostgresDB {
+		return PostgresDB{}
+	})
+	Provide[CockroachDB](ball, func() CockroachDB {
+		return CockroachDB{}
+	})
+
+	Implementation[[]DBAdapter, PostgresDB](ball)
+	Implementation[[]DBAdapter, CockroachDB](ball)
+	// PostgresDB is required
+	RemoveTag[PostgresDB, Optional](ball)
+
+	// initialize the non-optional components
+	// NOTE: optional is just a flag to find the right components
+	// it's the responsibility of the caller to initialize the right components.
+	for _, component := range Find(ball, And(All, func(c *Component) bool {
+		_, found := GetTagOf[Optional](c)
+		return !found
+	})) {
+		// this will ignore CockroachDB as it's not a required dependency
+		err := component.Init(ctx)
+		require.NoError(t, err)
+	}
+
+	// CockroachDB is not initialized, because it was optional
+	err := Execute0(ctx, ball, func(impl []DBAdapter) {
+		require.Len(t, impl, 1)
+		require.Equal(t, "postgres", impl[0].Name())
+	})
+	require.NoError(t, err)
+
+}
+
+type DBAdapter interface {
+	Name() string
+}
+type PostgresDB struct{}
+
+func (p PostgresDB) Name() string {
+	return "postgres"
+}
+
+type CockroachDB struct{}
+
+func (p CockroachDB) Name() string {
+	return "cockroach"
+}

private/mud/mud.go

@@ -23,7 +23,7 @@ func NewBall() *Ball {
 	return &Ball{}
 }
 
-// getLogger returns with the zap Logger, if component is registered.
+// getLogger returns with the zap Logger, i!f component is registered.
 // used for internal logging.
 func (ball *Ball) getLogger() *zap.Logger {
 	if logger := lookup[*zap.Logger](ball); logger != nil {
@@ -60,7 +60,7 @@ func RegisterManual[T any](
 
 // Tag attaches a tag to the component registration.
 func Tag[A any, Tag any](ball *Ball, tag Tag) {
-	c := mustLookup[A](ball)
+	c := MustLookupComponent[A](ball)
 
 	// we don't allow duplicated registrations, as we always return with the first value.
 	for ix, existing := range c.tags {
@@ -73,9 +73,28 @@ func Tag[A any, Tag any](ball *Ball, tag Tag) {
 	c.tags = append(c.tags, tag)
 }
 
+// RemoveTag removes all the Tag type of tags from the component.
+func RemoveTag[A any, Tag any](ball *Ball) {
+	c := MustLookupComponent[A](ball)
+	var filtered []any
+	// we don't allow duplicated registrations, as we always return with the first value.
+	for ix, existing := range c.tags {
+		_, ok := existing.(Tag)
+		if !ok {
+			filtered = append(filtered, c.tags[ix])
+		}
+	}
+	c.tags = filtered
+}
+
 // GetTag returns with attached tag (if attached).
 func GetTag[A any, Tag any](ball *Ball) (Tag, bool) {
-	c := mustLookup[A](ball)
+	c := MustLookupComponent[A](ball)
+	return findTag[Tag](c)
+}
+
+// GetTagOf returns with attached tag (if attached).
+func GetTagOf[Tag any](c *Component) (Tag, bool) {
 	return findTag[Tag](c)
 }
 
@@ -93,7 +112,7 @@ func findTag[Tag any](c *Component) (Tag, bool) {
 // DependsOn creates a dependency relation between two components.
 // With the help of the dependency graph, they can be executed in the right order.
 func DependsOn[BASE any, DEPENDENCY any](ball *Ball) {
-	c := mustLookup[BASE](ball)
+	c := MustLookupComponent[BASE](ball)
 	c.addRequirement(typeOf[DEPENDENCY]())
 }
 
@@ -288,13 +307,13 @@ func Supply[T any](ball *Ball, t T) {
 }
 
 // View is lightweight component, which provides a type based on a existing instances.
-func View[A any, B any](ball *Ball, convert func(A) B) {
-	RegisterManual[B](ball, func(ctx context.Context) (B, error) {
-		a := mustLookup[A](ball)
-		return convert(a.instance.(A)), nil
+func View[From any, To any](ball *Ball, convert func(From) To) {
+	RegisterManual[To](ball, func(ctx context.Context) (To, error) {
+		a := MustLookupComponent[From](ball)
+		return convert(a.instance.(From)), nil
 	})
-	component := mustLookup[B](ball)
-	component.requirements = append(component.requirements, mustLookup[A](ball).target)
+	component := MustLookupComponent[To](ball)
+	component.requirements = append(component.requirements, MustLookupComponent[From](ball).target)
 }
 
 // Dereference is a simple transformation to make real value from a pointer. Useful with View.
@@ -319,7 +338,8 @@ func lookup[T any](ball *Ball) *Component {
 	return nil
 }
 
-func mustLookup[T any](ball *Ball) *Component {
+// MustLookupComponent gets the component (or panics if doesn't exist) based on a type.
+func MustLookupComponent[T any](ball *Ball) *Component {
 	c := lookup[T](ball)
 	if c == nil {
 		panic("component is missing: " + name[T]())
@@ -346,7 +366,7 @@ func mustLookupByType(ball *Ball, tzpe reflect.Type) *Component {
 
 // MustLookup returns with the registered component instance (or panic).
 func MustLookup[T any](ball *Ball) T {
-	component := mustLookup[T](ball)
+	component := MustLookupComponent[T](ball)
 	if component.instance == nil {
 		panic("lookup of an uninitialized component " + name[T]())
 	}

private/mud/selector.go

@@ -6,7 +6,7 @@ package mud
 // Select is a component selector based on the specified type ([A]).
 func Select[A any](ball *Ball) ComponentSelector {
 	t := typeOf[A]()
-	mustLookup[A](ball)
+	MustLookupComponent[A](ball)
 	return func(c *Component) bool {
 		return c.target == t
 	}