pr

Author: sanych-sun

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

@@ -36,17 +36,15 @@ public static void EfficientCopyTo(this Stream input, Stream output)
             }
         }
 
-        public static int Read(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
-        {
-            return ExecuteOperationWithTimeout(
+        public static int Read(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken) =>
+            ExecuteOperationWithTimeout(
                 stream,
                 (str, state) => str.Read(state.Buffer, state.Offset, state.Count),
                 buffer,
                 offset,
                 count,
                 timeout,
                 cancellationToken);
-        }
 
         public static async Task<int> ReadAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
@@ -190,8 +188,7 @@ public static async Task ReadBytesAsync(this Stream stream, byte[] destination,
             }
         }
 
-        public static void Write(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) =>
             ExecuteOperationWithTimeout(
                 stream,
                 (str, state) =>
@@ -204,7 +201,6 @@ public static void Write(this Stream stream, byte[] buffer, int offset, int coun
                 count,
                 timeout,
                 cancellationToken);
-        }
 
         public static async Task WriteAsync(this Stream stream, byte[] buffer, int offset, int count, TimeSpan timeout, CancellationToken cancellationToken)
         {
@@ -304,7 +300,7 @@ private static TResult ExecuteOperationWithTimeout<TResult>(Stream stream, Func<
             }
             catch (IOException)
             {
-                if (callbackState?.OperationState == OperationState.Cancelled)
+                if (callbackState?.OperationState == OperationState.Interrupted)
                 {
                     cancellationToken.ThrowIfCancellationRequested();
                     throw new TimeoutException();
@@ -321,9 +317,9 @@ private static TResult ExecuteOperationWithTimeout<TResult>(Stream stream, Func<
             static void DisposeStreamCallback(object state)
             {
                 var disposeCallbackState = (StreamDisposeCallbackState)state;
-                if (!disposeCallbackState.TryChangeStateFromInProgress(OperationState.Cancelled))
+                if (!disposeCallbackState.TryChangeStateFromInProgress(OperationState.Interrupted))
                 {
-                    // if cannot change the state - then I/O was already succeeded
+                    // If the state can't be changed - then I/O had already succeeded
                     return;
                 }
 
@@ -352,7 +348,7 @@ private enum OperationState
         {
             InProgress = 0,
             Done,
-            Cancelled,
+            Interrupted,
         }
     }
 }

tests/MongoDB.Driver.Tests/Core/Misc/StreamExtensionMethodsTests.cs

@@ -1,4 +1,4 @@
-/* Copyright 2013-present MongoDB Inc.
+/* Copyright 2010-present MongoDB Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -201,6 +201,49 @@ await Record.ExceptionAsync(() => stream.ReadBytesAsync(OperationContext.NoTimeo
                 .ParamName.Should().Be("stream");
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public async Task ReadBytes_with_byte_array_throws_on_timeout([Values(true, false)]bool async)
+        {
+            var streamMock = new Mock<Stream>();
+            SetupStreamRead(streamMock);
+            var stream = streamMock.Object;
+
+            var destination = new byte[2];
+            var timeout = TimeSpan.FromMilliseconds(10);
+
+            var exception = async ?
+                await Record.ExceptionAsync(() => stream.ReadAsync(destination, 0, 2, timeout, CancellationToken.None)) :
+                Record.Exception(() => stream.Read(destination, 0, 2, timeout, CancellationToken.None));
+
+            exception.Should().BeOfType<TimeoutException>();
+        }
+
+        [Theory]
+        [ParameterAttributeData]
+        public async Task ReadBytes_with_byte_array_throws_on_cancellation([Values(true, false)]bool async)
+        {
+            var streamMock = new Mock<Stream>();
+            SetupStreamRead(streamMock);
+            var stream = streamMock.Object;
+
+            var destination = new byte[2];
+            using var cancellationTokenSource = new CancellationTokenSource(10);
+
+            var exception = async ?
+                await Record.ExceptionAsync(() => stream.ReadAsync(destination, 0, 2, Timeout.InfiniteTimeSpan, cancellationTokenSource.Token)) :
+                Record.Exception(() => stream.Read(destination, 0, 2, Timeout.InfiniteTimeSpan, cancellationTokenSource.Token));
+
+            if (async)
+            {
+                exception.Should().BeOfType<TaskCanceledException>();
+            }
+            else
+            {
+                exception.Should().BeOfType<OperationCanceledException>();
+            }
+        }
+
         [Theory]
         [InlineData(true, 0, new byte[] { 0, 0 })]
         [InlineData(true, 1, new byte[] { 1, 0 })]
@@ -529,5 +572,18 @@ private Mock<IByteBuffer> CreateMockByteBuffer(int length)
             mockBuffer.SetupGet(b => b.Length).Returns(length);
             return mockBuffer;
         }
+
+        private void SetupStreamRead(Mock<Stream> streamMock, TaskCompletionSource<int> readTaskCompletionSource = null)
+        {
+            readTaskCompletionSource ??= new TaskCompletionSource<int>();
+            streamMock.Setup(s => s.Close()).Callback(() =>
+            {
+                readTaskCompletionSource.SetException(new IOException());
+            });
+            streamMock.Setup(s => s.Read(It.IsAny<byte[]>(), It.IsAny<int>(), It.IsAny<int>())).Returns(() =>
+                readTaskCompletionSource.Task.GetAwaiter().GetResult());
+            streamMock.Setup(s => s.ReadAsync(It.IsAny<byte[]>(), It.IsAny<int>(), It.IsAny<int>(), It.IsAny<CancellationToken>())).Returns(() =>
+                readTaskCompletionSource.Task);
+        }
     }
 }