Merge pull request #995 from kasobol-msft/backport-net-framework-hang…

…-fix Eagerly dispose network stream if timeout occurs.
Azure · May 26, 2020 · 2afb44c · 2afb44c
2 parents 9ff772a + 5efcb63
commit 2afb44c
Showing 1 changed file with 35 additions and 6 deletions.
diff --git a/Lib/ClassLibraryCommon/Core/TimeoutStream.cs b/Lib/ClassLibraryCommon/Core/TimeoutStream.cs
@@ -32,7 +32,7 @@ internal class TimeoutStream : Stream
         private readonly Stream wrappedStream;
         private TimeSpan readTimeout;
         private TimeSpan writeTimeout;
-        private CancellationTokenSource cancellationTokenSource;
+        private CancellationTokenSource cancellationTokenSource = null!;
 
         public TimeoutStream(Stream wrappedStream, TimeSpan timeout)
             : this(wrappedStream, timeout, timeout) { }
@@ -47,7 +47,7 @@ public TimeoutStream(Stream wrappedStream, TimeSpan readTimeout, TimeSpan writeT
             this.writeTimeout = writeTimeout;
             this.UpdateReadTimeout();
             this.UpdateWriteTimeout();
-            this.cancellationTokenSource = new CancellationTokenSource();
+            this.InitializeTokenSource();
         }
 
         public override long Position
@@ -140,7 +140,13 @@ public override async Task FlushAsync(CancellationToken cancellationToken)
             var source = StartTimeout(cancellationToken, out bool dispose);
             try
             {
-                await this.wrappedStream.FlushAsync(source.Token);
+                await this.wrappedStream.FlushAsync(source.Token).ConfigureAwait(false);
+            }
+            // We dispose stream on timeout so catch and check if cancellation token was cancelled
+            catch (ObjectDisposedException)
+            {
+                source.Token.ThrowIfCancellationRequested();
+                throw;
             }
             finally
             {
@@ -158,7 +164,13 @@ public override async Task<int> ReadAsync(byte[] buffer, int offset, int count,
             var source = StartTimeout(cancellationToken, out bool dispose);
             try
             {
-                return await this.wrappedStream.ReadAsync(buffer, offset, count, source.Token);
+                return await this.wrappedStream.ReadAsync(buffer, offset, count, source.Token).ConfigureAwait(false);
+            }
+            // We dispose stream on timeout so catch and check if cancellation token was cancelled
+            catch (ObjectDisposedException)
+            {
+                source.Token.ThrowIfCancellationRequested();
+                throw;
             }
             finally
             {
@@ -191,7 +203,13 @@ public override async Task WriteAsync(byte[] buffer, int offset, int count, Canc
             var source = StartTimeout(cancellationToken, out bool dispose);
             try
             {
-                await this.wrappedStream.WriteAsync(buffer, offset, count, source.Token);
+                await this.wrappedStream.WriteAsync(buffer, offset, count, source.Token).ConfigureAwait(false);
+            }
+            // We dispose stream on timeout so catch and check if cancellation token was cancelled
+            catch (ObjectDisposedException)
+            {
+                source.Token.ThrowIfCancellationRequested();
+                throw;
             }
             finally
             {
@@ -204,11 +222,22 @@ public override void WriteByte(byte value)
             this.wrappedStream.WriteByte(value);
         }
 
+        private void InitializeTokenSource()
+        {
+            this.cancellationTokenSource = new CancellationTokenSource();
+            this.cancellationTokenSource.Token.Register(() => DisposeStream());
+        }
+
+        private void DisposeStream()
+        {
+            this.wrappedStream.Dispose();
+        }
+
         private CancellationTokenSource StartTimeout(CancellationToken additionalToken, out bool dispose)
         {
             if (this.cancellationTokenSource.IsCancellationRequested)
             {
-                this.cancellationTokenSource = new CancellationTokenSource();
+                this.InitializeTokenSource();
             }
 
             CancellationTokenSource source;