Wrapped tasks raise events in the correct order, also handle exceptions properly

Author: shana

src/GitHub.Api/NewTaskSystem/ActionTask.cs

@@ -10,6 +10,14 @@ class ActionTask : TaskBase
         protected Action<bool> Callback { get; }
         protected Action<bool, Exception> CallbackWithException { get; }
 
+        public ActionTask(CancellationToken token, Action action)
+            : base(token)
+        {
+            Guard.ArgumentNotNull(action, "action");
+            this.Callback = _ => action();
+            Name = "ActionTask";
+        }
+
         public ActionTask(CancellationToken token, Action<bool> action)
             : base(token)
         {
@@ -123,6 +131,14 @@ class FuncTask<T> : TaskBase<T>
         protected Func<bool, T> Callback { get; }
         protected Func<bool, Exception, T> CallbackWithException { get; }
 
+        public FuncTask(CancellationToken token, Func<T> action)
+            : base(token)
+        {
+            Guard.ArgumentNotNull(action, "action");
+            this.Callback = _ => action();
+            Name = $"FuncTask<{typeof(T)}>";
+        }
+
         public FuncTask(CancellationToken token, Func<bool, T> action)
             : base(token)
         {

src/GitHub.Api/NewTaskSystem/TaskBase.cs

@@ -80,12 +80,30 @@ public TaskBase(CancellationToken token)
 
         public TaskBase(Task task)
         {
-            task.ContinueWith(t =>
+            Task = new Task(t =>
             {
+                var scheduler = TaskManager.GetScheduler(Affinity);
                 RaiseOnStart();
-                RaiseOnEnd();
-            }, Token, runAlwaysOptions, TaskScheduler.Current);
-            Task = task;
+                var tk = ((Task)t);
+                try
+                {
+                    if (tk.Status == TaskStatus.Created && !tk.IsCompleted &&
+                      ((tk.CreationOptions & (TaskCreationOptions)512) == TaskCreationOptions.None))
+                    {
+                        tk.RunSynchronously(scheduler);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    Errors = ex.Message;
+                    if (!RaiseFaultHandlers(ex))
+                        throw;
+                }
+                finally
+                {
+                    RaiseOnEnd();
+                }
+            }, task, Token, TaskCreationOptions.None);
         }
 
         protected TaskBase() {}
@@ -395,7 +413,33 @@ public TaskBase(CancellationToken token)
         public TaskBase(Task<TResult> task)
             : base()
         {
-            Task = task;
+            var scheduler = TaskManager.GetScheduler(Affinity);
+            Task = new Task<TResult>(t =>
+            {
+                TResult ret = default(TResult);
+                RaiseOnStart();
+                var tk = ((Task<TResult>)t);
+                try
+                {
+                    if (tk.Status == TaskStatus.Created && !tk.IsCompleted &&
+                      ((tk.CreationOptions & (TaskCreationOptions)512) == TaskCreationOptions.None))
+                    {
+                        tk.RunSynchronously();
+                    }
+                    ret = tk.Result;
+                }
+                catch (Exception ex)
+                {
+                    Errors = ex.Message;
+                    if (!RaiseFaultHandlers(ex))
+                        throw;
+                }
+                finally
+                {
+                    RaiseOnEnd();
+                }
+                return ret;
+            }, task, Token, TaskCreationOptions.None);
         }
 
 

src/tests/TaskSystemIntegrationTests/Tests.cs

@@ -656,15 +656,11 @@ class TaskToActionTask : BaseTest
         [Test]
         public async Task CanWrapATask()
         {
-            var uiThread = 0;
-            await new ActionTask(Token, _ => uiThread = Thread.CurrentThread.ManagedThreadId) { Affinity = TaskAffinity.UI }
-                .StartAsAsync();
-
             var runOrder = new List<string>();
-            var task = new Task(() => runOrder.Add($"ran {Thread.CurrentThread.ManagedThreadId}"));
-            var act = new ActionTask(task) { Affinity = TaskAffinity.UI };
+            var task = new Task(() => runOrder.Add($"ran"));
+            var act = new ActionTask(task) { Affinity = TaskAffinity.Exclusive };
             await act.Start().Task;
-            CollectionAssert.AreEqual(new string[] { $"ran {uiThread}" }, runOrder);
+            CollectionAssert.AreEqual(new string[] { $"ran" }, runOrder);
         }
 
         /// <summary>
@@ -820,9 +816,10 @@ public TaskBase Test_GetTopMostTaskInCreatedState()
         }
 
         [Test]
-        public void GetTopMostTaskInCreatedState()
+        public async Task GetTopMostTaskInCreatedState()
         {
-            var task1 = new ActionTask(TaskEx.FromResult(true));
+            var task1 = new ActionTask(Token, () => { });
+            await task1.StartAwait();
             var task2 = new TestActionTask(Token, _ => { });
             var task3 = new TestActionTask(Token, _ => { });
 

src/tests/TaskSystemIntegrationTests/ThreadSynchronizationContext.cs

@@ -49,6 +49,7 @@ private void Wait(long id)
 
         private void Start()
         {
+            SetSynchronizationContext(this);
             threadId = Thread.CurrentThread.ManagedThreadId;
             var lastTime = DateTime.Now.Ticks;
             var wait = new ManualResetEventSlim(false);