Dispose of IDisposable context

Needs aspnet/Hosting#355
benaadams · Sep 15, 2015 · 14c6c85 · 14c6c85
1 parent 3a0b601
commit 14c6c85
Show file tree

Hide file tree

Showing 10 changed files with 25 additions and 14 deletions.
diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs
@@ -221,9 +221,10 @@ private void FireOnCompleted()
         private async Task ExecuteAsync()
         {
             Exception error = null;
+            IDisposable disposable = null;
             try
             {
-                await Application.Invoke(this).ConfigureAwait(false);
+                disposable = await Application.Invoke(this).ConfigureAwait(false);
 
                 // Trigger FireOnStarting if ProduceStart hasn't been called yet.
                 // We call it here, so it can go through our normal error handling
@@ -241,6 +242,7 @@ private async Task ExecuteAsync()
             {
                 FireOnCompleted();
                 ProduceEnd(error);
+                disposable?.Dispose();
             }
         }
 

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/Listener.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/Listener.cs
@@ -24,7 +24,7 @@ public Task StartAsync(
             string host,
             int port,
             KestrelThread thread,
-            Func<Frame, Task> application)
+            Func<Frame, Task<IDisposable>> application)
         {
             Thread = thread;
             Application = application;

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerContext.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerContext.cs
@@ -27,7 +27,7 @@ public ListenerContext(ListenerContext listenerContext)
 
         public KestrelThread Thread { get; set; }
 
-        public Func<Frame, Task> Application { get; set; }
+        public Func<Frame, Task<IDisposable>> Application { get; set; }
 
         public IMemoryPool Memory { get; set; }
 

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerPrimary.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerPrimary.cs
@@ -34,7 +34,7 @@ public async Task StartAsync(
             string host,
             int port,
             KestrelThread thread,
-            Func<Frame, Task> application)
+            Func<Frame, Task<IDisposable>> application)
         {
             await StartAsync(scheme, host, port, thread, application).ConfigureAwait(false);
 

diff --git a/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerSecondary.cs b/src/Microsoft.AspNet.Server.Kestrel/Http/ListenerSecondary.cs
@@ -25,7 +25,7 @@ protected ListenerSecondary(ServiceContext serviceContext) : base(serviceContext
         public Task StartAsync(
             string pipeName,
             KestrelThread thread,
-            Func<Frame, Task> application)
+            Func<Frame, Task<IDisposable>> application)
         {
             Thread = thread;
             Application = application;

diff --git a/src/Microsoft.AspNet.Server.Kestrel/KestrelEngine.cs b/src/Microsoft.AspNet.Server.Kestrel/KestrelEngine.cs
@@ -106,7 +106,7 @@ public void Dispose()
             Threads.Clear();
         }
 
-        public IDisposable CreateServer(string scheme, string host, int port, Func<Frame, Task> application)
+        public IDisposable CreateServer(string scheme, string host, int port, Func<Frame, Task<IDisposable>> application)
         {
             var listeners = new List<IDisposable>();
             var usingPipes = host.StartsWith(Constants.UnixPipeHostPrefix);

diff --git a/src/Microsoft.AspNet.Server.Kestrel/ServerFactory.cs b/src/Microsoft.AspNet.Server.Kestrel/ServerFactory.cs
@@ -37,7 +37,7 @@ public IFeatureCollection Initialize(IConfiguration configuration)
             return serverFeatures;
         }
 
-        public IDisposable Start(IFeatureCollection serverFeatures, Func<IFeatureCollection, Task> application)
+        public IDisposable Start(IFeatureCollection serverFeatures, Func<IFeatureCollection, Task<IDisposable>> application)
         {
             var disposables = new Stack<IDisposable>();
             var disposer = new Disposable(() =>

diff --git a/test/Microsoft.AspNet.Server.KestrelTests/ChunkedResponseTests.cs b/test/Microsoft.AspNet.Server.KestrelTests/ChunkedResponseTests.cs
@@ -18,6 +18,7 @@ public async Task ResponsesAreChunkedAutomatically()
                 frame.ResponseHeaders.Clear();
                 await frame.ResponseBody.WriteAsync(Encoding.ASCII.GetBytes("Hello "), 0, 6);
                 await frame.ResponseBody.WriteAsync(Encoding.ASCII.GetBytes("World!"), 0, 6);
+                return null;
             }))
             {
                 using (var connection = new TestConnection())
@@ -50,6 +51,7 @@ public async Task ZeroLengthWritesAreIgnored()
                 await frame.ResponseBody.WriteAsync(Encoding.ASCII.GetBytes("Hello "), 0, 6);
                 await frame.ResponseBody.WriteAsync(new byte[0], 0, 0);
                 await frame.ResponseBody.WriteAsync(Encoding.ASCII.GetBytes("World!"), 0, 6);
+                return null;
             }))
             {
                 using (var connection = new TestConnection())
@@ -80,6 +82,7 @@ public async Task EmptyResponseBodyHandledCorrectlyWithZeroLengthWrite()
             {
                 frame.ResponseHeaders.Clear();
                 await frame.ResponseBody.WriteAsync(new byte[0], 0, 0);
+                return null;
             }))
             {
                 using (var connection = new TestConnection())

diff --git a/test/Microsoft.AspNet.Server.KestrelTests/EngineTests.cs b/test/Microsoft.AspNet.Server.KestrelTests/EngineTests.cs
@@ -20,7 +20,7 @@ namespace Microsoft.AspNet.Server.KestrelTests
     /// </summary>
     public class EngineTests
     {
-        private async Task App(Frame frame)
+        private async Task<IDisposable> App(Frame frame)
         {
             frame.ResponseHeaders.Clear();
             for (; ;)
@@ -33,6 +33,7 @@ private async Task App(Frame frame)
                 }
                 await frame.ResponseBody.WriteAsync(buffer, 0, count);
             }
+            return null;
         }
 
         ILibraryManager LibraryManager
@@ -58,7 +59,7 @@ ILibraryManager LibraryManager
             }
         }
 
-        private async Task AppChunked(Frame frame)
+        private async Task<IDisposable> AppChunked(Frame frame)
         {
             frame.ResponseHeaders.Clear();
             var data = new MemoryStream();
@@ -75,12 +76,13 @@ private async Task AppChunked(Frame frame)
             var bytes = data.ToArray();
             frame.ResponseHeaders["Content-Length"] = new[] { bytes.Length.ToString() };
             await frame.ResponseBody.WriteAsync(bytes, 0, bytes.Length);
+            return null;
         }
 
-        private Task EmptyApp(Frame frame)
+        private Task<IDisposable> EmptyApp(Frame frame)
         {
             frame.ResponseHeaders.Clear();
-            return Task.FromResult<object>(null);
+            return Task.FromResult<IDisposable>(null);
         }
 
         [Fact]
@@ -472,6 +474,8 @@ public async Task ZeroContentLengthNotSetAutomaticallyForCertainStatusCodes()
                     var statusString = await reader.ReadLineAsync();
                     frame.StatusCode = int.Parse(statusString);
                 }
+
+                return null;
             }))
             {
                 using (var connection = new TestConnection())
@@ -715,7 +719,7 @@ public async Task ThrowingInOnStartingResultsIn500Response()
 
                 // If we write to the response stream, we will not get a 500.
 
-                return Task.FromResult<object>(null);
+                return Task.FromResult<IDisposable>(null);
             }))
             {
                 using (var connection = new TestConnection())
@@ -770,6 +774,8 @@ public async Task ThrowingInOnStartingResultsInFailedWrites()
 
                 // The second write should succeed since the OnStarting callback will not be called again
                 await frame.ResponseBody.WriteAsync(Encoding.ASCII.GetBytes("Exception!!"), 0, 11);
+
+                return null;
             }))
             {
                 using (var connection = new TestConnection())

diff --git a/test/Microsoft.AspNet.Server.KestrelTests/TestServer.cs b/test/Microsoft.AspNet.Server.KestrelTests/TestServer.cs
@@ -18,7 +18,7 @@ public class TestServer : IDisposable
         private KestrelEngine _engine;
         private IDisposable _server;
 
-        public TestServer(Func<Frame, Task> app)
+        public TestServer(Func<Frame, Task<IDisposable>> app)
         {
             Create(app);
         }
@@ -43,7 +43,7 @@ ILibraryManager LibraryManager
             }
         }
 
-        public void Create(Func<Frame, Task> app)
+        public void Create(Func<Frame, Task<IDisposable>> app)
         {
             _engine = new KestrelEngine(LibraryManager, new ShutdownNotImplemented(), new TestLogger());
             _engine.Start(1);