CSHARP-2884: Async Socket methods do not use Socket.ReceiveTimeout.

Author: rstam

src/MongoDB.Driver.Core/Core/Connections/BinaryConnection.cs

@@ -391,14 +391,15 @@ private async Task<IByteBuffer> ReceiveBufferAsync()
             try
             {
                 var messageSizeBytes = new byte[4];
-                await _stream.ReadBytesAsync(messageSizeBytes, 0, 4, _backgroundTaskCancellationToken).ConfigureAwait(false);
+                var readTimeout = _stream.CanTimeout ? TimeSpan.FromMilliseconds(_stream.ReadTimeout) : Timeout.InfiniteTimeSpan;
+                await _stream.ReadBytesAsync(messageSizeBytes, 0, 4, readTimeout, _backgroundTaskCancellationToken).ConfigureAwait(false);
                 var messageSize = BitConverter.ToInt32(messageSizeBytes, 0);
                 EnsureMessageSizeIsValid(messageSize);
                 var inputBufferChunkSource = new InputBufferChunkSource(BsonChunkPool.Default);
                 var buffer = ByteBufferFactory.Create(inputBufferChunkSource, messageSize);
                 buffer.Length = messageSize;
                 buffer.SetBytes(0, messageSizeBytes, 0, 4);
-                await _stream.ReadBytesAsync(buffer, 4, messageSize - 4, _backgroundTaskCancellationToken).ConfigureAwait(false);
+                await _stream.ReadBytesAsync(buffer, 4, messageSize - 4, readTimeout, _backgroundTaskCancellationToken).ConfigureAwait(false);
                 _lastUsedAtUtc = DateTime.UtcNow;
                 buffer.MakeReadOnly();
                 return buffer;
@@ -544,7 +545,8 @@ private async Task SendBufferAsync(IByteBuffer buffer, CancellationToken cancell
                 try
                 {
                     // don't use the caller's cancellationToken because once we start writing a message we have to write the whole thing
-                    await _stream.WriteBytesAsync(buffer, 0, buffer.Length, _backgroundTaskCancellationToken).ConfigureAwait(false);
+                    var writeTimeout = _stream.CanTimeout ? TimeSpan.FromMilliseconds(_stream.WriteTimeout) : Timeout.InfiniteTimeSpan;
+                    await _stream.WriteBytesAsync(buffer, 0, buffer.Length, writeTimeout, _backgroundTaskCancellationToken).ConfigureAwait(false);
                     _lastUsedAtUtc = DateTime.UtcNow;
                 }
                 catch (Exception ex)

src/MongoDB.Driver.Core/Core/Misc/StreamExtensionMethods.cs

@@ -36,6 +36,40 @@ public static void EfficientCopyTo(this Stream input, Stream output)
             }
         }
 
+        public static async Task<int> ReadAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
+        {
+            var state = 1; // 1 == reading, 2 == done reading, 3 == timedout, 4 == cancelled
+
+            var bytesRead = 0;
+            using (new Timer(_ => ChangeState(3), null, timeout, Timeout.InfiniteTimeSpan))
+            using (cancellationToken.Register(() => ChangeState(4)))
+            {
+                try
+                {
+                    bytesRead = await stream.ReadAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
+                    ChangeState(2);
+                }
+                catch when (state >= 3)
+                {
+                    // a different exception will be thrown instead below
+                }
+
+                if (state == 3) { throw new TimeoutException(); }
+                if (state == 4) { throw new OperationCanceledException(); }
+            }
+
+            return bytesRead;
+
+            void ChangeState(int to)
+            {
+                var from = Interlocked.CompareExchange(ref state, to, 1);
+                if (from == 1 && to >= 3)
+                {
+                    try { stream.Dispose(); } catch { } // disposing the stream aborts the read attempt
+                }
+            }
+        }
+
         public static void ReadBytes(this Stream stream, byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
@@ -76,7 +110,7 @@ public static void ReadBytes(this Stream stream, IByteBuffer buffer, int offset,
             }
         }
 
-        public static async Task ReadBytesAsync(this Stream stream, byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public static async Task ReadBytesAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -85,7 +119,7 @@ public static async Task ReadBytesAsync(this Stream stream, byte[] buffer, int o
 
             while (count > 0)
             {
-                var bytesRead = await stream.ReadAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
+                var bytesRead = await stream.ReadAsync(buffer, offset, count, timeout, cancellationToken).ConfigureAwait(false);
                 if (bytesRead == 0)
                 {
                     throw new EndOfStreamException();
@@ -95,7 +129,7 @@ public static async Task ReadBytesAsync(this Stream stream, byte[] buffer, int o
             }
         }
 
-        public static async Task ReadBytesAsync(this Stream stream, IByteBuffer buffer, int offset, int count, CancellationToken cancellationToken)
+        public static async Task ReadBytesAsync(this Stream stream, IByteBuffer buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -106,7 +140,7 @@ public static async Task ReadBytesAsync(this Stream stream, IByteBuffer buffer,
             {
                 var backingBytes = buffer.AccessBackingBytes(offset);
                 var bytesToRead = Math.Min(count, backingBytes.Count);
-                var bytesRead = await stream.ReadAsync(backingBytes.Array, backingBytes.Offset, bytesToRead, cancellationToken).ConfigureAwait(false);
+                var bytesRead = await stream.ReadAsync(backingBytes.Array, backingBytes.Offset, bytesToRead, timeout, cancellationToken).ConfigureAwait(false);
                 if (bytesRead == 0)
                 {
                     throw new EndOfStreamException();
@@ -116,6 +150,38 @@ public static async Task ReadBytesAsync(this Stream stream, IByteBuffer buffer,
             }
         }
 
+
+        public static async Task WriteAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
+        {
+            var state = 1; // 1 == writing, 2 == done writing, 3 == timedout, 4 == cancelled
+
+            using (new Timer(_ => ChangeState(3), null, timeout, Timeout.InfiniteTimeSpan))
+            using (cancellationToken.Register(() => ChangeState(4)))
+            {
+                try
+                {
+                    await stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
+                    ChangeState(2);
+                }
+                catch when (state >= 3)
+                {
+                    // a different exception will be thrown instead below
+                }
+
+                if (state == 3) { throw new TimeoutException(); }
+                if (state == 4) { throw new OperationCanceledException(); }
+            }
+
+            void ChangeState(int to)
+            {
+                var from = Interlocked.CompareExchange(ref state, to, 1);
+                if (from == 1 && to >= 3)
+                {
+                    try { stream.Dispose(); } catch { } // disposing the stream aborts the write attempt
+                }
+            }
+        }
+
         public static void WriteBytes(this Stream stream, IByteBuffer buffer, int offset, int count, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
@@ -134,7 +200,7 @@ public static void WriteBytes(this Stream stream, IByteBuffer buffer, int offset
             }
         }
 
-        public static async Task WriteBytesAsync(this Stream stream, IByteBuffer buffer, int offset, int count, CancellationToken cancellationToken)
+        public static async Task WriteBytesAsync(this Stream stream, IByteBuffer buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -145,7 +211,7 @@ public static async Task WriteBytesAsync(this Stream stream, IByteBuffer buffer,
             {
                 var backingBytes = buffer.AccessBackingBytes(offset);
                 var bytesToWrite = Math.Min(count, backingBytes.Count);
-                await stream.WriteAsync(backingBytes.Array, backingBytes.Offset, bytesToWrite, cancellationToken).ConfigureAwait(false);
+                await stream.WriteAsync(backingBytes.Array, backingBytes.Offset, bytesToWrite, timeout, cancellationToken).ConfigureAwait(false);
                 offset += bytesToWrite;
                 count -= bytesToWrite;
             }