Runtime/AgentCore/Concurrency/UnityMainThreadDispatcher.cs

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<Action> m_Queue = new ConcurrentQueue<Action>();

        // 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<Task> work)
        {
            if (work == null) throw new ArgumentNullException(nameof(work));

            var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

            if (IsOnMainThread)
            {
                RunAndComplete(work, tcs);
                return tcs.Task;
            }

            m_Queue.Enqueue(async () => await RunAndCompleteAsync(work, tcs));
            return tcs.Task;
        }

        public Task<T> RunOnMainAsync<T>(Func<Task<T>> work)
        {
            if (work == null) throw new ArgumentNullException(nameof(work));

            var tcs = new TaskCompletionSource<T>(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<Task> work, TaskCompletionSource<bool> 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<Task> work, TaskCompletionSource<bool> tcs)
        {
            try
            {
                await work().ConfigureAwait(false);
                tcs.TrySetResult(true);
            }
            catch (OperationCanceledException)
            {
                tcs.TrySetCanceled();
            }
            catch (Exception ex)
            {
                tcs.TrySetException(ex);
            }
        }

        static void RunAndComplete<T>(Func<Task<T>> work, TaskCompletionSource<T> 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<T>(Func<Task<T>> work, TaskCompletionSource<T> tcs)
        {
            try
            {
                var result = await work().ConfigureAwait(false);
                tcs.TrySetResult(result);
            }
            catch (OperationCanceledException)
            {
                tcs.TrySetCanceled();
            }
            catch (Exception ex)
            {
                tcs.TrySetException(ex);
            }
        }
    }
}