Fix flakiness in Kestrel ungraceful shutdown test (#6637)

src/Servers/Kestrel/shared/test/TestApplicationErrorLogger.cs

@@ -1,10 +1,11 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
@@ -16,6 +17,9 @@ namespace Microsoft.AspNetCore.Testing
         // Application errors are logged using 13 as the eventId.
         private const int ApplicationErrorEventId = 13;
 
+        private Func<LogMessage, bool> _messageFilter;
+        private TaskCompletionSource<LogMessage> _messageFilterTcs;
+
         public List<Type> IgnoredExceptions { get; } = new List<Type>();
 
         public bool ThrowOnCriticalErrors { get; set; } = true;
@@ -32,6 +36,19 @@ namespace Microsoft.AspNetCore.Testing
 
         public int ApplicationErrorsLogged => Messages.Count(message => message.EventId.Id == ApplicationErrorEventId);
 
+        public Task<LogMessage> WaitForMessage(Func<LogMessage, bool> messageFilter)
+        {
+            if (_messageFilterTcs != null)
+            {
+                throw new InvalidOperationException($"{nameof(WaitForMessage)} cannot be called concurrently.");
+            }
+
+            _messageFilterTcs = new TaskCompletionSource<LogMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _messageFilter = messageFilter;
+
+            return _messageFilterTcs.Task;
+        }
+
         public IDisposable BeginScope<TState>(TState state)
         {
             Scopes.Enqueue(state);
@@ -77,13 +94,21 @@ namespace Microsoft.AspNetCore.Testing
                 throw new Exception($"Unexpected connection error. {log}");
             }
 
-            Messages.Enqueue(new LogMessage
+            var logMessage = new LogMessage
             {
                 LogLevel = logLevel,
                 EventId = eventId,
                 Exception = exception,
                 Message = formatter(state, exception)
-            });
+            };
+
+            Messages.Enqueue(logMessage);
+
+            if (_messageFilter?.Invoke(logMessage) == true)
+            {
+                _messageFilterTcs.TrySetResult(logMessage);
+                _messageFilterTcs = null;
+            }
         }
 
         public class LogMessage

src/Servers/Kestrel/test/FunctionalTests/Http2/ShutdownTests.cs

@@ -134,7 +134,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests.Http2
                 Assert.False(requestTask.IsCompleted);
                 await requestStarted.Task.DefaultTimeout();
 
-                await server.StopAsync(new CancellationToken(true)).DefaultTimeout();
+                // Wait for the graceful shutdown log before canceling the token passed to StopAsync and triggering an ungraceful shutdown.
+                // Otherwise, graceful shutdown might be skipped causing there to be no corresponding log. https://github.com/aspnet/AspNetCore/issues/6556
+                var closingMessageTask = TestApplicationErrorLogger.WaitForMessage(m => m.Message.Contains("is closing.")).DefaultTimeout();
+
+                var cts = new CancellationTokenSource();
+                var stopServerTask = server.StopAsync(cts.Token).DefaultTimeout();
+
+                await closingMessageTask;
+                cts.Cancel();
+                await stopServerTask;
             }
 
             Assert.Contains(TestApplicationErrorLogger.Messages, m => m.Message.Contains("is closing."));