From 248989b3ebe9de5a1b24341774670be6533825d3 Mon Sep 17 00:00:00 2001
From: smoogipoo <smoogipoo@smgi.me>
Date: Wed, 27 Jan 2021 01:20:50 +0900
Subject: [PATCH] wip

---
 osu.Game.Tests/NonVisual/TaskChainTest.cs     | 38 ++++++++++++++++--
 .../Multiplayer/StatefulMultiplayerClient.cs  |  2 +-
 osu.Game/Utils/TaskChain.cs                   | 39 +++++++------------
 3 files changed, 51 insertions(+), 28 deletions(-)

diff --git a/osu.Game.Tests/NonVisual/TaskChainTest.cs b/osu.Game.Tests/NonVisual/TaskChainTest.cs
index d561fb4c1b..0a56468818 100644
--- a/osu.Game.Tests/NonVisual/TaskChainTest.cs
+++ b/osu.Game.Tests/NonVisual/TaskChainTest.cs
@@ -4,6 +4,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using NUnit.Framework;
+using osu.Game.Extensions;
 using osu.Game.Utils;
 
 namespace osu.Game.Tests.NonVisual
@@ -13,14 +14,22 @@ namespace osu.Game.Tests.NonVisual
     {
         private TaskChain taskChain;
         private int currentTask;
+        private CancellationTokenSource globalCancellationToken;
 
         [SetUp]
         public void Setup()
         {
+            globalCancellationToken = new CancellationTokenSource();
             taskChain = new TaskChain();
             currentTask = 0;
         }
 
+        [TearDown]
+        public void TearDown()
+        {
+            globalCancellationToken?.Cancel();
+        }
+
         [Test]
         public async Task TestChainedTasksRunSequentially()
         {
@@ -65,17 +74,40 @@ namespace osu.Game.Tests.NonVisual
             Assert.That(task3.task.Result, Is.EqualTo(2));
         }
 
+        [Test]
+        public async Task TestChainedTaskDoesNotCompleteBeforeChildTasks()
+        {
+            var mutex = new ManualResetEventSlim(false);
+
+            var task = taskChain.Add(async () =>
+            {
+                await Task.Run(() => mutex.Wait(globalCancellationToken.Token)).CatchUnobservedExceptions();
+            });
+
+            // Allow task to potentially complete
+            Thread.Sleep(1000);
+
+            Assert.That(task.IsCompleted, Is.False);
+
+            // Allow the task to complete.
+            mutex.Set();
+
+            await task;
+        }
+
         private (Task<int> task, ManualResetEventSlim mutex, CancellationTokenSource cancellation) addTask()
         {
             var mutex = new ManualResetEventSlim(false);
-            var cancellationSource = new CancellationTokenSource();
             var completionSource = new TaskCompletionSource<int>();
 
+            var cancellationSource = new CancellationTokenSource();
+            var token = CancellationTokenSource.CreateLinkedTokenSource(cancellationSource.Token, globalCancellationToken.Token);
+
             taskChain.Add(() =>
             {
-                mutex.Wait(CancellationToken.None);
+                mutex.Wait(globalCancellationToken.Token);
                 completionSource.SetResult(Interlocked.Increment(ref currentTask));
-            }, cancellationSource.Token);
+            }, token.Token);
 
             return (completionSource.Task, mutex, cancellationSource);
         }
diff --git a/osu.Game/Online/Multiplayer/StatefulMultiplayerClient.cs b/osu.Game/Online/Multiplayer/StatefulMultiplayerClient.cs
index 3d8ab4b4c7..5c6a0d34e0 100644
--- a/osu.Game/Online/Multiplayer/StatefulMultiplayerClient.cs
+++ b/osu.Game/Online/Multiplayer/StatefulMultiplayerClient.cs
@@ -182,7 +182,7 @@ namespace osu.Game.Online.Multiplayer
             return joinOrLeaveTaskChain.Add(async () =>
             {
                 await scheduledReset;
-                await LeaveRoomInternal().CatchUnobservedExceptions();
+                await LeaveRoomInternal();
             });
         }
 
diff --git a/osu.Game/Utils/TaskChain.cs b/osu.Game/Utils/TaskChain.cs
index 2bc2c00e28..30aea7578f 100644
--- a/osu.Game/Utils/TaskChain.cs
+++ b/osu.Game/Utils/TaskChain.cs
@@ -14,8 +14,8 @@ namespace osu.Game.Utils
     /// </summary>
     public class TaskChain
     {
-        private readonly object currentTaskLock = new object();
-        private Task? currentTask;
+        private readonly object finalTaskLock = new object();
+        private Task? finalTask;
 
         /// <summary>
         /// Adds a new task to the end of this <see cref="TaskChain"/>.
@@ -23,31 +23,22 @@ namespace osu.Game.Utils
         /// <param name="action">The action to be executed.</param>
         /// <param name="cancellationToken">The <see cref="CancellationToken"/> for this task. Does not affect further tasks in the chain.</param>
         /// <returns>The awaitable <see cref="Task"/>.</returns>
-        public Task Add(Action action, CancellationToken cancellationToken = default)
+        public async Task Add(Action action, CancellationToken cancellationToken = default)
         {
-            lock (currentTaskLock)
-            {
-                // Note: Attaching the cancellation token to the continuation could lead to re-ordering of tasks in the chain.
-                // Therefore, the cancellation token is not used to cancel the continuation but only the run of each task.
-                if (currentTask == null)
-                {
-                    currentTask = Task.Run(() =>
-                    {
-                        cancellationToken.ThrowIfCancellationRequested();
-                        action();
-                    }, CancellationToken.None);
-                }
-                else
-                {
-                    currentTask = currentTask.ContinueWith(_ =>
-                    {
-                        cancellationToken.ThrowIfCancellationRequested();
-                        action();
-                    }, CancellationToken.None);
-                }
+            Task? previousTask;
+            Task currentTask;
 
-                return currentTask;
+            lock (finalTaskLock)
+            {
+                previousTask = finalTask;
+                finalTask = currentTask = new Task(action, cancellationToken);
             }
+
+            if (previousTask != null)
+                await previousTask;
+
+            currentTask.Start();
+            await currentTask;
         }
     }
 }