Code cleanup

Azure · Dec 4, 2018 · 6abcb2d · 6abcb2d
1 parent 8a437d6
commit 6abcb2d
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 1,023 deletions.
diff --git a/src/Cbs/AmqpCbsLink.cs b/src/Cbs/AmqpCbsLink.cs
diff --git a/src/Fx/IteratorAsyncResult.cs b/src/Fx/IteratorAsyncResult.cs
diff --git a/src/Fx/Singleton.cs b/src/Fx/Singleton.cs
@@ -50,7 +50,7 @@ public Task CloseAsync()
         {
             this.Dispose();
 
-            return TaskHelpers.CompletedTask;
+            return Task.CompletedTask;
         }
 
         public void Close()

diff --git a/src/Fx/TaskHelpers.cs b/src/Fx/TaskHelpers.cs
diff --git a/src/ReceivingAmqpLink.cs b/src/ReceivingAmqpLink.cs
@@ -117,7 +117,7 @@ public IAsyncResult BeginReceiveRemoteMessages(int messageCount, TimeSpan batchW
 
         public Task<AmqpMessage> ReceiveMessageAsync(TimeSpan timeout)
         {
-            return TaskHelpers.CreateTask(
+            return Task.Factory.FromAsync(
                 (c, s) => ((ReceivingAmqpLink)s).BeginReceiveMessage(timeout, c, s),
                 (a) =>
                 {
@@ -244,8 +244,8 @@ public Task<Outcome> DisposeMessageAsync(ArraySegment<byte> deliveryTag, Outcome
 
         public Task<Outcome> DisposeMessageAsync(ArraySegment<byte> deliveryTag, ArraySegment<byte> txnId, Outcome outcome, bool batchable, TimeSpan timeout)
         {
-            return TaskHelpers.CreateTask(
-                (c, s) => this.BeginDisposeMessage(deliveryTag, txnId, outcome, batchable, timeout, c, s),
+            return Task.Factory.FromAsync(
+                (c, s) => ((ReceivingAmqpLink)s).BeginDisposeMessage(deliveryTag, txnId, outcome, batchable, timeout, c, s),
                 a => ((ReceivingAmqpLink)a.AsyncState).EndDisposeMessage(a),
                 this);
         }

diff --git a/src/RequestResponseAmqpLink.cs b/src/RequestResponseAmqpLink.cs
@@ -128,7 +128,7 @@ public AmqpSession Session
 
         public Task<AmqpMessage> RequestAsync(AmqpMessage request, TimeSpan timeout)
         {
-            return TaskHelpers.CreateTask(
+            return Task.Factory.FromAsync(
                 (c, s) => ((RequestResponseAmqpLink)s).BeginRequest(request, timeout, c, s),
                 (r) => ((RequestResponseAmqpLink)r.AsyncState).EndRequest(r),
                 this);

diff --git a/src/Transport/TransportStream.cs b/src/Transport/TransportStream.cs
@@ -92,9 +92,9 @@ public override void SetLength(long value)
 
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
-            return TaskHelpers.CreateTask(
-                (c, s) => this.BeginWrite(buffer, offset, count, c, s),
-                (a) => this.EndWrite(a),
+            return Task.Factory.FromAsync(
+                (c, s) => ((Stream)s).BeginWrite(buffer, offset, count, c, s),
+                (a) => ((Stream)a.AsyncState).EndWrite(a),
                 this);
         }
 
@@ -124,9 +124,9 @@ public override void EndWrite(IAsyncResult asyncResult)
 
         public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
-            return TaskHelpers.CreateTask(
-                (c, s) => this.BeginRead(buffer, offset, count, c, s),
-                (a) => this.EndRead(a),
+            return Task.Factory.FromAsync(
+                (c, s) => ((Stream)s).BeginRead(buffer, offset, count, c, s),
+                (a) => ((Stream)a.AsyncState).EndRead(a),
                 this);
         }
 

diff --git a/src/Transport/WebSocketTransportInitiator.cs b/src/Transport/WebSocketTransportInitiator.cs
@@ -27,7 +27,12 @@ public override bool ConnectAsync(TimeSpan timeout, TransportAsyncCallbackArgs c
                 cws.Options.Proxy = this.settings.Proxy;
             }
 
-            Task task = cws.ConnectAsync(this.settings.Uri, CancellationToken.None).WithTimeout(timeout, () => "timeout");
+            var task = new TimeoutTaskSource<ClientWebSocket>(
+                cws,
+                s => s.ConnectAsync(this.settings.Uri, CancellationToken.None),
+                s => s.Abort(),
+                timeout).Task;
+
             if (task.IsCompleted)
             {
                 callbackArgs.Transport = new WebSocketTransport(cws, this.settings.Uri);
@@ -53,5 +58,48 @@ public override bool ConnectAsync(TimeSpan timeout, TransportAsyncCallbackArgs c
             });
             return true;
         }
+
+        sealed class TimeoutTaskSource<T> : TaskCompletionSource<T> where T : class
+        {
+            readonly T t;
+            readonly TimeSpan timeout;
+            readonly ITimer timer;
+            readonly Action<T> onTimeout;
+
+            public TimeoutTaskSource(T t, Func<T, Task> onStart, Action<T> onTimeout, TimeSpan timeout)
+            {
+                this.t = t;
+                this.onTimeout = onTimeout;
+                this.timeout = timeout;
+                this.timer = SystemTimerFactory.Default.Create(OnTimer, this, timeout);
+
+                Task task = onStart(t);
+                task.ContinueWith((_t, _s) => ((TimeoutTaskSource<T>)_s).OnTask(_t), this);
+            }
+
+            static void OnTimer(object state)
+            {
+                var thisPtr = (TimeoutTaskSource<T>)state;
+                thisPtr.onTimeout(thisPtr.t);
+                thisPtr.TrySetException(new TimeoutException(AmqpResources.GetString(AmqpResources.AmqpTimeout, thisPtr.timeout, typeof(T).Name)));
+            }
+
+            void OnTask(Task inner)
+            {
+                this.timer.Cancel();
+                if (inner.IsFaulted)
+                {
+                    this.TrySetException(inner.Exception.InnerException);
+                }
+                else if (inner.IsCanceled)
+                {
+                    this.TrySetCanceled();
+                }
+                else
+                {
+                    this.TrySetResult(this.t);
+                }
+            }
+        }
     }
 }