Multithreaded Unity Recorder package PNG image sequence recorder class

ImageRecorder.cs

using System;
using System.IO;
using System.Collections.Generic;
using UnityEditor.Recorder.Input;
using UnityEngine;
using UnityEngine.Profiling;
using UnityEngine.Rendering;
using System.Threading;
using System.Threading.Tasks;
using System.Collections.Concurrent;
using Unity.Collections;

namespace UnityEditor.Recorder
{
    class ImageRecorder : BaseTextureRecorder<ImageRecorderSettings>
    {
        Queue<string> m_PathQueue = new Queue<string>();

        class ImageWriteRequest
        {
            public string path;
            public byte[] data;
            public int w;
            public int h;
            public UnityEngine.Experimental.Rendering.GraphicsFormat fmt;
        }

        BlockingCollection<ImageWriteRequest> m_images = new BlockingCollection<ImageWriteRequest>(8);
        Task[] m_writers = new Task[8];
        private long _shouldShutdown = 0;
        CustomSampler encodeSampler = CustomSampler.Create("EncodeToPNG");
        CustomSampler fileioSampler = CustomSampler.Create("WriteToFile");
        public bool ShouldShutdown
        {
            get
            {
                /* Interlocked.Read() is only available for int64,
                 * so we have to represent the bool as a long with 0's and 1's
                 */
                return Interlocked.Read(ref _shouldShutdown) == 1;
            }
            set
            {
                Interlocked.Exchange(ref _shouldShutdown, Convert.ToInt64(value));
            }
        }

        protected override TextureFormat ReadbackTextureFormat
        {
            get
            {
                return Settings.OutputFormat != ImageRecorderSettings.ImageRecorderOutputFormat.EXR ? TextureFormat.RGBA32 : TextureFormat.RGBAFloat;
            }
        }

        protected internal override bool BeginRecording(RecordingSession session)
        {
            if (!base.BeginRecording(session)) { return false; }

            //SYNC
            UseAsyncGPUReadback = false;

            Settings.fileNameGenerator.CreateDirectory(session);
            ShouldShutdown = false;
            for (int i = 0; i < m_writers.Length; ++i)
            {
                m_writers[i] = new Task(() => {
                    int id = i;
                    WriterTask(id);
                });
            }

            foreach (var writer in m_writers)
            {
                writer.Start();
            }

            return true;
        }

        protected internal override void EndRecording(RecordingSession session)
        {
            base.EndRecording(session);
            ShouldShutdown = true;
        }

        protected internal override void RecordFrame(RecordingSession session)
        {
            if (m_Inputs.Count != 1)
                throw new Exception("Unsupported number of sources");
            // Store path name for this frame into a queue, as WriteFrame may be called
            // asynchronously later on, when the current frame is no longer the same (thus creating
            // a file name that isn't in sync with the session's current frame).
            m_PathQueue.Enqueue(Settings.fileNameGenerator.BuildAbsolutePath(session));
            base.RecordFrame(session);
        }

        protected override void WriteFrame(Texture2D tex)
        {
            Profiler.BeginSample("ImageRecorder.EnqueImageRequestSync");

            ImageWriteRequest req = new ImageWriteRequest()
            {
                data = tex.GetRawTextureData(),
                w = tex.width,
                h = tex.height,
                path = m_PathQueue.Dequeue(),
                fmt = tex.graphicsFormat
            };

            if (ShouldShutdown)
            {
                m_images.Add(req);
                m_images.CompleteAdding();
                Task.WaitAll(m_writers);
            }
            else
            {
                m_images.Add(req);
            }

            if (m_Inputs[0] is BaseRenderTextureInput || Settings.OutputFormat != ImageRecorderSettings.ImageRecorderOutputFormat.JPEG)
                UnityHelpers.Destroy(tex);

            Profiler.EndSample();
        }

        private void WriteToFile(byte[] bytes)
        {
            Profiler.BeginSample("ImageRecorder.WriteToFile");
            File.WriteAllBytes(m_PathQueue.Dequeue(), bytes);
            Profiler.EndSample();
        }

        protected override void WriteFrame(AsyncGPUReadbackRequest r)
        {
            Profiler.BeginSample("ImageRecorder.EnqueImageRequestAsync");

            ImageWriteRequest req = new ImageWriteRequest()
            {
                data = r.GetData<byte>().ToArray(),
                w = r.width,
                h = r.height,
                path = m_PathQueue.Dequeue(),
                fmt = UnityEngine.Experimental.Rendering.GraphicsFormat.R8G8B8A8_SRGB
            };

            m_images.Add(req);

            Profiler.EndSample();

            if (ShouldShutdown)
            {
                m_images.CompleteAdding();
                Task.WaitAll(m_writers);
            }
        }

        void WriterTask(int id)
        {
            Profiler.BeginThreadProfiling("ImageRecorder.Writer", "Writer " + id);

            Debug.Log("thread " + id + " STARTED");
            while (!m_images.IsCompleted)
            {

                ImageWriteRequest req = null;
                // Blocks if dataItems.Count == 0.
                // IOE means that Take() was called on a completed collection.
                // Some other thread can call CompleteAdding after we pass the
                // IsCompleted check but before we call Take.
                // In this example, we can simply catch the exception since the
                // loop will break on the next iteration.
                try
                {
                    req = m_images.Take();
                }
                catch (InvalidOperationException) { }

                if (req != null)
                {
                    if (req.data != null && req.path != String.Empty)
                    {
                        //Debug.Log("writing from " + id + " to " + req.path);
                        encodeSampler.Begin();
                        byte[] bytes = ImageConversion.EncodeArrayToPNG(req.data, req.fmt, (uint)req.w, (uint)req.h);
                        encodeSampler.End();
                        fileioSampler.Begin();
                        File.WriteAllBytes(req.path, bytes);
                        fileioSampler.End();
                    }
                }
            }
            Debug.Log("thread " + id + " FINISHED");
            Profiler.EndThreadProfiling();
        }

    }
}