Ensuring that the ScriptHost is Disposed of in case of an exception in the run loop.

Author: fabiocav

src/WebJobs.Script/Host/ScriptHostManager.cs

@@ -19,6 +19,7 @@ namespace Microsoft.Azure.WebJobs.Script
     public class ScriptHostManager : IDisposable
     {
         private readonly ScriptHostConfiguration _config;
+        private readonly IScriptHostFactory _scriptHostFactory;
         private ScriptHost _currentInstance;
 
         // ScriptHosts are not thread safe, so be clear that only 1 thread at a time operates on each instance. 
@@ -33,8 +34,14 @@ public class ScriptHostManager : IDisposable
         private TraceWriter _traceWriter;
 
         public ScriptHostManager(ScriptHostConfiguration config)
+            : this(config, new ScriptHostFactory())
+        {
+        }
+
+        public ScriptHostManager(ScriptHostConfiguration config, IScriptHostFactory scriptHostFactory)
         {
             _config = config;
+            _scriptHostFactory = scriptHostFactory;
         }
 
         /// <summary>
@@ -57,6 +64,7 @@ public ScriptHost Instance
             // host level configuration files change
             do
             {
+                ScriptHost newInstance = null;
                 try
                 {
                     IsRunning = false;
@@ -66,7 +74,7 @@ public ScriptHost Instance
                     {
                         HostId = _config.HostConfig.HostId
                     };
-                    ScriptHost newInstance = ScriptHost.Create(_config);
+                    newInstance = _scriptHostFactory.Create(_config);
                     _traceWriter = newInstance.TraceWriter;
 
                     // write any function initialization errors to the log file
@@ -108,6 +116,20 @@ public ScriptHost Instance
                         _traceWriter.Error("A ScriptHost error occurred", ex);
                     }
 
+                    // If a ScriptHost instance was created before the exception was thrown
+                    // Orphan and cleanup that instance.
+                    if (newInstance != null)
+                    {
+                        Orphan(newInstance, forceStop: true)
+                            .ContinueWith(t =>
+                            {
+                                if (t.IsFaulted)
+                                {
+                                    t.Exception.Handle(e => true);
+                                }
+                            });
+                    }
+
                     // Wait for a short period of time before restarting to
                     // avoid cases where a host level config error might cause
                     // a rapid restart cycle
@@ -132,21 +154,33 @@ private static void LogErrors(ScriptHost host)
             }
         }
 
-        // Let the existing host instance finish currently executing functions.
-        private async Task Orphan(ScriptHost instance)
+        /// <summary>
+        /// Remove the <see cref="ScriptHost"/> instance from the live instances collection,
+        /// allowing it to finish currently executing functions before stopping and disposing of it.
+        /// </summary>
+        /// <param name="instance">The <see cref="ScriptHost"/> instance to remove</param>
+        /// <param name="forceStop">Forces the call to stop and dispose of the instance, even if it isn't present in the live instances collection.</param>
+        /// <returns></returns>
+        private async Task Orphan(ScriptHost instance, bool forceStop = false)
         {
             lock (_liveInstances)
             {
                 bool removed = _liveInstances.Remove(instance);
-                if (!removed)
+                if (!forceStop && !removed)
                 {
                     return; // somebody else is handling it
                 }
             }
 
-            // this thread now owns the instance
-            await instance.StopAsync();
-            instance.Dispose();
+            try
+            {
+                // this thread now owns the instance
+                await instance.StopAsync();
+            }
+            finally
+            {
+                instance.Dispose();
+            }
         }
 
         public void Stop()

src/WebJobs.Script/IScriptHostFactory.cs

@@ -0,0 +1,10 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace Microsoft.Azure.WebJobs.Script
+{
+    public interface IScriptHostFactory
+    {
+        ScriptHost Create(ScriptHostConfiguration config);
+    }
+}

src/WebJobs.Script/ScriptHostFactory.cs

@@ -0,0 +1,13 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace Microsoft.Azure.WebJobs.Script
+{
+    public sealed class ScriptHostFactory : IScriptHostFactory
+    {
+        public ScriptHost Create(ScriptHostConfiguration config)
+        {
+            return ScriptHost.Create(config);
+        }
+    }
+}

src/WebJobs.Script/WebJobs.Script.csproj

@@ -254,10 +254,12 @@
     <Compile Include="Extensions\ActionExtensions.cs" />
     <Compile Include="FileTraceWriter.cs" />
     <Compile Include="HttpMethodJsonConverter.cs" />
+    <Compile Include="IScriptHostFactory.cs" />
     <Compile Include="NullTraceWriter.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
     <Compile Include="Config\ScriptHostConfiguration.cs" />
     <Compile Include="Config\TypeLocator.cs" />
+    <Compile Include="ScriptHostFactory.cs" />
   </ItemGroup>
   <ItemGroup>
     <None Include="app.config">

test/WebJobs.Script.Tests/ScriptHostManagerTests.cs

@@ -7,6 +7,8 @@
 using System.Threading.Tasks;
 using Microsoft.Azure.WebJobs.Script;
 using Microsoft.WindowsAzure.Storage.Blob;
+using Moq;
+using Moq.Protected;
 using Xunit;
 
 namespace WebJobs.Script.Tests
@@ -59,6 +61,32 @@ public async Task UpdateFileAndRestart()
             }
         }
 
+        [Fact]
+        public void RunAndBlock_DisposesOfHost_WhenExceptionIsThrown()
+        {
+            ScriptHostConfiguration config = new ScriptHostConfiguration()
+            {
+                RootScriptPath = Environment.CurrentDirectory,
+                TraceWriter = NullTraceWriter.Instance
+            };
+
+            var hostMock = new Mock<TestScriptHost>(config);
+            var factoryMock = new Mock<IScriptHostFactory>();
+            factoryMock.Setup(f => f.Create(It.IsAny<ScriptHostConfiguration>()))
+                .Returns(hostMock.Object);
+
+            var target = new Mock<ScriptHostManager>(config, factoryMock.Object);
+            target.Protected().Setup("OnHostStarted")
+                .Throws(new Exception());
+
+            hostMock.Protected().Setup("Dispose", true)
+                .Callback(() => target.Object.Stop());
+
+            Task.Run(() => target.Object.RunAndBlock()).Wait(5000);
+
+            hostMock.Protected().Verify("Dispose", Times.Once(), true);
+        }
+
         // Update the manifest for the timer function
         // - this will cause a file touch which cause ScriptHostManager to notice and update
         // - set to a new output location so that we can ensure we're getting new changes. 
@@ -75,5 +103,12 @@ private static CloudBlockBlob UpdateOutputName(string prev, string hint, EndToEn
             blob.DeleteIfExists();
             return blob;
         }
+
+        public class TestScriptHost : ScriptHost
+        {
+            public TestScriptHost(ScriptHostConfiguration scriptConfig) : base(scriptConfig)
+            {
+            }
+        }
     }
 }