LightlessClient/LightlessSync/Services/Compression/BatchFileFragService.cs

using Microsoft.Extensions.Logging;
using System.Diagnostics;
using System.Text.RegularExpressions;
using System.Threading.Channels;

namespace LightlessSync.Services.Compression
{
    /// <summary>
    /// This batch service is made for the File Frag command, because of each file needing to use this command.
    /// It's better to combine into one big command in batches then doing each command on each compressed call.
    /// </summary>
    public sealed partial class BatchFilefragService : IDisposable
    {
        private readonly Channel<(string path, TaskCompletionSource<bool> tcs)> _ch;
        private readonly Task _worker;
        private readonly bool _useShell;
        private readonly ILogger _log;
        private readonly int _batchSize;
        private readonly TimeSpan _flushDelay;
        private readonly CancellationTokenSource _cts = new();

        public BatchFilefragService(bool useShell, ILogger log, int batchSize = 128, int flushMs = 25)
        {
            _useShell = useShell;
            _log = log;
            _batchSize = Math.Max(8, batchSize);
            _flushDelay = TimeSpan.FromMilliseconds(Math.Max(5, flushMs));
            _ch = Channel.CreateUnbounded<(string, TaskCompletionSource<bool>)>(new UnboundedChannelOptions { SingleReader = true, SingleWriter = false });
            _worker = Task.Run(ProcessAsync, _cts.Token);
        }

        /// <summary>
        /// Checks if the file is compressed using Btrfs using tasks
        /// </summary>
        /// <param name="linuxPath">Linux/Wine path given for the file.</param>
        /// <param name="ct">Cancellation Token</param>
        /// <returns>If it was compressed or not</returns>
        public Task<bool> IsCompressedAsync(string linuxPath, CancellationToken ct = default)
        {
            var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);

            if (!_ch.Writer.TryWrite((linuxPath, tcs)))
            {
                tcs.TrySetResult(false);
                return tcs.Task;
            }

            if (ct.CanBeCanceled)
            {
                var reg = ct.Register(() => tcs.TrySetCanceled(ct));
                _ = tcs.Task.ContinueWith(_ => reg.Dispose(), CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
            }

            return tcs.Task;
        }

        /// <summary>
        /// Process the pending compression tasks asynchronously
        /// </summary>
        /// <returns>Task</returns>
        private async Task ProcessAsync()
        {
            var reader = _ch.Reader;
            var pending = new List<(string path, TaskCompletionSource<bool> tcs)>(_batchSize);

            try
            {
                while (await reader.WaitToReadAsync(_cts.Token).ConfigureAwait(false))
                {
                    if (!reader.TryRead(out var first)) continue;
                    pending.Add(first);

                    var flushAt = DateTime.UtcNow + _flushDelay;
                    while (pending.Count < _batchSize && DateTime.UtcNow < flushAt)
                    {
                        if (reader.TryRead(out var item))
                        {
                            pending.Add(item);
                            continue;
                        }

                        if ((flushAt - DateTime.UtcNow) <= TimeSpan.Zero) break;
                        try 
                        { 
                            await Task.Delay(TimeSpan.FromMilliseconds(5), _cts.Token).ConfigureAwait(false); 
                        }
                        catch 
                        { 
                            break; 
                        }
                    }

                    try
                    {
                        var map = await RunBatchAsync(pending.Select(p => p.path)).ConfigureAwait(false);
                        foreach (var (path, tcs) in pending)
                        {
                            tcs.TrySetResult(map.TryGetValue(path, out var c) && c);
                        }
                    }
                    catch (Exception ex)
                    {
                        _log.LogDebug(ex, "filefrag batch failed. falling back to false");
                        foreach (var (_, tcs) in pending)
                        {
                            tcs.TrySetResult(false);
                        }
                    }
                    finally
                    {
                        pending.Clear();
                    }
                }
            }
            catch (OperationCanceledException) 
            { 
                //Shutting down worker, exception called
            }
        }

        /// <summary>
        /// Running the batch of each file in the queue in one file frag command.
        /// </summary>
        /// <param name="paths">Paths that are needed for the command building for the batch return</param>
        /// <returns>Path of the file and if it went correctly</returns>
        /// <exception cref="InvalidOperationException">Failing to start filefrag on the system if this exception is found</exception>
        private async Task<Dictionary<string, bool>> RunBatchAsync(IEnumerable<string> paths)
        {
            var list = paths.Distinct(StringComparer.Ordinal).ToList();
            var result = list.ToDictionary(p => p, _ => false, StringComparer.Ordinal);

            ProcessStartInfo psi;
            if (_useShell)
            {
                var inner = "filefrag -v -- " + string.Join(' ', list.Select(QuoteSingle));
                psi = new ProcessStartInfo
                {
                    FileName = "/bin/bash",
                    Arguments = "-c " + QuoteDouble(inner),
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true,
                    WorkingDirectory = "/"
                };
            }
            else
            {
                psi = new ProcessStartInfo
                {
                    FileName = "filefrag",
                    RedirectStandardOutput = true,
                    RedirectStandardError = true,
                    UseShellExecute = false,
                    CreateNoWindow = true
                };
                psi.ArgumentList.Add("-v");
                psi.ArgumentList.Add("--");
                foreach (var p in list) psi.ArgumentList.Add(p);
            }

            using var proc = Process.Start(psi) ?? throw new InvalidOperationException("Failed to start filefrag");
            var stdoutTask = proc.StandardOutput.ReadToEndAsync(_cts.Token);
            var stderrTask = proc.StandardError.ReadToEndAsync(_cts.Token);
            await Task.WhenAll(stdoutTask, stderrTask).ConfigureAwait(false);
            try 
            { 
                await proc.WaitForExitAsync(_cts.Token).ConfigureAwait(false); 
            } 
            catch (Exception ex)
            {
                _log.LogWarning(ex, "Error in the batch frag service. proc = {proc}", proc);
            }

            var stdout = await stdoutTask.ConfigureAwait(false);
            var stderr = await stderrTask.ConfigureAwait(false);

            if (proc.ExitCode != 0 && !string.IsNullOrWhiteSpace(stderr))
                _log.LogTrace("filefrag exited {code}: {err}", proc.ExitCode, stderr.Trim());

            ParseFilefrag(stdout, result);
            return result;
        }

        /// <summary>
        /// Parsing the string given from the File Frag command into mapping
        /// </summary>
        /// <param name="output">Output of the process from the File Frag</param>
        /// <param name="map">Mapping of the processed files</param>
        private static void ParseFilefrag(string output, Dictionary<string, bool> map)
        {
            var reHeaderColon = ColonRegex();
            var reHeaderSize = SizeRegex();

            string? current = null;
            using var sr = new StringReader(output);
            for (string? line = sr.ReadLine(); line != null; line = sr.ReadLine())
            {
                var m1 = reHeaderColon.Match(line);
                if (m1.Success) { current = m1.Groups[1].Value; continue; }

                var m2 = reHeaderSize.Match(line);
                if (m2.Success) { current = m2.Groups[1].Value; continue; }

                if (current is not null && line.Contains("flags:", StringComparison.OrdinalIgnoreCase) &&
                    line.Contains("compressed", StringComparison.OrdinalIgnoreCase) && map.ContainsKey(current))
                {
                    map[current] = true;
                }
            }
        }

        private static string QuoteSingle(string s) => "'" + s.Replace("'", "'\\''", StringComparison.Ordinal) + "'";
        private static string QuoteDouble(string s) => "\"" + s.Replace("\\", "\\\\", StringComparison.Ordinal).Replace("\"", "\\\"", StringComparison.Ordinal).Replace("$", "\\$", StringComparison.Ordinal).Replace("`", "\\`", StringComparison.Ordinal) + "\"";

        /// <summary>
        /// Regex of the File Size return on the Linux/Wine systems, giving back the amount
        /// </summary>
        /// <returns>Regex of the File Size</returns>
        [GeneratedRegex(@"^File size of (/.+?) is ", RegexOptions.ExplicitCapture | RegexOptions.CultureInvariant,matchTimeoutMilliseconds: 500)]
        private static partial Regex SizeRegex();

        /// <summary>
        /// Regex on colons return on the Linux/Wine systems
        /// </summary>
        /// <returns>Regex of the colons in the given path</returns>
        [GeneratedRegex(@"^(/.+?):\s", RegexOptions.ExplicitCapture | RegexOptions.CultureInvariant, matchTimeoutMilliseconds: 500)]
        private static partial Regex ColonRegex();

        public void Dispose()
        {
            _ch.Writer.TryComplete();
            _cts.Cancel();
            try
            {
                _worker.Wait(TimeSpan.FromSeconds(2), _cts.Token);
            }
            catch
            {
                // Ignore the catch in dispose
            }
            _cts.Dispose();
        }
    }
}