Apply Optimized GPU Readback

Util/FrameReader.cs

@@ -1,7 +1,181 @@
-namespace ADOFAIRenderer.Util
+using System;
+using System.Collections.Generic;
+using ADOFAIRenderer.Component;
+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
+    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;
+        
+        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 (!RenderMod.IsCapturing) return;
+            ReadbackRequestState state;
+
+            lock (_poolLock)
+            {
+                if (_statePool.Count == 0)
+                {
+                    state = new ReadbackRequestState(this);
+                }
+                else
+                {
+                    state = _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;
+
+            if (!RenderMod.IsCapturing)
+            {
+                RecycleState(state);
+                return;
+            }
+
+            if (request.hasError)
+            {
+                RenderMod.ReadbackFailureCount++;
+                Melon<RenderMod>.Logger.Warning($"[OOM Prevent] GPU Readback failed for frame {frame}. Queue deadlock prevented by sending an empty frame.");
+
+                PushToQueue(frame, Array.Empty<byte>());
+            }
+            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)
+        {
+            state.Recycle();
+            lock (_poolLock)
+            {
+                _statePool.Push(state);
+            }
+        }
+        
+        private static void PushToQueue(long frame, byte[] rawBytes)
+        {
+            lock (RenderManager.ImageQueue)
+            {
+                RenderManager.ImageQueue[frame] = rawBytes;
+            }
+        }
+        
+        public void Dispose()
+        {
+            lock (_poolLock)
+            {
+                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();
+        }
     }
 }