LightlessClient/LightlessSync/Services/PenumbraTempCollectionJanitor.cs

using System.Linq;
using LightlessSync.Interop.Ipc;
using LightlessSync.LightlessConfiguration.Models;
using LightlessSync.LightlessConfiguration;
using LightlessSync.Services.Mediator;
using Microsoft.Extensions.Logging;

namespace LightlessSync.Services;

public sealed class PenumbraTempCollectionJanitor : DisposableMediatorSubscriberBase
{
    private readonly IpcManager _ipc;
    private readonly TempCollectionConfigService _config;
    private readonly CancellationTokenSource _cleanupCts = new();
    private int _ran;
    private const int CleanupBatchSize = 50;
    private static readonly TimeSpan CleanupBatchDelay = TimeSpan.FromMilliseconds(50);
    private static readonly TimeSpan OrphanCleanupDelay = TimeSpan.FromDays(1);

    public PenumbraTempCollectionJanitor(
        ILogger<PenumbraTempCollectionJanitor> logger,
        LightlessMediator mediator,
        IpcManager ipc,
        TempCollectionConfigService config) : base(logger, mediator)
    {
        _ipc = ipc;
        _config = config;

        Mediator.Subscribe<PenumbraInitializedMessage>(this, _ => CleanupOrphansOnBoot());
    }

    public void Register(Guid id)
    {
        if (id == Guid.Empty) return;
        var changed = false;
        var config = _config.Current;

        var now = DateTime.UtcNow;
        var existing = config.OrphanableTempCollectionEntries.FirstOrDefault(entry => entry.Id == id);
        if (existing is null)
        {
            config.OrphanableTempCollectionEntries.Add(new OrphanableTempCollectionEntry
            {
                Id = id,
                RegisteredAtUtc = now
            });
            changed = true;
        }
        else if (existing.RegisteredAtUtc == DateTime.MinValue)
        {
            existing.RegisteredAtUtc = now;
            changed = true;
        }

        if (changed)
        {
            _config.Save();
        }
    }

    public void Unregister(Guid id)
    {
        if (id == Guid.Empty) return;
        var config = _config.Current;
        var changed = RemoveEntry(config.OrphanableTempCollectionEntries, id) > 0;
        if (changed)
        {
            _config.Save();
        }
    }

    private void CleanupOrphansOnBoot()
    {
        if (Interlocked.Exchange(ref _ran, 1) == 1)
            return;

        if (!_ipc.Penumbra.APIAvailable)
            return;

        _ = Task.Run(async () =>
        {
            try
            {
                await CleanupOrphansOnBootAsync(_cleanupCts.Token).ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
            }
            catch (Exception ex)
            {
                Logger.LogError(ex, "Error cleaning orphaned temp collections");
            }
        });
    }

    private async Task CleanupOrphansOnBootAsync(CancellationToken token)
    {
        var config = _config.Current;
        var entries = config.OrphanableTempCollectionEntries;
        if (entries.Count == 0)
            return;

        var now = DateTime.UtcNow;
        var changed = EnsureEntryTimes(entries, now);
        var cutoff = now - OrphanCleanupDelay;
        var expired = entries
            .Where(entry => entry.Id != Guid.Empty && entry.RegisteredAtUtc != DateTime.MinValue && entry.RegisteredAtUtc <= cutoff)
            .Select(entry => entry.Id)
            .Distinct()
            .ToList();
        if (expired.Count == 0)
        {
            if (changed)
            {
                _config.Save();
            }
            return;
        }

        var appId = Guid.NewGuid();
        Logger.LogInformation("Cleaning up {count} orphaned Lightless temp collections older than {delay}", expired.Count, OrphanCleanupDelay);

        List<Guid> removedIds = [];
        foreach (var id in expired)
        {
            if (token.IsCancellationRequested)
            {
                break;
            }

            try
            {
                await _ipc.Penumbra.RemoveTemporaryCollectionAsync(Logger, appId, id).ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                Logger.LogDebug(ex, "Failed removing orphaned temp collection {id}", id);
            }

            removedIds.Add(id);
            if (removedIds.Count % CleanupBatchSize == 0)
            {
                try
                {
                    await Task.Delay(CleanupBatchDelay, token).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
            }
        }

        if (removedIds.Count == 0)
        {
            if (changed)
            {
                _config.Save();
            }
            return;
        }

        foreach (var id in removedIds)
        {
            RemoveEntry(entries, id);
        }

        _config.Save();
    }

    protected override void Dispose(bool disposing)
    {
        if (disposing)
        {
            _cleanupCts.Cancel();
            _cleanupCts.Dispose();
        }

        base.Dispose(disposing);
    }

    private static int RemoveEntry(List<OrphanableTempCollectionEntry> entries, Guid id)
    {
        var removed = 0;
        for (var i = entries.Count - 1; i >= 0; i--)
        {
            if (entries[i].Id != id)
            {
                continue;
            }

            entries.RemoveAt(i);
            removed++;
        }

        return removed;
    }

    private static bool EnsureEntryTimes(List<OrphanableTempCollectionEntry> entries, DateTime now)
    {
        var changed = false;
        foreach (var entry in entries)
        {
            if (entry.Id == Guid.Empty || entry.RegisteredAtUtc != DateTime.MinValue)
            {
                continue;
            }

            entry.RegisteredAtUtc = now;
            changed = true;
        }

        return changed;
    }
}