FlowScope/Packages/com.flowscope.gamecore/Runtime/Resources/ResourceService.cs

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;

namespace FlowScope.Resources
{
    public sealed class ResourceService : IResourceService, ICompletedResourceService
    {
        private readonly object _gate = new();
        private readonly IResourceBackend _backend;
        private readonly Dictionary<ResourceKey, ResourceEntry> _entries = new();
        private readonly Dictionary<ResourceKey, SharedLoad> _loads = new();

        public ResourceService(IResourceBackend backend)
        {
            _backend = backend ?? throw new ArgumentNullException(nameof(backend));
        }

        public async Task<IResourceHandle<T>> LoadAsync<T>(
            string key,
            CancellationToken cancellationToken)
            where T : class
        {
            if (string.IsNullOrWhiteSpace(key))
            {
                throw new ArgumentException("Resource key is required.", nameof(key));
            }

            cancellationToken.ThrowIfCancellationRequested();
            var resourceKey = new ResourceKey(key, typeof(T));
            SharedLoad sharedLoad;

            lock (_gate)
            {
                if (_entries.TryGetValue(resourceKey, out var entry))
                {
                    return entry.CreateHandle<T>();
                }

                if (!_loads.TryGetValue(resourceKey, out sharedLoad))
                {
                    sharedLoad = new SharedLoad();
                    _loads.Add(resourceKey, sharedLoad);
                }

                sharedLoad.WaiterCount++;
                if (sharedLoad.LoadTask == null)
                {
                    sharedLoad.LoadTask = LoadAndOwnEntryAsync<T>(resourceKey, key, sharedLoad);
                }
            }

            try
            {
                var loaded = await WaitForSharedLoad(sharedLoad.LoadTask, cancellationToken);
                lock (_gate)
                {
                    return loaded.CreateHandle<T>();
                }
            }
            catch (OperationCanceledException)
            {
                lock (_gate)
                {
                    if (sharedLoad.WaiterCount > 0)
                    {
                        sharedLoad.WaiterCount--;
                    }
                }

                throw;
            }
        }

        public bool TryLoadCompleted<T>(string key, out IResourceHandle<T> handle)
            where T : class
        {
            if (string.IsNullOrWhiteSpace(key))
            {
                throw new ArgumentException("Resource key is required.", nameof(key));
            }

            lock (_gate)
            {
                if (_entries.TryGetValue(new ResourceKey(key, typeof(T)), out var entry))
                {
                    handle = entry.CreateHandle<T>();
                    return true;
                }
            }

            handle = null;
            return false;
        }

        public IResourceGroup CreateGroup()
        {
            return new ResourceGroup();
        }

        private async Task<ResourceEntry> LoadEntryAsync<T>(string key)
            where T : class
        {
            T asset;
            try
            {
                asset = await _backend.LoadAsync<T>(key, CancellationToken.None);
            }
            catch (Exception exception) when (exception is not OperationCanceledException)
            {
                throw new InvalidOperationException(
                    $"Failed to load resource '{key}' as {typeof(T).Name} from backend '{_backend.Name}'.",
                    exception);
            }

            if (asset == null)
            {
                throw new InvalidOperationException(
                    $"Resource '{key}' returned null for {typeof(T).Name} from backend '{_backend.Name}'.");
            }

            return new ResourceEntry(
                key,
                typeof(T),
                asset,
                (entryKey, assetType, releasedAsset) =>
                    ReleaseBackendEntry(entryKey, assetType, (T)releasedAsset));
        }

        private async Task<ResourceEntry> LoadAndOwnEntryAsync<T>(
            ResourceKey resourceKey,
            string key,
            SharedLoad sharedLoad)
            where T : class
        {
            ResourceEntry entry = null;
            try
            {
                entry = await LoadEntryAsync<T>(key);
                lock (_gate)
                {
                    _loads.Remove(resourceKey);
                    if (sharedLoad.WaiterCount > 0)
                    {
                        _entries.Add(resourceKey, entry);
                        return entry;
                    }
                }

                ReleaseEntry(entry);
                return entry;
            }
            catch
            {
                lock (_gate)
                {
                    _loads.Remove(resourceKey);
                }

                throw;
            }
        }

        private static async Task<ResourceEntry> WaitForSharedLoad(
            Task<ResourceEntry> loadTask,
            CancellationToken cancellationToken)
        {
            if (!cancellationToken.CanBeCanceled)
            {
                return await loadTask;
            }

            var cancellation = new TaskCompletionSource<bool>();
            using (cancellationToken.Register(() => cancellation.TrySetResult(true)))
            {
                if (await Task.WhenAny(loadTask, cancellation.Task) == cancellation.Task)
                {
                    throw new OperationCanceledException(cancellationToken);
                }
            }

            return await loadTask;
        }

        private void ReleaseEntry(ResourceEntry entry)
        {
            entry.ReleaseWithoutHandle();
        }

        private void ReleaseBackendEntry<T>(string key, Type assetType, T asset)
            where T : class
        {
            lock (_gate)
            {
                _entries.Remove(new ResourceKey(key, assetType));
            }

            _backend.Release(key, asset);
        }

        private sealed class SharedLoad
        {
            public Task<ResourceEntry> LoadTask;
            public int WaiterCount;
        }

        private readonly struct ResourceKey : IEquatable<ResourceKey>
        {
            private readonly string _key;
            private readonly Type _type;

            public ResourceKey(string key, Type type)
            {
                _key = key;
                _type = type;
            }

            public bool Equals(ResourceKey other)
            {
                return _key == other._key && _type == other._type;
            }

            public override bool Equals(object obj)
            {
                return obj is ResourceKey other && Equals(other);
            }

            public override int GetHashCode()
            {
                unchecked
                {
                    return ((_key != null ? _key.GetHashCode() : 0) * 397) ^
                           (_type != null ? _type.GetHashCode() : 0);
                }
            }
        }
    }
}