ShrinkDataSaver/Runtime/LocalStorageProvider.cs

using System;
using System.Collections.Concurrent;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Cysharp.Threading.Tasks;

namespace ShrinkDataSaver
{
    public enum LocalStorageCommitStrategy
    {
        FileReplace = 0,
        RenameOnly = 1
    }

    public class LocalStorageProvider : IStorageProvider
    {
        private const string BackupPrimarySuffix = ".bak1";
        private const string BackupSecondarySuffix = ".bak2";
        private const string TempWriteSuffix = ".tmp";

        private static readonly ConcurrentDictionary<string, SemaphoreSlim> PathLocks =
            new(StringComparer.OrdinalIgnoreCase);

        private readonly string _rootPath;
        private readonly LocalStorageCommitStrategy _commitStrategy;
        private readonly Action<string, string> _overwriteMove;

        public LocalStorageProvider(string rootPath = null)
            : this(rootPath, LocalStorageCommitStrategy.FileReplace)
        {
        }

        public LocalStorageProvider(string rootPath, LocalStorageCommitStrategy commitStrategy)
            : this(rootPath, commitStrategy, null)
        {
        }

        public LocalStorageProvider(string rootPath, LocalStorageCommitStrategy commitStrategy,
            Action<string, string> overwriteMove)
        {
            if (!Enum.IsDefined(typeof(LocalStorageCommitStrategy), commitStrategy))
                throw new ArgumentOutOfRangeException(nameof(commitStrategy));
            if (commitStrategy == LocalStorageCommitStrategy.RenameOnly && overwriteMove == null)
                throw new ArgumentException("RenameOnly requires an overwrite-rename implementation.", nameof(overwriteMove));
            _overwriteMove = overwriteMove;
            _rootPath = rootPath ?? ShrinkDataSaverPlatform.Paths?.PersistentDataPath
#if UNITY_5_3_OR_NEWER
                ?? UnityEngine.Application.persistentDataPath
#endif
                ?? throw new InvalidOperationException("ShrinkDataSaver persistent path is not configured.");
            _commitStrategy = commitStrategy;
        }

        public LocalStorageCommitStrategy CommitStrategy => _commitStrategy;

        private string Resolve(string path) =>
            Path.IsPathRooted(path) ? path : Path.Combine(_rootPath, path);

        public async UniTask WriteAsync(string path, byte[] data, CancellationToken ct = default)
        {
            var fullPath = Resolve(path);
            var dir = Path.GetDirectoryName(fullPath);
            if (!string.IsNullOrEmpty(dir))
            {
                await Task.Run(() => Directory.CreateDirectory(dir), ct);
            }

            var pathLock = GetPathLock(fullPath);
            await pathLock.WaitAsync(ct);
            try
            {
                await Task.Run(() => RecoverInterruptedWrite(fullPath), CancellationToken.None);
                ct.ThrowIfCancellationRequested();

                var tempPath = fullPath + TempWriteSuffix;
                try
                {
                    await WriteFileBytesAsync(tempPath, data, ct);
                    ct.ThrowIfCancellationRequested();
                    await Task.Run(() => Commit(tempPath, fullPath), CancellationToken.None);
                }
                catch
                {
                    await Task.Run(() => TryDelete(tempPath), CancellationToken.None);
                    throw;
                }
            }
            finally
            {
                pathLock.Release();
            }
        }

        public async UniTask<byte[]> ReadAsync(string path, CancellationToken ct = default)
        {
            var fullPath = Resolve(path);
            var pathLock = GetPathLock(fullPath);
            await pathLock.WaitAsync(ct);
            try
            {
                await Task.Run(() => RecoverInterruptedWrite(fullPath), CancellationToken.None);
                if (!File.Exists(fullPath))
                {
                    throw new FileNotFoundException($"ShrinkDataSaver: file not found: {fullPath}");
                }

                return await ReadFileBytesAsync(fullPath, ct);
            }
            finally
            {
                pathLock.Release();
            }
        }

        public async UniTask<bool> ExistsAsync(string path, CancellationToken ct = default) =>
            await ExistsCoreAsync(Resolve(path), ct);

        private async UniTask<bool> ExistsCoreAsync(string fullPath, CancellationToken ct)
        {
            var pathLock = GetPathLock(fullPath);
            await pathLock.WaitAsync(ct);
            try
            {
                await Task.Run(() => RecoverInterruptedWrite(fullPath), CancellationToken.None);
                return File.Exists(fullPath);
            }
            finally
            {
                pathLock.Release();
            }
        }

        public async UniTask DeleteAsync(string path, CancellationToken ct = default)
        {
            var fullPath = Resolve(path);
            var pathLock = GetPathLock(fullPath);
            await pathLock.WaitAsync(ct);
            try
            {
                await Task.Run(() =>
                {
                    File.Delete(fullPath + TempWriteSuffix);
                    File.Delete(fullPath + BackupPrimarySuffix);
                    File.Delete(fullPath + BackupSecondarySuffix);
                    File.Delete(fullPath);
                }, ct);
            }
            finally
            {
                pathLock.Release();
            }
        }

        public async UniTask<string[]> ListAsync(string prefix = "", CancellationToken ct = default)
        {
            return await Task.Run(() =>
            {
                var dir = string.IsNullOrEmpty(prefix) ? _rootPath : Path.Combine(_rootPath, prefix);
                if (!Directory.Exists(dir))
                {
                    return Array.Empty<string>();
                }

                return Directory.GetFiles(dir)
                    .Select(f => Path.GetRelativePath(_rootPath, f))
                    .ToArray();
            }, ct);
        }

        private void Commit(string tempPath, string fullPath)
        {
            if (_commitStrategy == LocalStorageCommitStrategy.RenameOnly)
                CommitRenameOnly(tempPath, fullPath);
            else
                CommitFileReplace(tempPath, fullPath);
        }

        private static void CommitFileReplace(string tempPath, string fullPath)
        {
            var backupPath = fullPath + BackupPrimarySuffix;
            var secondaryBackupPath = fullPath + BackupSecondarySuffix;
            if (!File.Exists(fullPath))
            {
                File.Move(tempPath, fullPath);
                return;
            }

            if (File.Exists(secondaryBackupPath))
            {
                File.Delete(secondaryBackupPath);
            }

            if (File.Exists(backupPath))
            {
                File.Move(backupPath, secondaryBackupPath);
            }

            File.Replace(tempPath, fullPath, backupPath, true);
        }

        private void CommitRenameOnly(string tempPath, string fullPath)
        {
            var backupPath = fullPath + BackupPrimarySuffix;
            var secondaryBackupPath = fullPath + BackupSecondarySuffix;
            if (File.Exists(fullPath))
            {
                if (File.Exists(backupPath))
                {
                    if (File.Exists(secondaryBackupPath))
                    {
                        File.Delete(secondaryBackupPath);
                    }

                    MoveFileReplacing(backupPath, secondaryBackupPath);
                }

                MoveFileReplacing(fullPath, backupPath);
            }

            MoveFileReplacing(tempPath, fullPath);
        }

        private void RecoverInterruptedWrite(string fullPath)
        {
            if (_commitStrategy != LocalStorageCommitStrategy.RenameOnly)
                return;

            var tempPath = fullPath + TempWriteSuffix;
            var primaryBackupPath = fullPath + BackupPrimarySuffix;
            var secondaryBackupPath = fullPath + BackupSecondarySuffix;
            if (!File.Exists(fullPath))
            {
                if (File.Exists(primaryBackupPath))
                {
                    MoveFileReplacing(primaryBackupPath, fullPath);
                }
                else if (File.Exists(secondaryBackupPath))
                {
                    MoveFileReplacing(secondaryBackupPath, fullPath);
                }
            }

            if (File.Exists(tempPath))
            {
                TryDelete(tempPath);
            }
        }

        private void MoveFileReplacing(string source, string destination) => _overwriteMove(source, destination);

        private static void TryDelete(string path)
        {
            try
            {
                if (File.Exists(path))
                {
                    File.Delete(path);
                }
            }
            catch (IOException)
            {
            }
            catch (UnauthorizedAccessException)
            {
            }
        }

        private static SemaphoreSlim GetPathLock(string fullPath)
        {
            return PathLocks.GetOrAdd(fullPath, _ => new SemaphoreSlim(1, 1));
        }

        private static async UniTask WriteFileBytesAsync(string fullPath, byte[] data, CancellationToken ct)
        {
            await using var stream = new FileStream(
                fullPath,
                FileMode.Create,
                FileAccess.Write,
                FileShare.None,
                4096,
                FileOptions.Asynchronous);
            await stream.WriteAsync(data, 0, data.Length, ct);
            stream.Flush(true);
        }

        private static async UniTask<byte[]> ReadFileBytesAsync(string fullPath, CancellationToken ct)
        {
            await using var stream = new FileStream(
                fullPath,
                FileMode.Open,
                FileAccess.Read,
                FileShare.Read,
                4096,
                FileOptions.Asynchronous | FileOptions.SequentialScan);

            var length = stream.Length;
            if (length > int.MaxValue)
            {
                throw new IOException($"ShrinkDataSaver: file too large to read into memory: {fullPath}");
            }

            var buffer = new byte[length];
            var offset = 0;
            while (offset < buffer.Length)
            {
                var read = await stream.ReadAsync(buffer, offset, buffer.Length - offset, ct);
                if (read <= 0)
                {
                    throw new EndOfStreamException($"ShrinkDataSaver: unexpected EOF while reading {fullPath}");
                }

                offset += read;
            }

            return buffer;
        }
    }
}