Merge pull request #3 from BastianBlokland/feature/destroy-runners-wh…

…en-done Destroy runners when work is done
BastianBlokland · Jul 23, 2019 · 10f0bd5 · 10f0bd5
2 parents 4eaeda6 + 4f0e866
commit 10f0bd5
Show file tree

Hide file tree

Showing 3 changed files with 197 additions and 21 deletions.
diff --git a/src/ComponentTask/Internal/MonoBehaviourTaskRunner.cs b/src/ComponentTask/Internal/MonoBehaviourTaskRunner.cs
@@ -15,56 +15,84 @@ public MonoBehaviourTaskRunner()
 
         public UnityEngine.Component ComponentToFollow { get; set; }
 
-        public Task StartTask(Func<Task> taskCreator) =>
-            this.taskRunner.StartTask(taskCreator);
+        public bool IsFinished { get; private set; }
 
-        public Task StartTask(Func<CancellationToken, Task> taskCreator) =>
-            this.taskRunner.StartTask(taskCreator);
+        public Task StartTask(Func<Task> taskCreator)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator);
+        }
 
-        public Task StartTask<TIn>(Func<TIn, Task> taskCreator, TIn data) =>
-            this.taskRunner.StartTask(taskCreator, data);
+        public Task StartTask(Func<CancellationToken, Task> taskCreator)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator);
+        }
 
-        public Task StartTask<TIn>(Func<TIn, CancellationToken, Task> taskCreator, TIn data) =>
-            this.taskRunner.StartTask(taskCreator, data);
+        public Task StartTask<TIn>(Func<TIn, Task> taskCreator, TIn data)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator, data);
+        }
 
-        public Task<TOut> StartTask<TOut>(Func<Task<TOut>> taskCreator) =>
-            this.taskRunner.StartTask(taskCreator);
+        public Task StartTask<TIn>(Func<TIn, CancellationToken, Task> taskCreator, TIn data)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator, data);
+        }
 
-        public Task<TOut> StartTask<TOut>(Func<CancellationToken, Task<TOut>> taskCreator) =>
-            this.taskRunner.StartTask(taskCreator);
+        public Task<TOut> StartTask<TOut>(Func<Task<TOut>> taskCreator)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator);
+        }
 
-        public Task<TOut> StartTask<TIn, TOut>(Func<TIn, Task<TOut>> taskCreator, TIn data) =>
-            this.taskRunner.StartTask(taskCreator, data);
+        public Task<TOut> StartTask<TOut>(Func<CancellationToken, Task<TOut>> taskCreator)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator);
+        }
 
-        public Task<TOut> StartTask<TIn, TOut>(Func<TIn, CancellationToken, Task<TOut>> taskCreator, TIn data) =>
-            this.taskRunner.StartTask(taskCreator, data);
+        public Task<TOut> StartTask<TIn, TOut>(Func<TIn, Task<TOut>> taskCreator, TIn data)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator, data);
+        }
+
+        public Task<TOut> StartTask<TIn, TOut>(Func<TIn, CancellationToken, Task<TOut>> taskCreator, TIn data)
+        {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Task was started on already finished runner");
+            return this.taskRunner.StartTask(taskCreator, data);
+        }
 
         // Dynamically called from the Unity runtime.
         private void LateUpdate()
         {
+            System.Diagnostics.Debug.Assert(!this.IsFinished, "Already finished runner was updated");
+
             // Check if we have a 'ComponentToFollow' assigned.
             if (this.ComponentToFollow is null)
             {
                 // If not then always just execute the runner.
-                this.taskRunner.Execute();
+                this.Execute();
             }
             else
             {
                 // If the component we are following has been destroyed then we destroy ourselves.
                 if (!this.ComponentToFollow)
-                    UnityEngine.Object.Destroy(this);
+                    this.Destroy();
                 else
                 {
                     // If the component is a 'Behaviour' then we update when its enabled.
                     if (ComponentToFollow is UnityEngine.Behaviour behaviour)
                     {
                         if (behaviour.isActiveAndEnabled)
-                            this.taskRunner.Execute();
+                            this.Execute();
                     }
                     else
                     {
                         // Otherwise we always update.
-                        this.taskRunner.Execute();
+                        this.Execute();
                     }
                 }
             }
@@ -73,6 +101,21 @@ private void LateUpdate()
         // Dynamically called from the Unity runtime.
         private void OnDestroy() => this.taskRunner.Dispose();
 
+        private void Execute()
+        {
+            var workRemaining = this.taskRunner.Execute();
+
+            // If we've finished all the work then destroy ourselves.
+            if (!workRemaining)
+                this.Destroy();
+        }
+
+        private void Destroy()
+        {
+            this.IsFinished = true;
+            UnityEngine.Object.Destroy(this);
+        }
+
         void IExceptionHandler.Handle(Exception exception)
         {
             if (exception is null)

diff --git a/src/ComponentTask/LocalTaskRunner.cs b/src/ComponentTask/LocalTaskRunner.cs
@@ -169,11 +169,13 @@ public Task<TOut> StartTask<TIn, TOut>(Func<TIn, CancellationToken, Task<TOut>>
         /// <summary>
         /// Execute all the work that was 'scheduled' by the tasks running on this runner.
         /// </summary>
-        public void Execute()
+        /// <returns>True if still running work, False if all work has finished.</returns>
+        public bool Execute()
         {
             if (this.isDisposed)
                 throw new ObjectDisposedException(nameof(LocalTaskRunner));
 
+            bool tasksRemaining;
             try
             {
                 // Execute all the work that was scheduled on this runner.
@@ -192,8 +194,12 @@ public void Execute()
                         if (this.runningTasks[i].IsFinished)
                             this.runningTasks.RemoveAt(i);
                     }
+
+                    tasksRemaining = this.runningTasks.Count > 0;
                 }
             }
+
+            return tasksRemaining;
         }
 
         /// <inheritdoc/>

diff --git a/tests/playmode/ComponentExtensionsTests.cs b/tests/playmode/ComponentExtensionsTests.cs
@@ -164,6 +164,50 @@ async Task IncrementCountAsync()
             }
         }
 
+        [UnityTest]
+        public IEnumerator TaskPausesWhenComponentIsDisabled()
+        {
+            var count = 0;
+            var go = new GameObject("TestGameObject");
+            var comp = go.AddComponent<MockComponent>();
+            comp.StartTask(IncrementCountAsync);
+
+            // Assert task is running.
+            yield return null;
+            Assert.AreEqual(1, count);
+
+            // Disable component.
+            comp.enabled = false;
+
+            // Assert task is paused.
+            yield return null;
+            Assert.AreEqual(1, count);
+            yield return null;
+            Assert.AreEqual(1, count);
+
+            // Enable component.
+            comp.enabled = true;
+
+            // Assert task is running.
+            yield return null;
+            Assert.AreEqual(2, count);
+            yield return null;
+            Assert.AreEqual(3, count);
+
+            // Cleanup.
+            Object.Destroy(go);
+
+            async Task IncrementCountAsync()
+            {
+                await Task.Yield();
+                count++;
+                await Task.Yield();
+                count++;
+                await Task.Yield();
+                count++;
+            }
+        }
+
         [UnityTest]
         public IEnumerator SameRunnerIsReusedForTheSameComponent()
         {
@@ -178,6 +222,89 @@ public IEnumerator SameRunnerIsReusedForTheSameComponent()
             Object.Destroy(go);
         }
 
+        [UnityTest]
+        public IEnumerator RunnerDestroysItselfWhenComponentIsDestroyed()
+        {
+            var go = new GameObject("TestGameObject");
+            var comp = go.AddComponent<MockComponent>();
+            comp.StartTask(TestAsync);
+
+            // Assert that runner was created.
+            Assert.AreEqual(2, go.GetComponents<MonoBehaviour>().Length);
+
+            // Destroy component.
+            Object.DestroyImmediate(comp);
+            yield return null;
+
+            // Assert that runner has destroyed itself.
+            Assert.AreEqual(0, go.GetComponents<MonoBehaviour>().Length);
+
+            // Cleanup.
+            Object.Destroy(go);
+
+            async Task TestAsync()
+            {
+                await Task.Yield();
+                await Task.Yield();
+            }
+        }
+
+        [UnityTest]
+        public IEnumerator RunnerDestroysItselfWhenDisabledComponentIsDestroyed()
+        {
+            var go = new GameObject("TestGameObject");
+            var comp = go.AddComponent<MockComponent>();
+            comp.StartTask(TestAsync);
+
+            // Disable component.
+            comp.enabled = false;
+
+            // Assert that runner was created.
+            Assert.AreEqual(2, go.GetComponents<MonoBehaviour>().Length);
+
+            // Destroy component.
+            Object.DestroyImmediate(comp);
+            yield return null;
+
+            // Assert that runner has destroyed itself.
+            Assert.AreEqual(0, go.GetComponents<MonoBehaviour>().Length);
+
+            // Cleanup.
+            Object.Destroy(go);
+
+            async Task TestAsync()
+            {
+                await Task.Yield();
+                await Task.Yield();
+            }
+        }
+
+        [UnityTest]
+        public IEnumerator RunnerDestroysItselfWhenWorkIsDone()
+        {
+            var go = new GameObject("TestGameObject");
+            var comp = go.AddComponent<MockComponent>();
+            comp.StartTask(TestAsync);
+
+            // Assert that runner was created.
+            Assert.AreEqual(2, go.GetComponents<MonoBehaviour>().Length);
+
+            // Wait for work to finish.
+            yield return null;
+
+            // Assert that runner has destroyed itself.
+            Assert.AreEqual(1, go.GetComponents<MonoBehaviour>().Length);
+
+            // Cleanup.
+            yield return null;
+            Object.Destroy(go);
+
+            async Task TestAsync()
+            {
+                await Task.Yield();
+            }
+        }
+
         [UnityTest]
         public IEnumerator ThrowsWhenCalledFromNonUnityThread()
         {