diff --git a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2ConnectionTests.cs b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2ConnectionTests.cs index 7733bd8a591b..c08edda0ea94 100644 --- a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2ConnectionTests.cs +++ b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2ConnectionTests.cs @@ -404,9 +404,11 @@ public async Task StreamPool_SingleStream_ReturnedToPool() { // Add stream to Http2Connection._completedStreams inline with SetResult(). var serverTcs = new TaskCompletionSource(); + var appTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); await InitializeConnectionAsync(async context => { + appTcs.TrySetResult(); await serverTcs.Task; await _echoApplication(context); }); @@ -415,6 +417,8 @@ await InitializeConnectionAsync(async context => await StartStreamAsync(1, _browserRequestHeaders, endStream: true); + // If app code is running we know the stream has been created + await appTcs.Task; var stream = _connection._streams[1]; serverTcs.SetResult(); @@ -490,9 +494,12 @@ public async Task StreamPool_MultipleStreamsInSequence_PooledStreamReused() TaskCompletionSource appDelegateTcs = null; object persistedState = null; var requestCount = 0; + var appTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); await InitializeConnectionAsync(async context => { + appTcs.TrySetResult(); + requestCount++; var persistentStateCollection = context.Features.Get().State; if (persistentStateCollection.TryGetValue("Counter", out var value)) @@ -509,6 +516,9 @@ await InitializeConnectionAsync(async context => appDelegateTcs = new TaskCompletionSource(); await StartStreamAsync(1, _browserRequestHeaders, endStream: true); + // If app code is running we know the stream has been created + await appTcs.Task; + appTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); // Get the in progress stream var stream = _connection._streams[1]; @@ -534,6 +544,8 @@ await ExpectAsync(Http2FrameType.HEADERS, appDelegateTcs = new TaskCompletionSource(); await StartStreamAsync(3, _browserRequestHeaders, endStream: true); + await appTcs.Task; + // New stream has been taken from the pool Assert.Equal(0, _connection.StreamPool.Count); @@ -564,9 +576,11 @@ public async Task StreamPool_StreamIsInvalidState_DontReturnedToPool() { // Add (or don't add) stream to Http2Connection._completedStreams inline with SetResult(). var serverTcs = new TaskCompletionSource(); + var appTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); await InitializeConnectionAsync(async context => { + appTcs.TrySetResult(); await serverTcs.Task.DefaultTimeout(); await context.Response.WriteAsync("Content"); @@ -575,6 +589,8 @@ await InitializeConnectionAsync(async context => await StartStreamAsync(1, _browserRequestHeaders, endStream: true); + // If app code is running we know the stream has been created + await appTcs.Task; var stream = _connection._streams[1]; serverTcs.SetResult();