Tigger caching header finalization on write #92

Author: JunTaoLuo

src/Microsoft.AspNetCore.ResponseCaching/Internal/Interfaces/IResponseCache.cs

@@ -8,7 +8,10 @@ namespace Microsoft.AspNetCore.ResponseCaching.Internal
 {
     public interface IResponseCache
     {
+        IResponseCacheEntry Get(string key);
         Task<IResponseCacheEntry> GetAsync(string key);
+
+        void Set(string key, IResponseCacheEntry entry, TimeSpan validFor);
         Task SetAsync(string key, IResponseCacheEntry entry, TimeSpan validFor);
     }
 }

src/Microsoft.AspNetCore.ResponseCaching/Internal/MemoryResponseCache.cs

@@ -4,7 +4,7 @@
 using System;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Caching.Memory;
-using Microsoft.Net.Http.Headers;
+using Microsoft.Extensions.Internal;
 
 namespace Microsoft.AspNetCore.ResponseCaching.Internal
 {
@@ -22,34 +22,39 @@ public MemoryResponseCache(IMemoryCache cache)
             _cache = cache;
         }
 
-        public Task<IResponseCacheEntry> GetAsync(string key)
+        public IResponseCacheEntry Get(string key)
         {
             var entry = _cache.Get(key);
 
             var memoryCachedResponse = entry as MemoryCachedResponse;
             if (memoryCachedResponse != null)
             {
-                return Task.FromResult<IResponseCacheEntry>(new CachedResponse
+                return new CachedResponse
                 {
                     Created = memoryCachedResponse.Created,
                     StatusCode = memoryCachedResponse.StatusCode,
                     Headers = memoryCachedResponse.Headers,
                     Body = new SegmentReadStream(memoryCachedResponse.BodySegments, memoryCachedResponse.BodyLength)
-                });
+                };
             }
             else
             {
-                return Task.FromResult(entry as IResponseCacheEntry);
+                return entry as IResponseCacheEntry;
             }
         }
 
-        public async Task SetAsync(string key, IResponseCacheEntry entry, TimeSpan validFor)
+        public Task<IResponseCacheEntry> GetAsync(string key)
+        {
+            return Task.FromResult(Get(key));
+        }
+
+        public void Set(string key, IResponseCacheEntry entry, TimeSpan validFor)
         {
             var cachedResponse = entry as CachedResponse;
             if (cachedResponse != null)
             {
                 var segmentStream = new SegmentWriteStream(StreamUtilities.BodySegmentSize);
-                await cachedResponse.Body.CopyToAsync(segmentStream);
+                cachedResponse.Body.CopyTo(segmentStream);
 
                 _cache.Set(
                     key,
@@ -77,5 +82,11 @@ public async Task SetAsync(string key, IResponseCacheEntry entry, TimeSpan valid
                     });
             }
         }
+
+        public Task SetAsync(string key, IResponseCacheEntry entry, TimeSpan validFor)
+        {
+            Set(key, entry, validFor);
+            return TaskCache.CompletedTask;
+        }
     }
 }

src/Microsoft.AspNetCore.ResponseCaching/Microsoft.AspNetCore.ResponseCaching.csproj

@@ -17,7 +17,7 @@
     <PackageReference Include="Microsoft.AspNetCore.Http" Version="1.2.0-*" />
     <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="1.2.0-*" />
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" Version="1.2.0-*" />
-    <PackageReference Include="Microsoft.Extensions.TaskCache.Sources" Version="1.2.0-*" PrivateAssets="All"/>
+    <PackageReference Include="Microsoft.Extensions.TaskCache.Sources" Version="1.2.0-*" PrivateAssets="All" />
   </ItemGroup>
 
 </Project>

src/Microsoft.AspNetCore.ResponseCaching/ResponseCachingMiddleware.cs

@@ -25,7 +25,6 @@ public class ResponseCachingMiddleware
         private readonly IResponseCachingPolicyProvider _policyProvider;
         private readonly IResponseCache _cache;
         private readonly IResponseCachingKeyProvider _keyProvider;
-        private readonly Func<object, Task> _onStartingCallback;
 
         public ResponseCachingMiddleware(
             RequestDelegate next,
@@ -66,7 +65,6 @@ public ResponseCachingMiddleware(
             _policyProvider = policyProvider;
             _cache = cache;
             _keyProvider = keyProvider;
-            _onStartingCallback = state => OnResponseStartingAsync((ResponseCachingContext)state);
         }
 
         public async Task Invoke(HttpContext httpContext)
@@ -90,13 +88,10 @@ public async Task Invoke(HttpContext httpContext)
 
                     try
                     {
-                        // Subscribe to OnStarting event
-                        httpContext.Response.OnStarting(_onStartingCallback, context);
-
                         await _next(httpContext);
 
                         // If there was no response body, check the response headers now. We can cache things like redirects.
-                        await OnResponseStartingAsync(context);
+                        await StartResponseAsync(context);
 
                         // Finalize the cache entry
                         await FinalizeCacheBodyAsync(context);
@@ -219,10 +214,17 @@ internal async Task<bool> TryServeFromCacheAsync(ResponseCachingContext context)
             return false;
         }
 
-        internal async Task FinalizeCacheHeadersAsync(ResponseCachingContext context)
+
+        /// <summary>
+        /// Finalize cache headers.
+        /// </summary>
+        /// <param name="context"></param>
+        /// <returns><c>true</c> if a vary by entry needs to be stored in the cache; otherwise <c>false</c>.</returns>
+        private bool OnFinalizeCacheHeaders(ResponseCachingContext context)
         {
             if (_policyProvider.IsResponseCacheable(context))
             {
+                var storeVaryByEntry = false;
                 context.ShouldCacheResponse = true;
 
                 // Create the cache entry now
@@ -262,7 +264,7 @@ internal async Task FinalizeCacheHeadersAsync(ResponseCachingContext context)
 
                     // Always overwrite the CachedVaryByRules to update the expiry information
                     _logger.LogVaryByRulesUpdated(normalizedVaryHeaders, normalizedVaryQueryKeys);
-                    await _cache.SetAsync(context.BaseKey, context.CachedVaryByRules, context.CachedResponseValidFor);
+                    storeVaryByEntry = true;
 
                     context.StorageVaryKey = _keyProvider.CreateStorageVaryByKey(context);
                 }
@@ -290,11 +292,29 @@ internal async Task FinalizeCacheHeadersAsync(ResponseCachingContext context)
                         context.CachedResponse.Headers[header.Key] = header.Value;
                     }
                 }
+
+                return storeVaryByEntry;
             }
-            else
+
+            context.ResponseCachingStream.DisableBuffering();
+            return false;
+        }
+
+        internal void FinalizeCacheHeaders(ResponseCachingContext context)
+        {
+            if (OnFinalizeCacheHeaders(context))
+            {
+                _cache.Set(context.BaseKey, context.CachedVaryByRules, context.CachedResponseValidFor);
+            }
+        }
+
+        internal Task FinalizeCacheHeadersAsync(ResponseCachingContext context)
+        {
+            if (OnFinalizeCacheHeaders(context))
             {
-                context.ResponseCachingStream.DisableBuffering();
+                return _cache.SetAsync(context.BaseKey, context.CachedVaryByRules, context.CachedResponseValidFor);
             }
+            return TaskCache.CompletedTask;
         }
 
         internal async Task FinalizeCacheBodyAsync(ResponseCachingContext context)
@@ -327,19 +347,38 @@ internal async Task FinalizeCacheBodyAsync(ResponseCachingContext context)
             }
         }
 
-        internal Task OnResponseStartingAsync(ResponseCachingContext context)
+        /// <summary>
+        /// Mark the response as started and set the response time if no reponse was started yet.
+        /// </summary>
+        /// <param name="context"></param>
+        /// <returns><c>true</c> if the response was not started before this call; otherwise <c>false</c>.</returns>
+        private bool OnStartResponse(ResponseCachingContext context)
         {
             if (!context.ResponseStarted)
             {
                 context.ResponseStarted = true;
                 context.ResponseTime = _options.SystemClock.UtcNow;
 
-                return FinalizeCacheHeadersAsync(context);
+                return true;
             }
-            else
+            return false;
+        }
+
+        internal void StartResponse(ResponseCachingContext context)
+        {
+            if (OnStartResponse(context))
             {
-                return TaskCache.CompletedTask;
+                FinalizeCacheHeaders(context);
+            }
+        }
+
+        internal Task StartResponseAsync(ResponseCachingContext context)
+        {
+            if (OnStartResponse(context))
+            {
+                return FinalizeCacheHeadersAsync(context);
             }
+            return TaskCache.CompletedTask;
         }
 
         internal static void AddResponseCachingFeature(HttpContext context)
@@ -355,7 +394,12 @@ internal void ShimResponseStream(ResponseCachingContext context)
         {
             // Shim response stream
             context.OriginalResponseStream = context.HttpContext.Response.Body;
-            context.ResponseCachingStream = new ResponseCachingStream(context.OriginalResponseStream, _options.MaximumBodySize, StreamUtilities.BodySegmentSize);
+            context.ResponseCachingStream = new ResponseCachingStream(
+                context.OriginalResponseStream,
+                _options.MaximumBodySize,
+                StreamUtilities.BodySegmentSize,
+                () => StartResponse(context),
+                () => StartResponseAsync(context));
             context.HttpContext.Response.Body = context.ResponseCachingStream;
 
             // Shim IHttpSendFileFeature

src/Microsoft.AspNetCore.ResponseCaching/Streams/ResponseCachingStream.cs

@@ -14,12 +14,16 @@ internal class ResponseCachingStream : Stream
         private readonly long _maxBufferSize;
         private readonly int _segmentSize;
         private SegmentWriteStream _segmentWriteStream;
+        private Action _startResponseCallback;
+        private Func<Task> _startResponseCallbackAsync;
 
-        internal ResponseCachingStream(Stream innerStream, long maxBufferSize, int segmentSize)
+        internal ResponseCachingStream(Stream innerStream, long maxBufferSize, int segmentSize, Action startResponseCallback, Func<Task> startResponseCallbackAsync)
         {
             _innerStream = innerStream;
             _maxBufferSize = maxBufferSize;
             _segmentSize = segmentSize;
+            _startResponseCallback = startResponseCallback;
+            _startResponseCallbackAsync = startResponseCallbackAsync;
             _segmentWriteStream = new SegmentWriteStream(_segmentSize);
         }
 
@@ -71,10 +75,32 @@ public override long Seek(long offset, SeekOrigin origin)
         }
 
         public override void Flush()
-            => _innerStream.Flush();
+        {
+            try
+            {
+                _startResponseCallback();
+                _innerStream.Flush();
+            }
+            catch
+            {
+                DisableBuffering();
+                throw;
+            }
+        }
 
-        public override Task FlushAsync(CancellationToken cancellationToken)
-            => _innerStream.FlushAsync();
+        public override async Task FlushAsync(CancellationToken cancellationToken)
+        {
+            try
+            {
+                await _startResponseCallbackAsync();
+                await _innerStream.FlushAsync();
+            }
+            catch
+            {
+                DisableBuffering();
+                throw;
+            }
+        }
 
         // Underlying stream is write-only, no need to override other read related methods
         public override int Read(byte[] buffer, int offset, int count)
@@ -84,6 +110,7 @@ public override void Write(byte[] buffer, int offset, int count)
         {
             try
             {
+                _startResponseCallback();
                 _innerStream.Write(buffer, offset, count);
             }
             catch
@@ -109,6 +136,7 @@ public override async Task WriteAsync(byte[] buffer, int offset, int count, Canc
         {
             try
             {
+                await _startResponseCallbackAsync();
                 await _innerStream.WriteAsync(buffer, offset, count, cancellationToken);
             }
             catch

test/Microsoft.AspNetCore.ResponseCaching.Tests/ResponseCachingMiddlewareTests.cs

@@ -359,7 +359,7 @@ public void ContentIsNotModified_IfNoneMatch_MatchesAtLeastOneValue_True()
         }
 
         [Fact]
-        public async Task OnResponseStartingAsync_IfAllowResponseCaptureIsTrue_SetsResponseTime()
+        public async Task StartResponsegAsync_IfAllowResponseCaptureIsTrue_SetsResponseTime()
         {
             var clock = new TestClock
             {
@@ -372,13 +372,13 @@ public async Task OnResponseStartingAsync_IfAllowResponseCaptureIsTrue_SetsRespo
             var context = TestUtils.CreateTestContext();
             context.ResponseTime = null;
 
-            await middleware.OnResponseStartingAsync(context);
+            await middleware.StartResponseAsync(context);
 
             Assert.Equal(clock.UtcNow, context.ResponseTime);
         }
 
         [Fact]
-        public async Task OnResponseStartingAsync_IfAllowResponseCaptureIsTrue_SetsResponseTimeOnlyOnce()
+        public async Task StartResponseAsync_IfAllowResponseCaptureIsTrue_SetsResponseTimeOnlyOnce()
         {
             var clock = new TestClock
             {
@@ -392,12 +392,12 @@ public async Task OnResponseStartingAsync_IfAllowResponseCaptureIsTrue_SetsRespo
             var initialTime = clock.UtcNow;
             context.ResponseTime = null;
 
-            await middleware.OnResponseStartingAsync(context);
+            await middleware.StartResponseAsync(context);
             Assert.Equal(initialTime, context.ResponseTime);
 
             clock.UtcNow += TimeSpan.FromSeconds(10);
 
-            await middleware.OnResponseStartingAsync(context);
+            await middleware.StartResponseAsync(context);
             Assert.NotEqual(clock.UtcNow, context.ResponseTime);
             Assert.Equal(initialTime, context.ResponseTime);
         }
@@ -790,10 +790,9 @@ public async Task FinalizeCacheBody_DoNotCache_IfShouldCacheResponseFalse()
             var middleware = TestUtils.CreateTestMiddleware(testSink: sink, cache: cache);
             var context = TestUtils.CreateTestContext();
 
-            context.ShouldCacheResponse = false;
             middleware.ShimResponseStream(context);
             await context.HttpContext.Response.WriteAsync(new string('0', 10));
-
+            context.ShouldCacheResponse = false;
 
             await middleware.FinalizeCacheBodyAsync(context);