Fixed #849: Store build callbacks in property bag so modules can add …

…to them.
autofac · May 2, 2017 · a12b3ab · a12b3ab
1 parent 43f390b
commit a12b3ab
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/src/Autofac/ContainerBuilder.cs b/src/Autofac/ContainerBuilder.cs
@@ -63,9 +63,10 @@ namespace Autofac
     public class ContainerBuilder
     {
         private readonly IList<DeferredCallback> _configurationCallbacks = new List<DeferredCallback>();
-        private readonly IList<Action<IContainer>> _buildCallbacks = new List<Action<IContainer>>();
         private bool _wasBuilt;
 
+        private const string BuildCallbackPropertyKey = "__BuildCallbackKey";
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ContainerBuilder"/> class.
         /// </summary>
@@ -81,6 +82,11 @@ public ContainerBuilder()
         internal ContainerBuilder(IDictionary<string, object> properties)
         {
             Properties = properties;
+
+            if (!Properties.ContainsKey(BuildCallbackPropertyKey))
+            {
+                Properties.Add(BuildCallbackPropertyKey, new List<Action<IContainer>>());
+            }
         }
 
         /// <summary>
@@ -115,7 +121,8 @@ public ContainerBuilder RegisterBuildCallback(Action<IContainer> buildCallback)
         {
             if (buildCallback == null) throw new ArgumentNullException(nameof(buildCallback));
 
-            _buildCallbacks.Add(buildCallback);
+            var buildCallbacks = GetBuildCallbacks();
+            buildCallbacks.Add(buildCallback);
 
             return this;
         }
@@ -140,7 +147,8 @@ public IContainer Build(ContainerBuildOptions options = ContainerBuildOptions.No
             if ((options & ContainerBuildOptions.IgnoreStartableComponents) == ContainerBuildOptions.None)
                 StartStartableComponents(result);
 
-            foreach (var buildCallback in _buildCallbacks)
+            var buildCallbacks = GetBuildCallbacks();
+            foreach (var buildCallback in buildCallbacks)
                 buildCallback(result);
 
             return result;
@@ -281,5 +289,10 @@ private void RegisterDefaultAdapters(IComponentRegistry componentRegistry)
             componentRegistry.AddRegistrationSource(new StronglyTypedMetaRegistrationSource());
             componentRegistry.AddRegistrationSource(new GeneratedFactoryRegistrationSource());
         }
+
+        private List<Action<IContainer>> GetBuildCallbacks()
+        {
+            return (List<Action<IContainer>>)Properties[BuildCallbackPropertyKey];
+        }
     }
 }
diff --git a/test/Autofac.Test/ContainerBuilderTests.cs b/test/Autofac.Test/ContainerBuilderTests.cs
@@ -589,5 +589,33 @@ void BuildCallback(IContainer c)
 
             Assert.Equal(2, called);
         }
+
+        [Fact]
+        public void BuildCallbacksInvokedWhenRegisteredInModuleLoad()
+        {
+            var module = new BuildCallbackModule();
+
+            var builder = new ContainerBuilder();
+            builder.RegisterModule(module);
+            builder.Build();
+
+            Assert.Equal(2, module.Called);
+        }
+
+        public class BuildCallbackModule : Module
+        {
+            public int Called { get; private set; }
+
+            protected override void Load(ContainerBuilder builder)
+            {
+                void BuildCallback(IContainer c)
+                {
+                    Called++;
+                }
+
+                builder.RegisterBuildCallback(BuildCallback)
+                    .RegisterBuildCallback(BuildCallback);
+            }
+        }
     }
 }