/
SafeTask.cs
131 lines (118 loc) · 5.11 KB
/
SafeTask.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
using System;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;
namespace UnityUtilities
{
/// <summary>
/// A replacement for `Task.Run()` that cancels tasks when entering or
/// exiting play mode in the Unity editor (which doesn't happen by default).
///
/// Also registers an UnobservedTaskException handler to prevent exceptions
/// from being swallowed in all Tasks (including SafeTasks), which would
/// happen when these are not awaited or are chained with `.ContinueWith()`.
///
/// Unity 2023.1 introduced `Awaitable` and its `BackgroundThreadAsync()`
/// method that is essentially a wrapper around `Task.Run()`, but the issues
/// addressed by this class remain - so it remains relevant (see
/// https://marcospereira.me/notes/#multithreading).
/// </summary>
public static class SafeTask
{
private static CancellationTokenSource cancellationTokenSource = new();
public static Task<TResult> Run<TResult>(Func<Task<TResult>> f) =>
SafeTask.Run<TResult>((object)f);
public static Task<TResult> Run<TResult>(Func<TResult> f) =>
SafeTask.Run<TResult>((object)f);
public static Task Run(Func<Task> f) => SafeTask.Run<object>((object)f);
public static Task Run(Action f) => SafeTask.Run<object>((object)f);
private static async Task<TResult> Run<TResult>(object f)
{
// We use tokens and not the cancellation source directly as it is
// replaced with a new one upon exiting play or edit mode.
CancellationToken token = CancellationToken.None;
TResult result = default;
// Pending tasks when entering/exiting play mode are only a problem
// in the editor.
if (Application.isEditor)
{
SafeTask.cancellationTokenSource ??= new();
token = SafeTask.cancellationTokenSource.Token;
}
try
{
// Pass token to Task.Run() as well, otherwise upon cancelling
// its status will change to faulted instead of cancelled.
// https://stackoverflow.com/a/72145763/2037431
if (f is Func<Task<TResult>> g)
{
result = await Task.Run(() => g(), token);
}
else if (f is Func<TResult> h)
{
result = await Task.Run(() => h(), token);
}
else if (f is Func<Task> i)
{
await Task.Run(() => i(), token);
}
else if (f is Action j)
{
await Task.Run(() => j(), token);
}
}
catch (Exception e)
{
// We log unobserved exceptions with an UnobservedTaskException
// handler, but those are only handled when garbage collection happens.
// We thus force exceptions to be logged here - at least for SafeTasks.
// If a failing SafeTask is awaited, the exception will be logged twice, but that's
// ok.
UnityEngine.Debug.LogException(e);
throw;
}
if (token.IsCancellationRequested)
{
throw new OperationCanceledException(
"An asynchronous task has been canceled due to entering or exiting play mode.",
token
);
}
return result;
}
#if UNITY_EDITOR
[UnityEditor.InitializeOnLoadMethod]
private static void OnLoadCallback()
{
// Prevent unobserved task exceptions from being swallowed.
// This happens when:
// * A Task that isn't awaited fails;
// * A Task chained with `.ContinueWith()` fails and exceptions are
// not explicitly handled in the function passed to it.
//
// This event handler works for both Tasks and SafeTasks.
//
// Note this only seems to run when garbage collection happens (such
// as after script reloading in the Unity editor).
// Calling `System.GC.Collect()` after the exception caused
// exceptions to be logged right away.
TaskScheduler.UnobservedTaskException += (_, e) =>
UnityEngine.Debug.LogException(e.Exception);
// Cancel pending `SafeTask.Run()` calls when exiting play or edit
// mode.
UnityEditor.EditorApplication.playModeStateChanged += (change) =>
{
if (
change == UnityEditor.PlayModeStateChange.ExitingPlayMode
|| change == UnityEditor.PlayModeStateChange.ExitingEditMode
)
{
SafeTask.cancellationTokenSource.Cancel();
SafeTask.cancellationTokenSource.Dispose();
SafeTask.cancellationTokenSource = new CancellationTokenSource();
}
};
}
#endif
}
}