ADOFAIRenderer/Util/Audio/Clip/AudioClipStore.cs

using System;
using System.Collections.Generic;
using System.Threading;
using MelonLoader;

namespace ADOFAIRenderer.Util.Audio.Clip;

/// <summary>
/// Clip sources of one render, keyed by clip name or path, with a shared pool of chunk arrays and one decoder thread.
///
/// Threads:
/// - Decoder thread (<see cref="Start"/>): the only thread that decodes, resamples, loads and frees chunks.
///   Each step loads one chunk, the one with the lowest priority value over all sources
///   (a chunk a consumer waits for, then cursor windows by slot, then resident chunks in order).
/// - Consumers (render thread, main thread for the spectrum): <see cref="Acquire"/> waits for the decoder thread and never loads.
/// - Main thread: <see cref="GetOrAdd"/> creates readers (for Unity clips this copies the clip data).
///
/// Memory has no limit; current and peak usage are tracked and logged by <see cref="LogMemoryUsage"/>.
/// An <see cref="OutOfMemoryException"/> while loading is kept in <see cref="FatalError"/> and rethrown by <see cref="Acquire"/>.
/// </summary>
internal sealed class AudioClipStore : IDisposable
{
    /// <summary>Longest wait of the decoder thread between checks when nothing signals it, in milliseconds.</summary>
    private const int DecoderIdleWaitMs = 50;

    /// <summary>Longest wait of a consumer between checks when no chunk is loaded, in milliseconds.</summary>
    private const int ConsumerWaitMs = 50;

    /// <summary>Longest wait for the decoder thread to stop in <see cref="Dispose"/>, in milliseconds.</summary>
    private const int DecoderStopTimeoutMs = 3000;

    private readonly Dictionary<string, AudioClipSource?> _sources = new();
    private readonly List<AudioClipSource> _sourceList = new();
    private readonly Stack<float[]> _chunkPool = new();
    private readonly object _poolLock = new();
    private readonly object _signal = new();
    private readonly AutoResetEvent _serviceRequested = new(false);

    private Thread? _decoderThread;
    private volatile bool _stopRequested;
    private long _allocatedBytes;
    private long _peakBytes;
    private volatile Exception? _fatalError;
    private volatile bool _disposed;

    /// <summary>Sample rate every source is resampled to, in Hz.</summary>
    public int OutputRate { get; }

    /// <summary>Bytes of sample arrays currently allocated (loaded chunks, pooled chunk arrays and copied Unity clip data).</summary>
    public long AllocatedBytes => Interlocked.Read(ref _allocatedBytes);

    /// <summary>Highest value of <see cref="AllocatedBytes"/> since creation.</summary>
    public long PeakBytes => Interlocked.Read(ref _peakBytes);

    /// <summary>Out-of-memory error raised while loading, or null.</summary>
    public Exception? FatalError => _fatalError;

    public AudioClipStore(int outputRate)
    {
        if (outputRate <= 0) throw new ArgumentOutOfRangeException(nameof(outputRate));
        OutputRate = outputRate;
    }

    // ── Decoder thread ──

    /// <summary>Starts the decoder thread. Called once.</summary>
    public void Start()
    {
        if (_decoderThread != null) throw new InvalidOperationException("Decoder thread already started.");
        _decoderThread = new Thread(DecoderLoop)
        {
            Name = "Audio Clip Decoder",
            IsBackground = true
        };
        _decoderThread.Start();
    }

    private void DecoderLoop()
    {
        try
        {
            while (!_stopRequested)
            {
                if (!LoadNext())
                    _serviceRequested.WaitOne(DecoderIdleWaitMs);
            }
        }
        catch (Exception e)
        {
            Melon<RenderMod>.Logger.Error($"[AudioClipStore] Decoder thread stopped: {e}");
            _fatalError ??= e;
            NotifyChunkLoaded();
        }
    }

    /// <summary>
    /// Frees unused windowed chunks, then loads the single chunk with the lowest priority value over all sources.
    /// Runs on the decoder thread (tests may call it directly while the thread is not started).
    /// </summary>
    /// <returns>True when a chunk was loaded or a load was attempted.</returns>
    internal bool LoadNext()
    {
        AudioClipSource[] sources;
        lock (_sources)
        {
            if (_disposed) return false;
            sources = _sourceList.ToArray();
        }

        AudioClipSource? best = null;
        var bestIndex = -1;
        var bestPriority = int.MaxValue;
        foreach (var source in sources)
        {
            source.FreeUnused();
            if (!source.TryGetNextLoad(out var index, out var priority) || priority >= bestPriority) continue;
            best = source;
            bestIndex = index;
            bestPriority = priority;
        }

        if (best == null) return false;
        best.LoadChunk(bestIndex);
        return true;
    }

    internal void RequestService() => _serviceRequested.Set();

    // ── Sources ──

    /// <summary>
    /// Returns the source registered under <paramref name="key"/>, creating it with <paramref name="createReader"/> on first use.
    /// A key whose reader could not be created is remembered as null and not retried.
    /// For readers that can seek, additional decode streams (see <see cref="AudioClipSource.MaxDecodeStreams"/>) are opened on the
    /// decoder thread with <paramref name="openAdditionalReader"/>, or with <paramref name="createReader"/> when it is null.
    /// </summary>
    public AudioClipSource? GetOrAdd(string key, Func<IPCMReader?> createReader, Func<IPCMReader?>? openAdditionalReader = null)
    {
        lock (_sources)
        {
            if (_disposed) return null;
            if (_sources.TryGetValue(key, out var existing)) return existing;
        }

        AudioClipSource? source = null;
        try
        {
            var reader = createReader();
            if (reader != null)
            {
                try
                {
                    source = new AudioClipSource(this, key, reader, reader.IsForwardOnly ? null : openAdditionalReader ?? createReader, OutputRate);
                }
                catch
                {
                    reader.Dispose();
                    throw;
                }
            }
        }
        catch (OutOfMemoryException e)
        {
            Melon<RenderMod>.Logger.Error($"[AudioClipStore] Out of memory while opening clip '{key}' ({AllocatedBytes / 1048576.0:F1} MiB allocated): {e}");
            _fatalError ??= e;
        }
        catch (Exception e)
        {
            Melon<RenderMod>.Logger.Error($"[AudioClipStore] Failed to open clip '{key}': {e}");
        }

        lock (_sources)
        {
            if (_disposed || _sources.TryGetValue(key, out var raced))
            {
                source?.DisposeInternal();
                return _disposed ? null : _sources[key];
            }
            _sources[key] = source;
            if (source != null) _sourceList.Add(source);
        }

        if (source != null)
        {
            Melon<RenderMod>.Logger.Msg($"[AudioClipStore] Opened clip '{key}': {source.Channels}ch, {source.Duration:F3}s, {source.TotalFrames} frames at {OutputRate}Hz, {source.ChunkCount} chunk(s), {source.Policy}");
            RequestService();
        }
        return source;
    }

    public bool TryGet(string key, out AudioClipSource? source)
    {
        lock (_sources)
        {
            return _sources.TryGetValue(key, out source);
        }
    }

    // ── Consumers ──

    /// <summary>
    /// Holds the chunk containing <paramref name="frame"/>, waiting for the decoder thread until it is loaded.
    /// Returns an invalid view when the source failed, the frame is outside the source, or the store was disposed.
    /// </summary>
    /// <exception cref="InvalidOperationException">Loading ran out of memory or the decoder thread stopped.</exception>
    /// <exception cref="OperationCanceledException"><paramref name="cancellationToken"/> was canceled while waiting.</exception>
    public ChunkView Acquire(AudioClipSource source, long frame, CancellationToken cancellationToken = default)
    {
        if (frame < 0 || frame >= source.TotalFrames) return default;

        while (true)
        {
            ThrowIfFatal();
            if (source.IsFailed || _disposed) return default;
            if (source.TryAcquire(frame, out var view)) return view;

            lock (_signal)
            {
                if (source.TryAcquire(frame, out view)) return view;
                if (source.IsFailed || _fatalError != null || _disposed) continue;
                Monitor.Wait(_signal, ConsumerWaitMs);
            }
            cancellationToken.ThrowIfCancellationRequested();
        }
    }

    private void ThrowIfFatal()
    {
        var error = _fatalError;
        if (error != null)
            throw new InvalidOperationException("Audio clip loading failed: " + error.Message, error);
    }

    internal void NotifyChunkLoaded()
    {
        lock (_signal)
        {
            Monitor.PulseAll(_signal);
        }
    }

    internal void ReportSourceFailure(AudioClipSource source, Exception error)
    {
        Melon<RenderMod>.Logger.Error($"[AudioClipStore] Failed to load clip '{source.Key}', it will be silent: {error}");
        NotifyChunkLoaded();
    }

    internal void ReportFatal(AudioClipSource source, OutOfMemoryException error)
    {
        Melon<RenderMod>.Logger.Error($"[AudioClipStore] Out of memory while loading clip '{source.Key}' ({AllocatedBytes / 1048576.0:F1} MiB allocated): {error}");
        _fatalError ??= error;
        NotifyChunkLoaded();
    }

    // ── Arrays ──

    internal float[] RentChunkArray()
    {
        lock (_poolLock)
        {
            if (_chunkPool.Count > 0) return _chunkPool.Pop();
        }
        var data = new float[AudioClipSource.ChunkFloats];
        AddAllocated(data.LongLength * sizeof(float));
        return data;
    }

    internal void ReturnChunkArray(float[] data)
    {
        if (_disposed)
        {
            AddAllocated(-data.LongLength * sizeof(float));
            return;
        }
        lock (_poolLock)
        {
            _chunkPool.Push(data);
        }
    }

    internal float[] AllocateExactArray(int floats)
    {
        var data = new float[floats];
        AddAllocated(data.LongLength * sizeof(float));
        return data;
    }

    internal void FreeExactArray(float[] data) => AddAllocated(-data.LongLength * sizeof(float));

    private void AddAllocated(long bytes)
    {
        var current = Interlocked.Add(ref _allocatedBytes, bytes);
        long peak;
        while (current > (peak = Interlocked.Read(ref _peakBytes)))
        {
            if (Interlocked.CompareExchange(ref _peakBytes, current, peak) == peak) break;
        }
    }

    public void LogMemoryUsage()
    {
        int sourceCount;
        lock (_sources) sourceCount = _sourceList.Count;
        Melon<RenderMod>.Logger.Msg($"[AudioClipStore] Clips: {sourceCount}, sample memory now {AllocatedBytes / 1048576.0:F1} MiB, peak {PeakBytes / 1048576.0:F1} MiB");
    }

    /// <summary>Stops the decoder thread, then releases every source and pooled array.</summary>
    public void Dispose()
    {
        AudioClipSource[] sources;
        lock (_sources)
        {
            if (_disposed) return;
            _disposed = true;
            sources = _sourceList.ToArray();
        }

        _stopRequested = true;
        _serviceRequested.Set();
        NotifyChunkLoaded();

        // A decoder thread blocked inside a read (for example a file on an unresponsive drive) is left running:
        // it is a background thread, and its sources keep their readers so the blocked call does not touch freed FFmpeg contexts.
        var decoderStopped = _decoderThread == null || _decoderThread.Join(DecoderStopTimeoutMs);
        if (!decoderStopped)
            Melon<RenderMod>.Logger.Warning($"[AudioClipStore] Decoder thread did not stop within {DecoderStopTimeoutMs} ms; its readers are left open.");

        LogMemoryUsage();

        lock (_sources)
        {
            _sourceList.Clear();
            _sources.Clear();
        }

        if (decoderStopped)
        {
            foreach (var source in sources)
                source.DisposeInternal();
        }

        lock (_poolLock)
        {
            while (_chunkPool.Count > 0)
                AddAllocated(-_chunkPool.Pop().LongLength * sizeof(float));
        }
    }
}