CSHARP-3662: MongoClientSettings.SocketTimeout does not work for values under 500ms on Windows for sync code (#1690)

Author: sanych-sun

src/MongoDB.Driver/Core/Compression/SnappyCompressor.cs

@@ -34,7 +34,7 @@ public void Compress(Stream input, Stream output)
         {
             var uncompressedSize = (int)(input.Length - input.Position);
             var uncompressedBytes = new byte[uncompressedSize]; // does not include uncompressed message headers
-            input.ReadBytes(uncompressedBytes, offset: 0, count: uncompressedSize, CancellationToken.None);
+            input.ReadBytes(uncompressedBytes, offset: 0, count: uncompressedSize, Timeout.InfiniteTimeSpan, CancellationToken.None);
             var maxCompressedSize = Snappy.GetMaxCompressedLength(uncompressedSize);
             var compressedBytes = new byte[maxCompressedSize];
             var compressedSize = Snappy.Compress(uncompressedBytes, compressedBytes);
@@ -50,7 +50,7 @@ public void Decompress(Stream input, Stream output)
         {
             var compressedSize = (int)(input.Length - input.Position);
             var compressedBytes = new byte[compressedSize];
-            input.ReadBytes(compressedBytes, offset: 0, count: compressedSize, CancellationToken.None);
+            input.ReadBytes(compressedBytes, offset: 0, count: compressedSize, Timeout.InfiniteTimeSpan, CancellationToken.None);
             var uncompressedSize = Snappy.GetUncompressedLength(compressedBytes);
             var decompressedBytes = new byte[uncompressedSize];
             var decompressedSize = Snappy.Decompress(compressedBytes, decompressedBytes);

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

@@ -333,14 +333,15 @@ private IByteBuffer ReceiveBuffer(CancellationToken cancellationToken)
             try
             {
                 var messageSizeBytes = new byte[4];
-                _stream.ReadBytes(messageSizeBytes, 0, 4, cancellationToken);
+                var readTimeout = _stream.CanTimeout ? TimeSpan.FromMilliseconds(_stream.ReadTimeout) : Timeout.InfiniteTimeSpan;
+                _stream.ReadBytes(messageSizeBytes, 0, 4, readTimeout, cancellationToken);
                 var messageSize = BinaryPrimitives.ReadInt32LittleEndian(messageSizeBytes);
                 EnsureMessageSizeIsValid(messageSize);
                 var inputBufferChunkSource = new InputBufferChunkSource(BsonChunkPool.Default);
                 var buffer = ByteBufferFactory.Create(inputBufferChunkSource, messageSize);
                 buffer.Length = messageSize;
                 buffer.SetBytes(0, messageSizeBytes, 0, 4);
-                _stream.ReadBytes(buffer, 4, messageSize - 4, cancellationToken);
+                _stream.ReadBytes(buffer, 4, messageSize - 4, readTimeout, cancellationToken);
                 _lastUsedAtUtc = DateTime.UtcNow;
                 buffer.MakeReadOnly();
                 return buffer;
@@ -535,7 +536,8 @@ private void SendBuffer(IByteBuffer buffer, CancellationToken cancellationToken)
 
                 try
                 {
-                    _stream.WriteBytes(buffer, 0, buffer.Length, cancellationToken);
+                    var writeTimeout = _stream.CanTimeout ? TimeSpan.FromMilliseconds(_stream.WriteTimeout) : Timeout.InfiniteTimeSpan;
+                    _stream.WriteBytes(buffer, 0, buffer.Length, writeTimeout, cancellationToken);
                     _lastUsedAtUtc = DateTime.UtcNow;
                 }
                 catch (Exception ex)

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

@@ -138,7 +138,10 @@ private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancell
 
             if (!connectOperation.IsCompleted)
             {
-                try { socket.Dispose(); } catch { }
+                try
+                {
+                    socket.Dispose();
+                } catch { }
 
                 cancellationToken.ThrowIfCancellationRequested();
                 throw new TimeoutException($"Timed out connecting to {endPoint}. Timeout was {_settings.ConnectTimeout}.");
@@ -164,7 +167,12 @@ private async Task ConnectAsync(Socket socket, EndPoint endPoint, CancellationTo
 
             if (!connectTask.IsCompleted)
             {
-                try { socket.Dispose(); } catch { }
+                try
+                {
+                    socket.Dispose();
+                    // should await on the read task to avoid UnobservedTaskException
+                    await connectTask.ConfigureAwait(false);
+                } catch { }
 
                 cancellationToken.ThrowIfCancellationRequested();
                 throw new TimeoutException($"Timed out connecting to {endPoint}. Timeout was {_settings.ConnectTimeout}.");

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

@@ -36,46 +36,80 @@ 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)
+        public static int Read(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
             {
-                try
-                {
-                    bytesRead = await stream.ReadAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
-                    ChangeState(2); // note: might not actually go to state 2 if already in state 3 or 4
-                }
-                catch when (state == 1)
-                {
-                    try { stream.Dispose(); } catch { }
-                    throw;
-                }
-                catch when (state >= 3)
+                using var manualResetEvent = new ManualResetEventSlim();
+                var readOperation = stream.BeginRead(
+                    buffer,
+                    offset,
+                    count,
+                    state => ((ManualResetEventSlim)state.AsyncState).Set(),
+                    manualResetEvent);
+
+                if (readOperation.IsCompleted || manualResetEvent.Wait(timeout, cancellationToken))
                 {
-                    // a timeout or operation cancelled exception will be thrown instead
+                    return stream.EndRead(readOperation);
                 }
+            }
+            catch (OperationCanceledException)
+            {
+                // Have to suppress OperationCanceledException here, it will be thrown after the stream will be disposed.
+            }
+            catch (ObjectDisposedException)
+            {
+                throw new IOException();
+            }
 
-                if (state == 3) { throw new TimeoutException(); }
-                if (state == 4) { throw new OperationCanceledException(); }
+            try
+            {
+                stream.Dispose();
+            }
+            catch
+            {
+                // Ignore any exceptions
             }
 
-            return bytesRead;
+            cancellationToken.ThrowIfCancellationRequested();
+            throw new TimeoutException();
+        }
 
-            void ChangeState(int to)
+        public static async Task<int> ReadAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
+        {
+            Task<int> readTask = null;
+            try
+            {
+                readTask = stream.ReadAsync(buffer, offset, count);
+                return await readTask.WaitAsync(timeout, cancellationToken).ConfigureAwait(false);
+            }
+            catch (ObjectDisposedException)
             {
-                var from = Interlocked.CompareExchange(ref state, to, 1);
-                if (from == 1 && to >= 3)
+                // It's possible to get ObjectDisposedException when the connection pool was closed with interruptInUseConnections set to true.
+                throw new IOException();
+            }
+            catch (Exception ex) when (ex is OperationCanceledException or TimeoutException)
+            {
+                // await Task.WaitAsync() throws OperationCanceledException in case of cancellation and TimeoutException in case of timeout
+                try
                 {
-                    try { stream.Dispose(); } catch { } // disposing the stream aborts the read attempt
+                    stream.Dispose();
+                    if (readTask != null)
+                    {
+                        // Should await on the task to avoid UnobservedTaskException
+                        await readTask.ConfigureAwait(false);
+                    }
                 }
+                catch
+                {
+                    // Ignore any exceptions
+                }
+
+                throw;
             }
         }
 
-        public static void ReadBytes(this Stream stream, byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public static void ReadBytes(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -84,7 +118,7 @@ public static void ReadBytes(this Stream stream, byte[] buffer, int offset, int
 
             while (count > 0)
             {
-                var bytesRead = stream.Read(buffer, offset, count); // TODO: honor cancellationToken?
+                var bytesRead = stream.Read(buffer, offset, count, timeout, cancellationToken);
                 if (bytesRead == 0)
                 {
                     throw new EndOfStreamException();
@@ -94,7 +128,7 @@ public static void ReadBytes(this Stream stream, byte[] buffer, int offset, int
             }
         }
 
-        public static void ReadBytes(this Stream stream, IByteBuffer buffer, int offset, int count, CancellationToken cancellationToken)
+        public static void ReadBytes(this Stream stream, IByteBuffer buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -105,7 +139,7 @@ public static void ReadBytes(this Stream stream, IByteBuffer buffer, int offset,
             {
                 var backingBytes = buffer.AccessBackingBytes(offset);
                 var bytesToRead = Math.Min(count, backingBytes.Count);
-                var bytesRead = stream.Read(backingBytes.Array, backingBytes.Offset, bytesToRead); // TODO: honor cancellationToken?
+                var bytesRead = stream.Read(backingBytes.Array, backingBytes.Offset, bytesToRead, timeout, cancellationToken);
                 if (bytesRead == 0)
                 {
                     throw new EndOfStreamException();
@@ -155,44 +189,82 @@ 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)
+        public static void Write(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
             {
-                try
-                {
-                    await stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
-                    ChangeState(2); // note: might not actually go to state 2 if already in state 3 or 4
-                }
-                catch when (state == 1)
-                {
-                    try { stream.Dispose(); } catch { }
-                    throw;
-                }
-                catch when (state >= 3)
+                using var manualResetEvent = new ManualResetEventSlim();
+                var writeOperation = stream.BeginWrite(
+                    buffer,
+                    offset,
+                    count,
+                    state => ((ManualResetEventSlim)state.AsyncState).Set(),
+                    manualResetEvent);
+
+                if (writeOperation.IsCompleted || manualResetEvent.Wait(timeout, cancellationToken))
                 {
-                    // a timeout or operation cancelled exception will be thrown instead
+                    stream.EndWrite(writeOperation);
+                    return;
                 }
+            }
+            catch (OperationCanceledException)
+            {
+                // Have to suppress OperationCanceledException here, it will be thrown after the stream will be disposed.
+            }
+            catch (ObjectDisposedException)
+            {
+                // It's possible to get ObjectDisposedException when the connection pool was closed with interruptInUseConnections set to true.
+                throw new IOException();
+            }
 
-                if (state == 3) { throw new TimeoutException(); }
-                if (state == 4) { throw new OperationCanceledException(); }
+            try
+            {
+                stream.Dispose();
             }
+            catch
+            {
+                // Ignore any exceptions
+            }
+
+            cancellationToken.ThrowIfCancellationRequested();
+            throw new TimeoutException();
+        }
 
-            void ChangeState(int to)
+        public static async Task WriteAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
+        {
+            Task writeTask = null;
+            try
+            {
+                writeTask = stream.WriteAsync(buffer, offset, count);
+                await writeTask.WaitAsync(timeout, cancellationToken).ConfigureAwait(false);
+            }
+            catch (ObjectDisposedException)
+            {
+                // It's possible to get ObjectDisposedException when the connection pool was closed with interruptInUseConnections set to true.
+                throw new IOException();
+            }
+            catch (Exception ex) when (ex is OperationCanceledException or TimeoutException)
             {
-                var from = Interlocked.CompareExchange(ref state, to, 1);
-                if (from == 1 && to >= 3)
+                // await Task.WaitAsync() throws OperationCanceledException in case of cancellation and TimeoutException in case of timeout
+                try
                 {
-                    try { stream.Dispose(); } catch { } // disposing the stream aborts the write attempt
+                    stream.Dispose();
+                    // Should await on the task to avoid UnobservedTaskException
+                    if (writeTask != null)
+                    {
+                        await writeTask.ConfigureAwait(false);
+                    }
                 }
+                catch
+                {
+                    // Ignore any exceptions
+                }
+
+                throw;
             }
         }
 
-        public static void WriteBytes(this Stream stream, IByteBuffer buffer, int offset, int count, CancellationToken cancellationToken)
+        public static void WriteBytes(this Stream stream, IByteBuffer buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
             Ensure.IsNotNull(stream, nameof(stream));
             Ensure.IsNotNull(buffer, nameof(buffer));
@@ -204,7 +276,7 @@ public static void WriteBytes(this Stream stream, IByteBuffer buffer, int offset
                 cancellationToken.ThrowIfCancellationRequested();
                 var backingBytes = buffer.AccessBackingBytes(offset);
                 var bytesToWrite = Math.Min(count, backingBytes.Count);
-                stream.Write(backingBytes.Array, backingBytes.Offset, bytesToWrite); // TODO: honor cancellationToken?
+                stream.Write(backingBytes.Array, backingBytes.Offset, bytesToWrite, timeout, cancellationToken);
                 offset += bytesToWrite;
                 count -= bytesToWrite;
             }