PR feedback.

Author: codemzs

src/Microsoft.ML.Core/Data/IHostEnvironment.cs

@@ -72,7 +72,7 @@ public interface IHostEnvironment : IChannelProvider, IProgressChannelProvider
     internal interface ICancelableHost : IHost
     {
         /// <summary>
-        /// Signal to stop exection in this host and all its children.
+        /// Signal to stop exection in this host.
         /// </summary>
         void CancelExecution();
 

src/Microsoft.ML.Core/Environment/HostEnvironmentBase.cs

@@ -96,22 +96,17 @@ internal interface IMessageSource
     internal abstract class HostEnvironmentBase<TEnv> : ChannelProviderBase, IHostEnvironment, IChannelProvider
         where TEnv : HostEnvironmentBase<TEnv>
     {
-        public bool IsCanceled { get; protected set; }
-
-        private readonly List<IHost> _hosts;
-
-        private readonly object _cancelEnvLock;
-
         [BestFriend]
         internal void CancelExecutionHosts()
         {
-            lock (_cancelEnvLock)
+            lock (_cancelLock)
             {
-                foreach (var host in _hosts)
-                    if (host is ICancelableHost)
-                        ((ICancelableHost)host).CancelExecution();
+                foreach (var child in _children)
+                    if (child.TryGetTarget(out IHost host))
+                        if (host is ICancelableHost cancelableHost)
+                            cancelableHost.CancelExecution();
 
-                _hosts.Clear();
+                _children.Clear();
             }
         }
 
@@ -126,14 +121,10 @@ public abstract class HostBase : HostEnvironmentBase<TEnv>, ICancelableHost
 
             public Random Rand => _rand;
 
-            // We don't have dispose mechanism for hosts, so to let GC collect children hosts we make them WeakReference.
-            private readonly List<WeakReference<IHost>> _children;
-
             public HostBase(HostEnvironmentBase<TEnv> source, string shortName, string parentFullName, Random rand, bool verbose)
                 : base(source, rand, verbose, shortName, parentFullName)
             {
                 Depth = source.Depth + 1;
-                _children = new List<WeakReference<IHost>>();
             }
 
             public void CancelExecution()
@@ -195,7 +186,7 @@ protected PipeBase(ChannelProviderBase parent, string shortName,
 
             public void Dispose()
             {
-                if(!_disposed)
+                if (!_disposed)
                 {
                     Dispose(true);
                     _disposed = true;
@@ -363,6 +354,11 @@ public void RemoveListener(Action<IMessageSource, TMessage> listenerFunc)
 
         public override int Depth => 0;
 
+        public bool IsCanceled { get; protected set; }
+
+        // We don't have dispose mechanism for hosts, so to let GC collect children hosts we make them WeakReference.
+        private readonly List<WeakReference<IHost>> _children;
+
         /// <summary>
         ///  The main constructor.
         /// </summary>
@@ -375,10 +371,9 @@ protected HostEnvironmentBase(Random rand, bool verbose,
             ListenerDict = new ConcurrentDictionary<Type, Dispatcher>();
             ProgressTracker = new ProgressReporting.ProgressTracker(this);
             _cancelLock = new object();
-            _cancelEnvLock = new object();
             Root = this as TEnv;
             ComponentCatalog = new ComponentCatalog();
-            _hosts = new List<IHost>();
+            _children = new List<WeakReference<IHost>>();
         }
 
         /// <summary>
@@ -392,26 +387,25 @@ protected HostEnvironmentBase(HostEnvironmentBase<TEnv> source, Random rand, boo
             Contracts.CheckValueOrNull(rand);
             _rand = rand ?? RandomUtils.Create();
             _cancelLock = new object();
-            _cancelEnvLock = new object();
-            _hosts = new List<IHost>();
 
             // This fork shares some stuff with the master.
             Master = source;
             Root = source.Root;
             ListenerDict = source.ListenerDict;
             ProgressTracker = source.ProgressTracker;
             ComponentCatalog = source.ComponentCatalog;
+            _children = new List<WeakReference<IHost>>();
         }
 
         public IHost Register(string name, int? seed = null, bool? verbose = null)
         {
             Contracts.CheckNonEmpty(name, nameof(name));
             IHost host;
-            lock (_cancelEnvLock)
+            lock (_cancelLock)
             {
                 Random rand = (seed.HasValue) ? RandomUtils.Create(seed.Value) : RandomUtils.Create(_rand);
                 host = RegisterCore(this, name, Master?.FullName, rand, verbose ?? Verbose);
-                _hosts.Add(host);
+                _children.Add(new WeakReference<IHost>(host));
             }
             return host;
         }

src/Microsoft.ML.Data/MLContext.cs

@@ -73,8 +73,6 @@ public sealed class MLContext : IHostEnvironment
         /// </summary>
         public ComponentCatalog ComponentCatalog => _env.ComponentCatalog;
 
-        private List<IHost> _hosts;
-
         /// <summary>
         /// Create the ML context.
         /// </summary>
@@ -93,7 +91,6 @@ public MLContext(int? seed = null)
             Transforms = new TransformsCatalog(_env);
             Model = new ModelOperationsCatalog(_env);
             Data = new DataOperationsCatalog(_env);
-            _hosts = new List<IHost>();
         }
 
         private void ProcessMessage(IMessageSource source, ChannelMessage message)

test/Microsoft.ML.Core.Tests/UnitTests/TestHosts.cs

@@ -84,18 +84,14 @@ public void TestCancellationApi()
             hosts.Add(new Tuple<IHost, int>(mainHost.Register("3"), 1));
             hosts.Add(new Tuple<IHost, int>(mainHost.Register("4"), 1));
             hosts.Add(new Tuple<IHost, int>(mainHost.Register("5"), 1));
-            var addThread = new Thread(
-            () =>
+
+            for (int i = 0; i < 5; i++)
             {
-                for (int i = 0; i < 5; i++)
-                {
-                    var tupple = hosts.ElementAt(i);
-                    var newHost = tupple.Item1.Register((tupple.Item2 + 1).ToString());
-                    hosts.Add(new Tuple<IHost, int>(newHost, tupple.Item2 + 1));
-                }
-            });
-            addThread.Start();
-            addThread.Join();
+                var tupple = hosts.ElementAt(i);
+                var newHost = tupple.Item1.Register((tupple.Item2 + 1).ToString());
+                hosts.Add(new Tuple<IHost, int>(newHost, tupple.Item2 + 1));
+            }
+
             ((MLContext)env).CancelExecution();
 
             //Ensure all created hosts are cancelled.