using System; using System.Collections.Concurrent; using System.Threading; using System.Threading.Tasks; using UnityEngine; namespace OnDeviceAgent.AgentCore { public sealed class UnityMainThreadDispatcher : MonoBehaviour, IUnityMainThreadDispatcher { static Thread s_MainThread; readonly ConcurrentQueue m_Queue = new ConcurrentQueue(); // Captured at startup on the main thread, before any Awake, so IsOnMainThread is reliable regardless of init order. [RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)] static void CaptureMainThread() => s_MainThread = Thread.CurrentThread; void Update() { while (m_Queue.TryDequeue(out var work)) { try { work(); } catch (Exception ex) { Debug.LogException(ex); } } } public bool IsOnMainThread => Thread.CurrentThread == s_MainThread; public void Post(Action action) { if (action == null) return; if (IsOnMainThread) { try { action(); } catch (Exception ex) { Debug.LogException(ex); } return; } m_Queue.Enqueue(action); } public Task RunOnMainAsync(Func work) { if (work == null) throw new ArgumentNullException(nameof(work)); var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); if (IsOnMainThread) { RunAndComplete(work, tcs); return tcs.Task; } m_Queue.Enqueue(async () => await RunAndCompleteAsync(work, tcs)); return tcs.Task; } public Task RunOnMainAsync(Func> work) { if (work == null) throw new ArgumentNullException(nameof(work)); var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); if (IsOnMainThread) { RunAndComplete(work, tcs); return tcs.Task; } m_Queue.Enqueue(async () => await RunAndCompleteAsync(work, tcs)); return tcs.Task; } static void RunAndComplete(Func work, TaskCompletionSource tcs) { try { var inner = work(); inner.ContinueWith(t => { if (t.IsFaulted) tcs.TrySetException(t.Exception!.InnerExceptions); else if (t.IsCanceled) tcs.TrySetCanceled(); else tcs.TrySetResult(true); }, TaskContinuationOptions.ExecuteSynchronously); } catch (Exception ex) { tcs.TrySetException(ex); } } static async Task RunAndCompleteAsync(Func work, TaskCompletionSource tcs) { try { await work().ConfigureAwait(false); tcs.TrySetResult(true); } catch (OperationCanceledException) { tcs.TrySetCanceled(); } catch (Exception ex) { tcs.TrySetException(ex); } } static void RunAndComplete(Func> work, TaskCompletionSource tcs) { try { var inner = work(); inner.ContinueWith(t => { if (t.IsFaulted) tcs.TrySetException(t.Exception!.InnerExceptions); else if (t.IsCanceled) tcs.TrySetCanceled(); else tcs.TrySetResult(t.Result); }, TaskContinuationOptions.ExecuteSynchronously); } catch (Exception ex) { tcs.TrySetException(ex); } } static async Task RunAndCompleteAsync(Func> work, TaskCompletionSource tcs) { try { var result = await work().ConfigureAwait(false); tcs.TrySetResult(result); } catch (OperationCanceledException) { tcs.TrySetCanceled(); } catch (Exception ex) { tcs.TrySetException(ex); } } } }