Fix race in TransportBase.SetConnected (#253)

* Fix race in TransportBase.SetConnected

- Previously the ChannelWriter could already be completed if multiple
  threads raced in SetConnected(false)

* Fix WriteMessageAsync IsConnected test
This commit is contained in:
Stephen Halter 2025-04-09 11:10:51 -07:00 • committed by GitHub
commit 08f9cdbf7d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -13,7 +13,7 @@ public abstract class TransportBase : ITransport
{ {
private readonly Channel<IJsonRpcMessage> _messageChannel; private readonly Channel<IJsonRpcMessage> _messageChannel;
private readonly ILogger _logger; private readonly ILogger _logger;
private bool _isConnected; private int _isConnected;
/// <summary> /// <summary>
/// Initializes a new instance of the <see cref="TransportBase"/> class. /// Initializes a new instance of the <see cref="TransportBase"/> class.
@ -30,7 +30,7 @@ public abstract class TransportBase : ITransport
} }
/// <inheritdoc/> /// <inheritdoc/>
public bool IsConnected => _isConnected; public bool IsConnected => _isConnected == 1;
/// <inheritdoc/> /// <inheritdoc/>
public ChannelReader<IJsonRpcMessage> MessageReader => _messageChannel.Reader; public ChannelReader<IJsonRpcMessage> MessageReader => _messageChannel.Reader;
@ -48,7 +48,7 @@ public abstract class TransportBase : ITransport
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param> /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
protected async Task WriteMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default) protected async Task WriteMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{ {
if (!_isConnected) if (!IsConnected)
{ {
throw new McpTransportException("Transport is not connected"); throw new McpTransportException("Transport is not connected");
} }
@ -64,12 +64,12 @@ public abstract class TransportBase : ITransport
/// <param name="isConnected">Whether the transport is connected.</param> /// <param name="isConnected">Whether the transport is connected.</param>
protected void SetConnected(bool isConnected) protected void SetConnected(bool isConnected)
{ {
if (_isConnected == isConnected) var newIsConnected = isConnected ? 1 : 0;
if (Interlocked.Exchange(ref _isConnected, newIsConnected) == newIsConnected)
{ {
return; return;
} }
_isConnected = isConnected;
if (!isConnected) if (!isConnected)
{ {
_messageChannel.Writer.Complete(); _messageChannel.Writer.Complete();