ADOFAIRenderer/Util/FrameReader.cs

using System;
using System.Buffers;
using System.Collections.Generic;
using ADOFAIRenderer.Core;
using ADOFAIRenderer.Util.FFmpeg;
using MelonLoader;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;
using UnityEngine.Rendering;

namespace ADOFAIRenderer.Util
{
    public class FrameReader : IDisposable
    {
        private const int InitialBufferSize = 16;
        
        private readonly Stack<ReadbackRequestState> _statePool = new Stack<ReadbackRequestState>();
        private readonly object _poolLock = new object();
        private bool _isInitialized;
        private volatile bool _isDisposed;
        
        public FrameReader()
        {
            InitializePool();
        }

        private void InitializePool()
        {
            
            lock (_poolLock)
            {
                if (_isInitialized) return;
                for (var i = 0; i < InitialBufferSize; i++)
                {
                    _statePool.Push(new ReadbackRequestState(this));
                }

                _isInitialized = true;
            }
        }

        public void RequestFrameCapture(Texture texture, long currentFrame)
        {
            if (!RenderManager.IsCapturingActive) return;
            ReadbackRequestState state;

            lock (_poolLock)
            {
                state = _statePool.Count == 0 ? new ReadbackRequestState(this) : _statePool.Pop();
            }

            var bufferSize = texture.width * texture.height * 4; // RGBA = 4 channel
            state.Setup(currentFrame, bufferSize);

            AsyncGPUReadback.RequestIntoNativeArray(
                ref state.NativeBuffer,
                texture,
                0,
                state.CachedCallback
            );
        }

        internal void OnReadbackComplete(ReadbackRequestState state, AsyncGPUReadbackRequest request)
        {
            var frame = state.FrameIndex;

            lock (_poolLock)
            {
                if (_isDisposed)
                {
                    state.Dispose();
                    return;
                }
            }

            if (RenderManager.Instance && RenderManager.Instance.isRenderCancelled)
            {
                RecycleState(state);
                return;
            }

            if (request.hasError)
            {
                if (RenderManager.Instance)
                {
                    RenderManager.Instance.readbackFailureCount++;
                }
                Melon<RenderMod>.Logger.Warning($"[OOM Prevent] GPU Readback failed for frame {frame}. Queue deadlock prevented by sending a black frame.");

                var byteSize = state.NativeBuffer.Length;
                var blackFrame = ArrayPool<byte>.Shared.Rent(byteSize);
                Array.Clear(blackFrame, 0, byteSize);
                PushToQueue(frame, blackFrame);
            }
            else
            {
                unsafe
                {
                    var ptr = state.NativeBuffer.GetUnsafeReadOnlyPtr();
                    var rawBytes = FFmpegProcess.RentRawBytes(state.NativeBuffer.Length);

                    fixed (byte* destPtr = rawBytes)
                    {
                        Buffer.MemoryCopy(ptr, destPtr, rawBytes.Length, state.NativeBuffer.Length);
                    }

                    PushToQueue(frame, rawBytes);
                }
            }

            RecycleState(state);
        }

        private void RecycleState(ReadbackRequestState state)
        {
            lock (_poolLock)
            {
                if (_isDisposed || !_isInitialized)
                {
                    state.Dispose();
                    return;
                }
                
                state.Recycle();
                _statePool.Push(state);
            }
        }
        
        private static void PushToQueue(long frame, byte[] rawBytes)
        {
            if (!RenderManager.Instance) return;
            lock (RenderManager.Instance.ImageQueue)
            {
                RenderManager.Instance.ImageQueue[frame] = rawBytes;
            }
        }
        
        public void Dispose()
        {
            lock (_poolLock)
            {
                if (_isDisposed) return;
                _isDisposed = true;
                while (_statePool.Count > 0)
                {
                    var state = _statePool.Pop();
                    state.Dispose();
                }
                _isInitialized = false;
            }
        }
    }

    internal class ReadbackRequestState : IDisposable
    {
        private readonly FrameReader _owner;
        public long FrameIndex { get; private set; }
        public NativeArray<byte> NativeBuffer;

        public readonly Action<AsyncGPUReadbackRequest> CachedCallback;

        public ReadbackRequestState(FrameReader owner)
        {
            _owner = owner;
            CachedCallback = OnComplete;
            FrameIndex = -1;
        }

        public void Setup(long frameIndex, int bufferSize)
        {
            FrameIndex = frameIndex;
            
            // 버퍼가 생성되지 않았거나 크기가 다른 경우 재생성
            if (!NativeBuffer.IsCreated || NativeBuffer.Length != bufferSize)
            {
                // 이전 메모리 해제
                if (NativeBuffer.IsCreated) NativeBuffer.Dispose();

                NativeBuffer = new NativeArray<byte>(bufferSize, Allocator.Persistent,
                    NativeArrayOptions.UninitializedMemory);
            }
        }

        private void OnComplete(AsyncGPUReadbackRequest request)
        {
            _owner.OnReadbackComplete(this, request);
        }

        public void Recycle()
        {
            FrameIndex = -1;
        }

        public void Dispose()
        {
            if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
        }
    }
}