BlockingReadback with AsyncGPUReadbackRequest

BlockingReadback.cs

/* 
*   NatRender
*   Copyright (c) 2019 Yusuf Olokoba
*/

namespace NatCorder.Readback
{

    using UnityEngine;
    using System;
    using System.Runtime.InteropServices;
    using Dispatch;
    using System.Collections.Generic;
    using UnityEngine.Rendering;
    using System.Collections;
    using System.Threading;
    using System.Threading.Tasks;

    public sealed class BlockingReadback : ReadbackContext
    {

        #region --Op vars--
        private readonly Texture2D framebuffer;
        private readonly IDispatcher readbackDispatcher, handlerDispatcher;
        private Queue<Action<IntPtr>> intQueue = new Queue<Action<IntPtr>>();
        private Queue<RenderTexture> texQueue = new Queue<RenderTexture>();
        private Action<IntPtr> finishAction;
        private int cacheSize;
        public int remain = 0;
        public bool finish = false;
        private bool inited = false;
        #endregion


        #region --Client API--

        public BlockingReadback(TextureFormat format, int cacheSize = 10)
        {
            framebuffer = new Texture2D(16, 16, format, false, false);
            readbackDispatcher = new MainDispatcher();
            handlerDispatcher = new WorkDispatcher();
            this.cacheSize = cacheSize;
        }

        public override RenderTexture Allocate(int width, int height)
        {
            if (!inited)
            {
                for (int i = 0; i < cacheSize; i++)
                {
                    var frameTexture = RenderTexture.GetTemporary(width, height, 24, RenderTextureFormat.Default, RenderTextureReadWrite.sRGB);
                    frameTexture.antiAliasing = 1;
                    var _ = frameTexture.colorBuffer;
                    texQueue.Enqueue(frameTexture);
                }
                inited = true;
            }
            if (texQueue.Count > 0)
            {
                return texQueue.Dequeue();
            }
            return null;
        }

        public override void Readback(RenderTexture frame, Action<IntPtr> handler)
        {
            // Null checking
            if (!frame)
            {
                // handlerDispatcher.Dispatch(() => handler(IntPtr.Zero));
                finishAction = handler;
                return;
            }
            // State checking
            if (framebuffer.width != frame.width || framebuffer.height != frame.height)
                framebuffer.Resize(frame.width, frame.height);
            remain++;
            AsyncGPUReadback.Request(frame, 0, TextureFormat.RGBA32, ReadbackCompleted);
            texQueue.Enqueue(frame);
            intQueue.Enqueue(handler);

            // Readback
            // var currentRT = RenderTexture.active;
            // RenderTexture.active = frame;
            // framebuffer.ReadPixels(new Rect(0, 0, framebuffer.width, framebuffer.height), 0, 0, false);
            // RenderTexture.active = currentRT;
            // RenderTexture.ReleaseTemporary(frame);
            // // Invoke handler
            // var data = framebuffer.GetRawTextureData();
            // handlerDispatcher.Dispatch(() =>
            // {
            //     var dataHandle = GCHandle.Alloc(data, GCHandleType.Pinned);
            //     handler(dataHandle.AddrOfPinnedObject());
            //     dataHandle.Free();
            //     GC.Collect();
            // });
        }

        private void ReadbackCompleted(AsyncGPUReadbackRequest request)
        {
            var image = request.GetData<Color32>();
            framebuffer.LoadRawTextureData(image);
            framebuffer.Apply();
            var data = framebuffer.GetRawTextureData();
            var handler = intQueue.Dequeue();
            handlerDispatcher.Dispatch(() =>
            {
                var dataHandle = GCHandle.Alloc(data, GCHandleType.Pinned);
                handler(dataHandle.AddrOfPinnedObject());
                dataHandle.Free();
                GC.Collect();
                remain--;
                if (finish && remain <= 0) this.DisposeMulti();
            });
        }

        //old version
        public override void Dispose()
        {
            // Readback(null, ptr =>
            // {
            //     using (var dispatcher = new MainDispatcher())
            //         dispatcher.Dispatch(() =>
            //         {
            //             Texture2D.Destroy(framebuffer);
            //             readbackDispatcher.Dispose();
            //             handlerDispatcher.Dispose();
            //         });
            // });
        }

        //new version
        private void DisposeMulti()
        {
            using (var dispatcher = new MainDispatcher())
                dispatcher.Dispatch(() =>
                {
                    Texture2D.Destroy(framebuffer);
                    readbackDispatcher.Dispose();
                    handlerDispatcher.Dispose();
                    finishAction(IntPtr.Zero);
                    foreach (var tex in texQueue)
                    {
                        tex.Release();
                    }
                });
        }

        #endregion
    }
}