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 Unity 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;

    /// <summary>Sources by key; a null value marks a key whose reader could not be created. Also the lock of the source tables.</summary>
    private readonly Dictionary<string, AudioClipSource?> _sources = new();

    /// <summary>Non-null sources in creation order, scanned by the decoder thread. Guarded by the <see cref="_sources"/> lock.</summary>
    private readonly List<AudioClipSource> _sourceList = new();

    /// <summary>Free windowed chunk arrays (<see cref="AudioClipSource.ChunkFloats"/> floats each), guarded by <see cref="_poolLock"/>.</summary>
    private readonly Stack<float[]> _chunkPool = new();

    private readonly object _poolLock = new();

    /// <summary>Monitor the decoder thread pulses after each load (or failure); consumers in <see cref="Acquire"/> wait on it.</summary>
    private readonly object _signal = new();

    /// <summary>Wakes the decoder thread when a cursor, demand or source changes (<see cref="RequestService"/>).</summary>
    private readonly AutoResetEvent _serviceRequested = new(false);

    private Thread? _decoderThread;
    private volatile bool _stopRequested;

    /// <summary>Bytes of allocated sample arrays (see <see cref="AllocatedBytes"/>), updated with Interlocked.</summary>
    private long _allocatedBytes;

    private long _peakBytes;

    /// <summary>First fatal error (out of memory, or an exception that stopped the decoder thread); null while none.</summary>
    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 opening or loading a clip, or the exception that stopped the decoder thread; null while none.
    /// </summary>
    public Exception? FatalError => _fatalError;

    /// <param name="outputRate">Sample rate in Hz every source is resampled to (positive). Call <see cref="Start"/> to begin decoding.</param>
    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();
    }

    /// <summary>
    /// Decoder thread body: loads one chunk per iteration, or sleeps until <see cref="RequestService"/> (at most
    /// <see cref="DecoderIdleWaitMs"/> ms) when nothing is missing. An unexpected exception ends the thread and is kept as the
    /// fatal error, so waiting consumers fail instead of waiting forever.
    /// </summary>
    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;
    }

    /// <summary>Wakes the decoder thread to re-evaluate what to free and load. Any thread.</summary>
    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}");
        }

        // The reader is created outside the lock (it may decode or copy a whole clip). If another caller registered the
        // same key meanwhile, its entry wins and this source is discarded.
        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;
    }

    /// <summary>
    /// True when <paramref name="key"/> has been registered; <paramref name="source"/> is then the source, or null when its reader
    /// could not be created. Any thread.
    /// </summary>
    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;

            // Retry under the signal lock so a load completed between the first try and the wait is not missed;
            // the timed wait also rechecks cancellation and failure states periodically.
            lock (_signal)
            {
                if (source.TryAcquire(frame, out view)) return view;
                if (source.IsFailed || _fatalError != null || _disposed) continue;
                Monitor.Wait(_signal, ConsumerWaitMs);
            }
            cancellationToken.ThrowIfCancellationRequested();
        }
    }

    /// <summary>Throws <see cref="InvalidOperationException"/> wrapping <see cref="FatalError"/> when one is set.</summary>
    private void ThrowIfFatal()
    {
        var error = _fatalError;
        if (error != null)
            throw new InvalidOperationException("Audio clip loading failed: " + error.Message, error);
    }

    /// <summary>Wakes every consumer waiting in <see cref="Acquire"/> (after a load, a failure, or on dispose).</summary>
    internal void NotifyChunkLoaded()
    {
        lock (_signal)
        {
            Monitor.PulseAll(_signal);
        }
    }

    /// <summary>Logs a load error of one source (the source is already marked failed and plays silence) and wakes waiting consumers.</summary>
    internal void ReportSourceFailure(AudioClipSource source, Exception error)
    {
        Melon<RenderMod>.Logger.Error($"[AudioClipStore] Failed to load clip '{source.Key}', it will be silent: {error}");
        NotifyChunkLoaded();
    }

    /// <summary>Logs an out-of-memory error, keeps it as <see cref="FatalError"/> (first one wins) and wakes waiting consumers.</summary>
    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 ──

    /// <summary>Takes a pooled array of <see cref="AudioClipSource.ChunkFloats"/> floats, or allocates (and counts) a new one. Contents are undefined.</summary>
    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;
    }

    /// <summary>Puts a rented array back into the pool (it stays counted); after dispose it is dropped and uncounted.</summary>
    internal void ReturnChunkArray(float[] data)
    {
        if (_disposed)
        {
            AddAllocated(-data.LongLength * sizeof(float));
            return;
        }
        lock (_poolLock)
        {
            _chunkPool.Push(data);
        }
    }

    /// <summary>Allocates a counted array of exactly <paramref name="floats"/> floats (resident chunks, copied Unity clip blocks).</summary>
    internal float[] AllocateExactArray(int floats)
    {
        var data = new float[floats];
        AddAllocated(data.LongLength * sizeof(float));
        return data;
    }

    /// <summary>Uncounts an array from <see cref="AllocateExactArray"/>; the array itself is left to the garbage collector.</summary>
    internal void FreeExactArray(float[] data) => AddAllocated(-data.LongLength * sizeof(float));

    /// <summary>Adds <paramref name="bytes"/> (negative to subtract) to the allocation counter and raises the peak lock-free.</summary>
    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;
        }
    }

    /// <summary>Logs the source count and the current and peak sample memory in MiB.</summary>
    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));
        }
    }
}