From 3f9f0e7a9eb37ca696487198d6e6c992fa813db4 Mon Sep 17 00:00:00 2001
From: Miha Zupan <mihazupan.zupan1@gmail.com>
Date: Mon, 25 Nov 2019 17:11:30 +0100
Subject: [PATCH] Add SmtpClient.SendMailAsync overloads with cancellation

---
 .../System.Net.Mail/ref/System.Net.Mail.cs    |  2 +
 .../src/System/Net/Mail/SmtpClient.cs         | 61 ++++++++++++++++---
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/src/libraries/System.Net.Mail/ref/System.Net.Mail.cs b/src/libraries/System.Net.Mail/ref/System.Net.Mail.cs
index d1b280aec9489..b3c1f32996dca 100644
--- a/src/libraries/System.Net.Mail/ref/System.Net.Mail.cs
+++ b/src/libraries/System.Net.Mail/ref/System.Net.Mail.cs
@@ -186,6 +186,8 @@ public void SendAsync(string from, string recipients, string subject, string bod
         public void SendAsyncCancel() { }
         public System.Threading.Tasks.Task SendMailAsync(System.Net.Mail.MailMessage message) { throw null; }
         public System.Threading.Tasks.Task SendMailAsync(string from, string recipients, string subject, string body) { throw null; }
+        public System.Threading.Tasks.Task SendMailAsync(System.Net.Mail.MailMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }
+        public System.Threading.Tasks.Task SendMailAsync(string from, string recipients, string subject, string body, System.Threading.CancellationToken cancellationToken) { throw null; }
     }
     public enum SmtpDeliveryFormat
     {
diff --git a/src/libraries/System.Net.Mail/src/System/Net/Mail/SmtpClient.cs b/src/libraries/System.Net.Mail/src/System/Net/Mail/SmtpClient.cs
index d67a2660eda60..3ddd9b283e7fc 100644
--- a/src/libraries/System.Net.Mail/src/System/Net/Mail/SmtpClient.cs
+++ b/src/libraries/System.Net.Mail/src/System/Net/Mail/SmtpClient.cs
@@ -786,17 +786,35 @@ public void SendAsyncCancel()
         public Task SendMailAsync(string from, string recipients, string subject, string body)
         {
             var message = new MailMessage(from, recipients, subject, body);
-            return SendMailAsync(message);
+            return SendMailAsync(message, cancellationToken: default);
         }
 
         public Task SendMailAsync(MailMessage message)
         {
+            return SendMailAsync(message, cancellationToken: default);
+        }
+
+        public Task SendMailAsync(string from, string recipients, string subject, string body, CancellationToken cancellationToken)
+        {
+            var message = new MailMessage(from, recipients, subject, body);
+            return SendMailAsync(message, cancellationToken);
+        }
+
+        public Task SendMailAsync(MailMessage message, CancellationToken cancellationToken)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return Task.FromCanceled(cancellationToken);
+            }
+
             // Create a TaskCompletionSource to represent the operation
-            var tcs = new TaskCompletionSource<object>();
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            CancellationTokenRegistration ctr = default;
 
             // Register a handler that will transfer completion results to the TCS Task
             SendCompletedEventHandler handler = null;
-            handler = (sender, e) => HandleCompletion(tcs, e, handler);
+            handler = (sender, e) => HandleCompletion(tcs, ctr, e, handler);
             SendCompleted += handler;
 
             // Start the async operation.
@@ -810,20 +828,45 @@ public Task SendMailAsync(MailMessage message)
                 throw;
             }
 
+            // Only register on the CT if HandleCompletion hasn't started to ensure the CTR is disposed
+            bool lockTaken = false;
+            try
+            {
+                Monitor.TryEnter(tcs, ref lockTaken);
+                if (lockTaken && !tcs.Task.IsCompleted)
+                {
+                    ctr = cancellationToken.Register(s =>
+                    {
+                        ((SmtpClient)s).SendAsyncCancel();
+                    }, this);
+                }
+            }
+            finally
+            {
+                if (lockTaken) Monitor.Exit(tcs);
+            }
+
             // Return the task to represent the asynchronous operation
             return tcs.Task;
         }
 
-        private void HandleCompletion(TaskCompletionSource<object> tcs, AsyncCompletedEventArgs e, SendCompletedEventHandler handler)
+        private void HandleCompletion(TaskCompletionSource<object> tcs, CancellationTokenRegistration ctr, AsyncCompletedEventArgs e, SendCompletedEventHandler handler)
         {
             if (e.UserState == tcs)
             {
-                try { SendCompleted -= handler; }
-                finally
+                lock (tcs)
                 {
-                    if (e.Error != null) tcs.TrySetException(e.Error);
-                    else if (e.Cancelled) tcs.TrySetCanceled();
-                    else tcs.TrySetResult(null);
+                    try
+                    {
+                        SendCompleted -= handler;
+                        ctr.Dispose();
+                    }
+                    finally
+                    {
+                        if (e.Error != null) tcs.TrySetException(e.Error);
+                        else if (e.Cancelled) tcs.TrySetCanceled();
+                        else tcs.TrySetResult(null);
+                    }
                 }
             }
         }