非同期で行うスクリーンショットの取得・永続化・読み出し

AsyncTexture.cs

public class AsyncTexture
{
    public async UniTask Save(Texture2D source, string path, CancellationToken cancellationToken)
    {
        var data = source.GetRawTextureData<byte>();
        var png = ImageConversion.EncodeNativeArrayToPNG(data,source.graphicsFormat,
            (uint)source.width,(uint)source.height,0);

        var memory = default(Memory<byte>);
        unsafe
        {
            var mgr = new NativeArrayManager<byte>((byte*)png.GetUnsafePtr(), png.Length);
            memory = mgr.Memory;
        }

        var stream = File.OpenWrite(path);
        await stream.WriteAsync(memory,cancellationToken);
        stream.Dispose();
        
        png.Dispose();
    }
    public async UniTask<Texture2D> Load(string path, CancellationToken cancellationToken)
    {
        var uri =$"file://{path}";
        using var request = UnityWebRequestTexture.GetTexture(uri);
        await request.SendWebRequest();
        return DownloadHandlerTexture.GetContent(request);
    }
    
}

LocalTexture.cs

/// <summary>
/// 非同期的に読み込めるテクスチャ
/// </summary>
public class LocalTexture
{
    #region save
    
    public async UniTask Save(Texture2D source, string path, CancellationToken cancellationToken)
    {
        // 方針
        // Texture2DのRawデータ読み込みが使えるようにGetRawTextureDataを取得して吐き出す
        // それ以外だと、各テクスチャフォーマットのデコーダーを非同期で書く必要があるので、それは避けたい

        var rawTextureData = source.GetRawTextureData<byte>();
        var headerSize = Header.Size;
        var outputSize = headerSize + rawTextureData.Length;
        var outputData = ArrayPool<byte>.Shared.Rent(outputSize);

        CreateOutputData(source.width, source.height, source.format, rawTextureData, outputData);

        //poolのRentはジャストサイズ以上を返す場合があるらしいので保存長を指定する
        //https://ikorin2.hatenablog.jp/entry/2020/07/25/113904
        using var stream = File.OpenWrite(path);
        await stream.WriteAsync(outputData, 0, outputSize, cancellationToken);
        
        ArrayPool<byte>.Shared.Return(outputData);
    }

    private void CreateOutputData(int width, int height, TextureFormat format, NativeArray<byte> rawTextureData,
        byte[] output)
    {
        Span<byte> widthData = stackalloc byte[sizeof(int)];
        Span<byte> heightData = stackalloc byte[sizeof(int)];
        Span<byte> formatData = stackalloc byte[sizeof(int)];

        BitConverter.TryWriteBytes(widthData, width);
        BitConverter.TryWriteBytes(heightData, height);
        BitConverter.TryWriteBytes(formatData, (int)format);

        var headerSize = Header.Size;
        var outputSize = headerSize + rawTextureData.Length;

        if (output.Length < outputSize)
        {
            throw new ArgumentException("output data is too small");
        }

        var outputDataSpan = output.AsSpan();
        widthData.CopyTo(outputDataSpan);
        heightData.CopyTo(outputDataSpan.Slice(sizeof(int)));
        formatData.CopyTo(outputDataSpan.Slice(sizeof(int) * 2));

        unsafe
        {
            Span<byte> dataSpan = new Span<byte>(rawTextureData.GetUnsafePtr(), rawTextureData.Length);
            dataSpan.CopyTo(outputDataSpan.Slice(headerSize));
        }
    }

    #endregion

    #region load

    public async UniTask<Texture2D> Load(string path, CancellationToken cancellationToken)
    {
        if (!File.Exists(path))
        {
            Debug.LogError("File Not Found");
            return null;
        }
        
        var fileInfo = new FileInfo(path);
        var fileSize = fileInfo.Length;
        if (fileSize <= 0)
        {

            Debug.LogError("File Size is 0");
            return null;
        }

        ReadHandle readHandle;
        //非同期読み込みが4フレーム以内に終わる保証はないのでPersistent
        var readCmds = new NativeArray<ReadCommand>(1, Allocator.Persistent);
        unsafe
        {
            var cmd = new ReadCommand();

            cmd.Offset = 0;
            cmd.Size = fileSize;
            cmd.Buffer =
                (byte*)UnsafeUtility.Malloc(cmd.Size, UnsafeUtility.AlignOf<byte>(), Allocator.Persistent);

            readCmds[0] = cmd;

            readHandle = AsyncReadManager.Read(path, (ReadCommand*)readCmds.GetUnsafePtr(), (uint)readCmds.Length);
        }

        var awaitable = new ReadHandleAwaitable(readHandle,cancellationToken);
        var isSuccess = await awaitable;

        awaitable.Dispose();

        if (!isSuccess)
        {
            Debug.LogError("Read Failed");

            unsafe
            {
                UnsafeUtility.Free(readCmds[0].Buffer, Allocator.Persistent);
            }

            readCmds.Dispose();
            readHandle.Dispose();

            return null;
        }

        var pBuffer = IntPtr.Zero;
        unsafe
        {
            pBuffer = (IntPtr)readCmds[0].Buffer;
        }

        var header = ReadHeader(pBuffer);
        var texture = new Texture2D(header.Width, header.Height, TextureFormat.ARGB32, false, true);
        var dataSize = fileSize - Header.Size;

        //header分抜く
        pBuffer = IntPtr.Add(pBuffer, Header.Size);

        //texture.LoadImage()だと重くてスパイクするので、直接Rawデータを読む
        texture.LoadRawTextureData(pBuffer, (int)dataSize);
        texture.Apply();

        unsafe
        {
            UnsafeUtility.Free(readCmds[0].Buffer, Allocator.Persistent);
        }

        readCmds.Dispose();
        readHandle.Dispose();

        return texture;
    }

    private struct Header
    {
        public int Width { get; set; }
        public int Height { get; set; }
        public TextureFormat Format { get; set; }

        public static int Size => sizeof(int) * 3;
    }

    private Header ReadHeader(IntPtr pData)
    {
        var header = new Header();
        var buffer = ArrayPool<byte>.Shared.Rent(Header.Size);
        Marshal.Copy(pData, buffer, 0, Header.Size);

        var width = BitConverter.ToInt32(buffer, 0);
        var height = BitConverter.ToInt32(buffer, sizeof(int));
        var format = (TextureFormat)BitConverter.ToInt32(buffer, sizeof(int) * 2);

        ArrayPool<byte>.Shared.Return(buffer);

        header.Width = width;
        header.Height = height;
        header.Format = format;

        return header;
    }

    private class ReadHandleAwaitable : IDisposable
    {
        private ReadHandle _handle;
        private TaskCompletionSource<bool> _taskCompletionSource;

        private bool _isDisposed;
        private bool _isSuccess;

        private readonly uint _waitForMilliseconds;
        
        private readonly object _lockObject = new object();
        private CancellationToken _cancellationToken;

        public ReadHandleAwaitable(ReadHandle handle, CancellationToken cancellationToken,uint waitForMilliseconds = 32)
        {
            _handle = handle;
            _waitForMilliseconds = waitForMilliseconds;
            _cancellationToken = cancellationToken;
        }

        ~ReadHandleAwaitable()
        {
            Dispose();
        }

        private void WaitReadComplete(object state)
        {
            while (true)
            {
                lock (_lockObject)
                {
                    if (_isDisposed)
                    {
                        return;
                    }
                }
                
                if (_cancellationToken.IsCancellationRequested)
                {
                    _taskCompletionSource?.TrySetCanceled();
                    return;
                }

                var isProgress = _handle.Status == ReadStatus.InProgress;
                if (isProgress)
                {
                    Thread.Sleep((int)_waitForMilliseconds);
                    continue;
                }

                lock (_lockObject)
                {
                    if (_isDisposed)
                    {
                        return;
                    }
                }
                
                if (_cancellationToken.IsCancellationRequested)
                {
                    _taskCompletionSource?.TrySetCanceled();
                    return;
                }

                _isSuccess = _handle.Status == ReadStatus.Complete;
                break;
            }

            _taskCompletionSource?.TrySetResult(_isSuccess);
        }

        public TaskAwaiter<bool> GetAwaiter()
        {
            _taskCompletionSource = new TaskCompletionSource<bool>();
            ThreadPool.QueueUserWorkItem(WaitReadComplete);

            return _taskCompletionSource.Task.GetAwaiter();
        }

        public void Dispose()
        {
            lock (_lockObject)
            {
                if (_isDisposed)
                {
                    return;
                }

                _isDisposed = true;
            }
        }
    }

    #endregion
}

ScenarioTest.cs

using System.Collections;
using System.IO;
using System.Threading;
using Cysharp.Threading.Tasks;
using NUnit.Framework;
using Shared.ScreenShot;
using Unity.PerformanceTesting;
using UnityEngine;
using UnityEngine.TestTools;

namespace AsyncSample.Test
{
    public class ScenarioTest
    {
        [SetUp]
        public void Setup()
        {
            var path = Path.Combine(Application.persistentDataPath, "sampleRaw.png");
            var hasLoadTex = File.Exists(path);
            if (!hasLoadTex)
            {
                var sourcePath = Path.Combine(Application.persistentDataPath, "sample.png");
                var loadTex = new Texture2D(1,1);
                loadTex.LoadImage(File.ReadAllBytes(sourcePath));
                
                var loader = new LocalTexture();
                loader.Save(loadTex, path,CancellationToken.None);
            }
        }
        
        [UnityTest, Performance,Category("Load")]
        public IEnumerator TextureLoadWebRequest() => UniTask.ToCoroutine(async () =>
        {
            var path = Path.Combine(Application.persistentDataPath, "sample.png");
            var sampleGroup = new SampleGroup("TextureLoadWebRequest", SampleUnit.Millisecond);

            var asyncTexture = new AsyncTexture();
            
            for (var i = 0; i < 1000; i++)
            {
                var loadTex = default(Texture2D);
                using (Measure.Scope(sampleGroup))
                {
                    await asyncTexture.Load(path, CancellationToken.None);
                }
                MonoBehaviour.DestroyImmediate(loadTex);
            }
        });
        
        [UnityTest, Performance,Category("Load")]
        public IEnumerator TextureLoadRaw() => UniTask.ToCoroutine(async () =>
        {
            var path = Path.Combine(Application.persistentDataPath, "sampleRaw.png");
            var loader = new LocalTexture();
            
            var sampleGroup = new SampleGroup("TextureLoadRaw", SampleUnit.Millisecond);
            for (var i = 0; i < 1000; i++)
            {
                var loadTex = default(Texture2D);
                using (Measure.Scope(sampleGroup))
                {
                    loadTex = await loader.Load(path, CancellationToken.None);
                }
                MonoBehaviour.DestroyImmediate(loadTex);
            }
        });
        
        [UnityTest, Performance,Category("Load")]
        public IEnumerator TextureLoadSync() => UniTask.ToCoroutine(async () =>
        {
            var sourcePath = Path.Combine(Application.persistentDataPath, "sample.png");

            var sampleGroup = new SampleGroup("TextureLoadSync", SampleUnit.Millisecond);
            for (var i = 0; i < 1000; i++)
            {
                var loadTex = default(Texture2D);
                using (Measure.Scope(sampleGroup))
                { 
                    loadTex = new Texture2D(1,1);
                    loadTex.LoadImage(File.ReadAllBytes(sourcePath));
                }
                MonoBehaviour.DestroyImmediate(loadTex);
            }
        });
        
        [UnityTest, Performance,Category("Save")]
        public IEnumerator TextureSaveSync() => UniTask.ToCoroutine(async () =>
        {
            var sourcePath = Path.Combine(Application.persistentDataPath, "sample.png");
            var loadTex = new Texture2D(1,1);
            loadTex.LoadImage(File.ReadAllBytes(sourcePath));
            
            var path = Path.Combine(Application.persistentDataPath, "reSaveSample.png");

            var sampleGroup = new SampleGroup("TextureSaveSync", SampleUnit.Millisecond);
            for (var i = 0; i < 1000; i++)
            {
                using (Measure.Scope(sampleGroup))
                { 
                    var test = loadTex.EncodeToPNG();
                    File.WriteAllBytes(path,test);
                }
            }
        });
        
        [UnityTest, Performance,Category("Save")]
        public IEnumerator TextureSaveRaw() => UniTask.ToCoroutine(async () =>
        {
            var sourcePath = Path.Combine(Application.persistentDataPath, "sample.png");
            var loadTex = new Texture2D(1,1);
            loadTex.LoadImage(File.ReadAllBytes(sourcePath));
            
            var path = Path.Combine(Application.persistentDataPath, "reSaveSampleRaw.png");
            var loader = new LocalTexture();

            var sampleGroup = new SampleGroup("TextureSaveRaw", SampleUnit.Millisecond);
            for (var i = 0; i < 1000; i++)
            {
                using (Measure.Scope(sampleGroup))
                { 
                    await loader.Save(loadTex,path,CancellationToken.None);
                }
            }
        });
        
        [UnityTest, Performance,Category("Save")]
        public IEnumerator TextureSaveAsync() => UniTask.ToCoroutine(async () =>
        {
            var sourcePath = Path.Combine(Application.persistentDataPath, "sample.png");
            var loadTex = new Texture2D(1,1);
            loadTex.LoadImage(File.ReadAllBytes(sourcePath));
            
            var path = Path.Combine(Application.persistentDataPath, "reSaveSampleAsync.png");
            var asyncTexture = new AsyncTexture();
            var sampleGroup = new SampleGroup("TextureSaveAsync", SampleUnit.Millisecond);
            for (var i = 0; i < 1000; i++)
            {
                using (Measure.Scope(sampleGroup))
                {
                    await asyncTexture.Save(loadTex, path, CancellationToken.None);
                }
            }
        });
    }
}

ScreenShotPublisher.cs

/// <summary>
/// スクリーンショットを非同期で発行する
/// </summary>
public class ScreenShotPublisher
{
    private Camera _target;
    private int _width;
    private int _height;
    
    public ScreenShotPublisher(Camera target,int width,int height)
    {
        _target = target;
        _width = width;
        _height = height;
        
        //手動描画の場合はカメラのenabledをfalseにしておく
        target.enabled = false;
        target.aspect = (float)width / height;
        target.ResetAspect();
    }
    
    /// <summary>
    /// スクショ
    /// </summary>
    public async UniTask<Texture2D> Publish(CancellationToken cancellationToken)
    {
        var rt = RenderTexture.GetTemporary(_width, _height, 24, RenderTextureFormat.ARGB32);
        _target.enabled = true;
        _target.targetTexture = rt;
        _target.Render();
        
        await UniTask.WaitForEndOfFrame(cancellationToken);
        
        var screenShot = await ReadPixels(rt);
        
        _target.targetTexture = null;
        _target.enabled = false;
        RenderTexture.ReleaseTemporary(rt);
        
        return screenShot;
    }
    
    private async UniTask<Texture2D> ReadPixels(RenderTexture source)
    {
        
        Texture2D SyncReadPixels()
        {
            var currentRT = RenderTexture.active;
            RenderTexture.active = source;
            
            Texture2D screenShot = new Texture2D(_width, _height, TextureFormat.ARGB32, false);
            screenShot.ReadPixels(new Rect(0, 0, _width, _height), 0, 0);
            screenShot.Apply();
            
            RenderTexture.active = currentRT;
            return screenShot;
        }

        //ReadPixelsが重いので非同期に逃がせる場合逃す
        //参考 https://light11.hatenadiary.com/entry/2021/01/18/203431
        if (!SystemInfo.supportsAsyncGPUReadback)
        {
            return SyncReadPixels();
        }
        
        var request = await AsyncGPUReadback.Request(source, 0, TextureFormat.ARGB32);
        if (request.hasError)
        {
            Debug.LogError("AsyncGPUReadback has error use SyncReadPixels");
            return SyncReadPixels();
        }
        
        var data = request.GetData<Color32>();
        var screenShot = new Texture2D(source.width, source.height, TextureFormat.ARGB32, false);
        screenShot.LoadRawTextureData(data);
        screenShot.Apply();
        
        return screenShot;
    }
}