fix!(container): fix issue with providing a module-scoped dependency from within a module (#11913)

container/container.go

@@ -133,6 +133,8 @@ func (c *container) getResolver(typ reflect.Type) (resolver, error) {
 	return c.resolvers[typ], nil
 }
 
+var stringType = reflect.TypeOf("")
+
 func (c *container) addNode(provider *ProviderDescriptor, key *moduleKey) (interface{}, error) {
 	providerGraphNode, err := c.locationGraphNode(provider.Location, key)
 	if err != nil {
@@ -140,12 +142,17 @@ func (c *container) addNode(provider *ProviderDescriptor, key *moduleKey) (inter
 	}
 
 	hasModuleKeyParam := false
+	hasOwnModuleKeyParam := false
 	for _, in := range provider.Inputs {
 		typ := in.Type
 		if typ == moduleKeyType {
 			hasModuleKeyParam = true
 		}
 
+		if typ == ownModuleKeyType {
+			hasOwnModuleKeyParam = true
+		}
+
 		if isAutoGroupType(typ) {
 			return nil, fmt.Errorf("auto-group type %v can't be used as an input parameter", typ)
 		} else if isOnePerModuleType(typ) {
@@ -170,7 +177,7 @@ func (c *container) addNode(provider *ProviderDescriptor, key *moduleKey) (inter
 		c.addGraphEdge(typeGraphNode, providerGraphNode)
 	}
 
-	if key != nil || !hasModuleKeyParam {
+	if !hasModuleKeyParam {
 		c.logf("Registering %s", provider.Location.String())
 		c.indentLogger()
 		defer c.dedentLogger()
@@ -227,6 +234,11 @@ func (c *container) addNode(provider *ProviderDescriptor, key *moduleKey) (inter
 
 		return sp, nil
 	} else {
+		if hasOwnModuleKeyParam {
+			return nil, errors.Errorf("%T and %T must not be declared as dependencies on the same provided",
+				ModuleKey{}, OwnModuleKey{})
+		}
+
 		c.logf("Registering module-scoped provider: %s", provider.Location.String())
 		c.indentLogger()
 		defer c.dedentLogger()
@@ -314,6 +326,15 @@ func (c *container) resolve(in ProviderInput, moduleKey *moduleKey, caller Locat
 		return reflect.ValueOf(ModuleKey{moduleKey}), nil
 	}
 
+	if in.Type == ownModuleKeyType {
+		if moduleKey == nil {
+			return reflect.Value{}, errors.Errorf("trying to resolve %T for %s but not inside of any module's scope", moduleKey, caller)
+		}
+		c.logf("Providing OwnModuleKey %s", moduleKey.name)
+		markGraphNodeAsUsed(typeGraphNode)
+		return reflect.ValueOf(OwnModuleKey{moduleKey}), nil
+	}
+
 	vr, err := c.getResolver(in.Type)
 	if err != nil {
 		return reflect.Value{}, err

container/container_test.go

@@ -15,14 +15,13 @@ type KVStoreKey struct {
 	name string
 }
 
-type ModuleKey string
-
 type MsgClientA struct {
-	key ModuleKey
+	key string
 }
 
 type KeeperA struct {
 	key  KVStoreKey
+	name string
 }
 
 type KeeperB struct {
@@ -46,18 +45,14 @@ func ProvideKVStoreKey(moduleKey container.ModuleKey) KVStoreKey {
 	return KVStoreKey{name: moduleKey.Name()}
 }
 
-func ProvideModuleKey(moduleKey container.ModuleKey) (ModuleKey, error) {
-	return ModuleKey(moduleKey.Name()), nil
-}
-
-func ProvideMsgClientA(_ container.ModuleKey, key ModuleKey) MsgClientA {
-	return MsgClientA{key}
+func ProvideMsgClientA(key container.ModuleKey) MsgClientA {
+	return MsgClientA{key.Name()}
 }
 
 type ModuleA struct{}
 
-func (ModuleA) Provide(key KVStoreKey) (KeeperA, Handler, Command) {
-	return KeeperA{key}, Handler{}, Command{}
+func (ModuleA) Provide(key KVStoreKey, moduleKey container.OwnModuleKey) (KeeperA, Handler, Command) {
+	return KeeperA{key: key, name: container.ModuleKey(moduleKey).Name()}, Handler{}, Command{}
 }
 
 type ModuleB struct{}
@@ -76,7 +71,7 @@ type BProvides struct {
 	Commands []Command
 }
 
-func (ModuleB) Provide(dependencies BDependencies, _ container.ModuleKey) (BProvides, Handler, error) {
+func (ModuleB) Provide(dependencies BDependencies) (BProvides, Handler, error) {
 	return BProvides{
 		KeeperB: KeeperB{
 			key:        dependencies.Key,
@@ -96,6 +91,7 @@ func TestScenario(t *testing.T) {
 				require.Len(t, commands, 3)
 				require.Equal(t, KeeperA{
 					key:  KVStoreKey{name: "a"},
+					name: "a",
 				}, a)
 				require.Equal(t, KeeperB{
 					key: KVStoreKey{name: "b"},
@@ -104,11 +100,8 @@ func TestScenario(t *testing.T) {
 					},
 				}, b)
 			},
-			container.Provide(
-				ProvideKVStoreKey,
-				ProvideModuleKey,
-				ProvideMsgClientA,
-			),
+			container.Provide(ProvideMsgClientA),
+			container.ProvideInModule("runtime", ProvideKVStoreKey),
 			container.ProvideInModule("a", wrapMethod0(ModuleA{})),
 			container.ProvideInModule("b", wrapMethod0(ModuleB{})),
 		))

container/module_key.go

@@ -6,14 +6,17 @@ import (
 
 // ModuleKey is a special type used to scope a provider to a "module".
 //
-// Special module-scoped providers can be used with Provide by declaring a
-// provider with an input parameter of type ModuleKey. These providers
-// may construct a unique value of a dependency for each module and will
-// be called at most once per module.
+// Special module-scoped providers can be used with Provide and ProvideInModule
+// by declaring a provider with an input parameter of type ModuleKey. These
+// providers may construct a unique value of a dependency for each module and
+// will be called at most once per module.
 //
-// Providers passed to ProvideInModule can also declare an input parameter
-// of type ModuleKey to retrieve their module key but these providers will be
-// called at most once.
+// When being used with ProvideInModule, the provider will not receive its
+// own ModuleKey but rather the key of the module requesting the dependency
+// so that modules can provide module-scoped dependencies to other modules.
+//
+// In order for a module to retrieve their own module key they can define
+// a provider which requires the OwnModuleKey type and DOES NOT require ModuleKey.
 type ModuleKey struct {
 	*moduleKey
 }
@@ -28,4 +31,8 @@ func (k ModuleKey) Name() string {
 
 var moduleKeyType = reflect.TypeOf(ModuleKey{})
 
-var stringType = reflect.TypeOf("")
+// OwnModuleKey is a type which can be used in a module to retrieve its own
+// ModuleKey. It MUST NOT be used together with a ModuleKey dependency.
+type OwnModuleKey ModuleKey
+
+var ownModuleKeyType = reflect.TypeOf((*OwnModuleKey)(nil)).Elem()