Release SSE response stream reference when GET request ends (#1519)
This commit is contained in:
parent
df0c102d75
commit
acb4cc9fd1
1 changed files with 141 additions and 29 deletions
|
|
@ -44,7 +44,7 @@ public sealed partial class StreamableHttpServerTransport : ITransport
|
|||
private TaskCompletionSource<bool>? _httpResponseTcs;
|
||||
private string? _negotiatedProtocolVersion;
|
||||
private bool _getHttpRequestStarted;
|
||||
private bool _getHttpResponseCompleted;
|
||||
private bool _disposed;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="StreamableHttpServerTransport"/> class.
|
||||
|
|
@ -137,33 +137,53 @@ public sealed partial class StreamableHttpServerTransport : ITransport
|
|||
throw new InvalidOperationException("GET requests are not supported in stateless mode.");
|
||||
}
|
||||
|
||||
using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
|
||||
try
|
||||
{
|
||||
if (_getHttpRequestStarted)
|
||||
using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
|
||||
if (_getHttpRequestStarted)
|
||||
{
|
||||
throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
|
||||
}
|
||||
|
||||
_getHttpRequestStarted = true;
|
||||
_httpSseWriter = new SseEventWriter(sseResponseStream);
|
||||
_httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
_storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
|
||||
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
else
|
||||
{
|
||||
// If there's no priming write, flush the stream to ensure HTTP response headers are
|
||||
// sent to the client now that the transport is ready to accept messages via SendMessageAsync.
|
||||
await sseResponseStream.FlushAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
_getHttpRequestStarted = true;
|
||||
_httpSseWriter = new SseEventWriter(sseResponseStream);
|
||||
_httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
_storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
|
||||
if (_storeSseWriter is not null)
|
||||
// Wait for the response to be written before returning from the handler.
|
||||
// This keeps the HTTP response open until the final response message is sent.
|
||||
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
// Release the SseEventWriter's reference to the response stream promptly when the GET
|
||||
// request ends, regardless of how it exits. Otherwise the response stream (and the
|
||||
// underlying Kestrel connection and associated memory pool buffers) remains pinned
|
||||
// in memory until the session itself is disposed (via explicit DELETE or idle timeout).
|
||||
// Clients that disconnect without sending DELETE — common with long-lived SSE — would
|
||||
// otherwise accumulate significant unmanaged memory per session during that interval.
|
||||
using (await _unsolicitedMessageLock.LockAsync(CancellationToken.None).ConfigureAwait(false))
|
||||
{
|
||||
var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
|
||||
await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
else
|
||||
{
|
||||
// If there's no priming write, flush the stream to ensure HTTP response headers are
|
||||
// sent to the client now that the transport is ready to accept messages via SendMessageAsync.
|
||||
await sseResponseStream.FlushAsync(cancellationToken).ConfigureAwait(false);
|
||||
if (_httpSseWriter is { } writer)
|
||||
{
|
||||
_httpSseWriter = null;
|
||||
writer.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for the response to be written before returning from the handler.
|
||||
// This keeps the HTTP response open until the final response message is sent.
|
||||
await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
|
@ -219,23 +239,22 @@ public sealed partial class StreamableHttpServerTransport : ITransport
|
|||
return;
|
||||
}
|
||||
|
||||
Debug.Assert(_httpSseWriter is not null);
|
||||
Debug.Assert(_httpResponseTcs is not null);
|
||||
|
||||
var item = SseItem.Message(message);
|
||||
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
// Always record the message in the event store (if configured) — even when the GET
|
||||
// response stream is gone — so a reconnecting client can replay it via Last-Event-ID.
|
||||
item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
if (!_getHttpResponseCompleted)
|
||||
if (_httpSseWriter is { } writer)
|
||||
{
|
||||
// Only write the message to the response if the response has not completed.
|
||||
|
||||
try
|
||||
{
|
||||
await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
await writer.WriteAsync(item, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
|
||||
{
|
||||
|
|
@ -249,12 +268,12 @@ public sealed partial class StreamableHttpServerTransport : ITransport
|
|||
{
|
||||
using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false);
|
||||
|
||||
if (_getHttpResponseCompleted)
|
||||
if (_disposed)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
_getHttpResponseCompleted = true;
|
||||
_disposed = true;
|
||||
|
||||
try
|
||||
{
|
||||
|
|
@ -266,7 +285,11 @@ public sealed partial class StreamableHttpServerTransport : ITransport
|
|||
try
|
||||
{
|
||||
_httpResponseTcs?.TrySetResult(true);
|
||||
_httpSseWriter?.Dispose();
|
||||
if (_httpSseWriter is { } writer)
|
||||
{
|
||||
_httpSseWriter = null;
|
||||
writer.Dispose();
|
||||
}
|
||||
|
||||
if (_storeSseWriter is not null)
|
||||
{
|
||||
|
|
|
|||
Loading…
Reference in a new issue