MAke sure in-thread finally handlers are always called when appending chains

Author: shana

src/GitHub.Api/Git/RepositoryManager.cs

@@ -186,17 +186,15 @@ public int WaitForEvents()
 
         public ITask CommitAllFiles(string message, string body)
         {
-            var task = GitClient
-                .AddAll()
+            var task = GitClient.AddAll()
                 .Then(GitClient.Commit(message, body));
 
             return HookupHandlers(task, true);
         }
 
         public ITask CommitFiles(List<string> files, string message, string body)
         {
-            var task = GitClient
-                .Add(files)
+            var task = GitClient.Add(files)
                 .Then(GitClient.Commit(message, body));
 
             return HookupHandlers(task, true);
@@ -332,18 +330,9 @@ public ITask DiscardChanges(GitStatusEntry[] gitStatusEntries)
                         }
                     }
 
-                    ITask<string> gitDiscardTask = null;
                     if (itemsToRevert.Any())
                     {
-                        gitDiscardTask = GitClient.Discard(itemsToRevert);
-                        task.Then(gitDiscardTask)
-                            // we're appending a new continuation after HookupHandlers has run,
-                            // we need to set the finally handler manually so it runs at the end of everything
-                            .Finally(s =>
-                            {
-                                watcher.Start();
-                                isBusy = false;
-                            });
+                        task.Then(GitClient.Discard(itemsToRevert));
                     }
                 }
                 , () => gitStatusEntries);

src/GitHub.Api/Tasks/TaskBase.cs

@@ -1,5 +1,6 @@
 using GitHub.Logging;
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -160,6 +161,7 @@ public virtual T Then<T>(T nextTask, TaskRunOptions runOptions = TaskRunOptions.
                 nextTaskBase = nextTaskBase.GetTopMostTask() ?? nextTaskBase;
             // make the next task dependent on this one so it can get values from us
             nextTaskBase.SetDependsOn(this);
+            var nextTaskFinallyHandler = nextTaskBase.finallyHandler;
 
             if (runOptions == TaskRunOptions.OnSuccess)
             {
@@ -173,8 +175,8 @@ public virtual T Then<T>(T nextTask, TaskRunOptions runOptions = TaskRunOptions.
                     SetFaultHandler(nextTaskBase.continuationOnAlways);
                 if (nextTaskBase.catchHandler != null)
                     Catch(nextTaskBase.catchHandler);
-                if (nextTaskBase.finallyHandler != null)
-                    Finally(nextTaskBase.finallyHandler);
+                if (nextTaskFinallyHandler != null)
+                    Finally(nextTaskFinallyHandler);
             }
             else if (runOptions == TaskRunOptions.OnFailure)
             {
@@ -186,6 +188,18 @@ public virtual T Then<T>(T nextTask, TaskRunOptions runOptions = TaskRunOptions.
                 this.continuationOnAlways = nextTaskBase;
                 DependsOn?.SetFaultHandler(nextTaskBase);
             }
+
+            // if the current task has a fault handler, attach it to the chain we're appending
+            if (finallyHandler != null)
+            {
+                var endOfChainTask = nextTaskBase.GetBottomMostTask();
+                while (endOfChainTask != this && endOfChainTask != null)
+                {
+                    endOfChainTask.finallyHandler += finallyHandler;
+                    endOfChainTask = endOfChainTask.DependsOn;
+                }
+            }
+
             return nextTask;
         }
 
@@ -240,11 +254,7 @@ public ITask Finally(Action<bool, Exception> actionToContinueWith, TaskAffinity
         public T Finally<T>(T taskToContinueWith)
             where T : ITask
         {
-            Guard.ArgumentNotNull(taskToContinueWith, nameof(taskToContinueWith));
-            continuationOnAlways = (TaskBase)(object)taskToContinueWith;
-            continuationOnAlways.SetDependsOn(this);
-            DependsOn?.SetFaultHandler(continuationOnAlways);
-            return taskToContinueWith;
+            return Then(taskToContinueWith, TaskRunOptions.OnAlways);
         }
 
         /// <summary>
@@ -363,6 +373,15 @@ protected TaskBase GetTopMostTask()
             return depends.GetTopMostTask(null, false);
         }
 
+        protected TaskBase GetBottomMostTask()
+        {
+            if (continuationOnSuccess != null)
+                return continuationOnSuccess.GetBottomMostTask();
+            else if (continuationOnAlways != null)
+                return continuationOnAlways.GetBottomMostTask();
+            return this;
+        }
+
         protected TaskBase GetTopMostTask(TaskBase ret, bool onlyCreatedState)
         {
             ret = (!onlyCreatedState || Task.Status == TaskStatus.Created ? this : ret);