Only dispose the stream pipe wrapper and more tests

dotnet · Feb 8, 2019 · 67023fa · 67023fa
1 parent 1a7b2c2
commit 67023fa
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 12 deletions.
diff --git a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.FeatureCollection.cs b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.FeatureCollection.cs
@@ -205,8 +205,7 @@ PipeWriter IResponseBodyPipeFeature.ResponseBodyPipe
                 {
                     var responseBody = new WriteOnlyPipeStream(ResponsePipeWriter);
                     ResponseBody = responseBody;
-
-                    RegisterWrapperForDisposal(responseBody);
+                    _cachedResponsePipeWriter = ResponsePipeWriter;
                 }
             }
         }
@@ -225,18 +224,15 @@ Stream IHttpResponseFeature.Body
                     var responsePipeWriter = new StreamPipeWriter(ResponseBody, minimumSegmentSize: 4096, _context.MemoryPool);
                     ResponsePipeWriter = responsePipeWriter;
                     _cachedResponseBodyStream = ResponseBody;
-                    RegisterWrapperForDisposal(responsePipeWriter);
-                }
-            }
-        }
 
-        private void RegisterWrapperForDisposal(IDisposable responseBody)
-        {
-            if (_wrapperObjectsToDispose == null)
-            {
-                _wrapperObjectsToDispose = new List<IDisposable>();
+                    // The StreamPipeWrapper needs to be disposed as it hold onto blocks of memory
+                    if (_wrapperObjectsToDispose == null)
+                    {
+                        _wrapperObjectsToDispose = new List<IDisposable>();
+                    }
+                    _wrapperObjectsToDispose.Add(responsePipeWriter);
+                }
             }
-            _wrapperObjectsToDispose.Add(responseBody);
         }
 
         protected void ResetHttp1Features()

diff --git a/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs b/src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs
@@ -3517,6 +3517,74 @@ await connection.Receive(
             }
         }
 
+        [Fact]
+        public async Task ResponseSetBodyToSameValueTwiceGetPipeMultipleTimesSameObject()
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                var memoryStream = new MemoryStream();
+                httpContext.Response.Body = memoryStream;
+                var bodyPipe1 = httpContext.Response.BodyPipe;
+
+                httpContext.Response.Body = memoryStream;
+                var bodyPipe2 = httpContext.Response.BodyPipe;
+
+                Assert.Equal(bodyPipe1, bodyPipe2);
+                await Task.CompletedTask;
+            }, new TestServiceContext(LoggerFactory)))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "Host:",
+                        "",
+                        "");
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "Content-Length: 0",
+                        "",
+                        "");
+                }
+                await server.StopAsync();
+            }
+        }
+
+        [Fact]
+        public async Task ResponseSetPipeToSameValueTwiceGetBodyMultipleTimesSameObject()
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                var pipeWriter = new Pipe().Writer;
+                httpContext.Response.BodyPipe = pipeWriter;
+                var body1 = httpContext.Response.Body;
+
+                httpContext.Response.BodyPipe = pipeWriter;
+                var body2 = httpContext.Response.Body;
+
+                Assert.Equal(body1, body2);
+                await Task.CompletedTask;
+            }, new TestServiceContext(LoggerFactory)))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "Host:",
+                        "",
+                        "");
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "Content-Length: 0",
+                        "",
+                        "");
+                }
+                await server.StopAsync();
+            }
+        }
+
         [Fact]
         public async Task ResponseSetPipeAndBodyPipeIsWrapped()
         {