DmbFactory.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Api.Models;
using Tgstation.Server.Common.Extensions;
using Tgstation.Server.Host.Components.Deployment.Remote;
using Tgstation.Server.Host.Components.Events;
using Tgstation.Server.Host.Database;
using Tgstation.Server.Host.IO;
using Tgstation.Server.Host.Models;
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host.Components.Deployment
{
    sealed class DmbFactory : IDmbFactory, ICompileJobSink
    {
        public Task OnNewerDmb
        {
            get
            {
                lock (jobLockCounts)
                    return newerDmbTcs.Task;
            }
        }
 
        [MemberNotNullWhen(true, nameof(nextDmbProvider))]
        public bool DmbAvailable => nextDmbProvider != null;
 
        readonly IDatabaseContextFactory databaseContextFactory;
 
        readonly IIOManager ioManager;
 
        readonly IRemoteDeploymentManagerFactory remoteDeploymentManagerFactory;
 
        readonly ILogger<DmbFactory> logger;
 
        readonly IEventConsumer eventConsumer;
 
        readonly Api.Models.Instance metadata;
 
        readonly CancellationTokenSource cleanupCts;
 
        readonly Dictionary<long, int> jobLockCounts;
 
        volatile TaskCompletionSource newerDmbTcs;
 
        Task cleanupTask;
 
        IDmbProvider? nextDmbProvider;
 
        bool started;
 
        public DmbFactory(
            IDatabaseContextFactory databaseContextFactory,
            IIOManager ioManager,
            IRemoteDeploymentManagerFactory remoteDeploymentManagerFactory,
            IEventConsumer eventConsumer,
            ILogger<DmbFactory> logger,
            Api.Models.Instance metadata)
        {
            this.databaseContextFactory = databaseContextFactory ?? throw new ArgumentNullException(nameof(databaseContextFactory));
            this.ioManager = ioManager ?? throw new ArgumentNullException(nameof(ioManager));
            this.remoteDeploymentManagerFactory = remoteDeploymentManagerFactory ?? throw new ArgumentNullException(nameof(remoteDeploymentManagerFactory));
            this.eventConsumer = eventConsumer ?? throw new ArgumentNullException(nameof(eventConsumer));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
            this.metadata = metadata ?? throw new ArgumentNullException(nameof(metadata));
 
            cleanupTask = Task.CompletedTask;
            newerDmbTcs = new TaskCompletionSource();
            cleanupCts = new CancellationTokenSource();
            jobLockCounts = new Dictionary<long, int>();
        }
 
        public void Dispose() => cleanupCts.Dispose(); // we don't dispose nextDmbProvider here, since it might be the only thing we have
 
        public async ValueTask LoadCompileJob(CompileJob job, Action<bool>? activationAction, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(job);
 
            var newProvider = await FromCompileJob(job, cancellationToken);
            if (newProvider == null)
                return;
 
            // Do this first, because it's entirely possible when we set the tcs it will immediately need to be applied
            if (started)
            {
                var remoteDeploymentManager = remoteDeploymentManagerFactory.CreateRemoteDeploymentManager(
                    metadata,
                    job);
                await remoteDeploymentManager.StageDeployment(
                    newProvider.CompileJob,
                    activationAction,
                    cancellationToken);
            }
 
            ValueTask dmbDisposeTask;
            lock (jobLockCounts)
            {
                dmbDisposeTask = nextDmbProvider?.DisposeAsync() ?? ValueTask.CompletedTask;
                nextDmbProvider = newProvider;
 
                // Oh god dammit
                var temp = Interlocked.Exchange(ref newerDmbTcs, new TaskCompletionSource());
                temp.SetResult();
            }
 
            await dmbDisposeTask;
        }
 
        public IDmbProvider LockNextDmb(int lockCount)
        {
            if (!DmbAvailable)
                throw new InvalidOperationException("No .dmb available!");
            if (lockCount < 0)
                throw new ArgumentOutOfRangeException(nameof(lockCount), lockCount, "lockCount must be greater than or equal to 0!");
            lock (jobLockCounts)
            {
                var jobId = nextDmbProvider.CompileJob.Require(x => x.Id);
                var incremented = jobLockCounts[jobId] += lockCount;
                logger.LogTrace("Compile job {jobId} lock count now: {lockCount}", jobId, incremented);
                return nextDmbProvider;
            }
        }
 
        public async Task StartAsync(CancellationToken cancellationToken)
        {
            CompileJob? cj = null;
            await databaseContextFactory.UseContext(
                async (db) =>
                    cj = await db
                        .CompileJobs
                        .AsQueryable()
                        .Where(x => x.Job.Instance!.Id == metadata.Id)
                        .OrderByDescending(x => x.Job.StoppedAt)
                        .FirstOrDefaultAsync(cancellationToken));
 
            try
            {
                if (cj == default(CompileJob))
                    return;
                await LoadCompileJob(cj, null, cancellationToken);
            }
            finally
            {
                started = true;
            }
 
            // we dont do CleanUnusedCompileJobs here because the watchdog may have plans for them yet
        }
 
        public async Task StopAsync(CancellationToken cancellationToken)
        {
            try
            {
                lock (jobLockCounts)
                    remoteDeploymentManagerFactory.ForgetLocalStateForCompileJobs(jobLockCounts.Keys);
 
                using (cancellationToken.Register(() => cleanupCts.Cancel()))
                    await cleanupTask;
            }
            finally
            {
                started = false;
            }
        }
 
#pragma warning disable CA1506 // TODO: Decomplexify
        public async ValueTask<IDmbProvider?> FromCompileJob(CompileJob compileJob, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(compileJob);
 
            // ensure we have the entire metadata tree
            var compileJobId = compileJob.Require(x => x.Id);
            logger.LogTrace("Loading compile job {id}...", compileJobId);
            await databaseContextFactory.UseContext(
                async db => compileJob = await db
                    .CompileJobs
                    .AsQueryable()
                    .Where(x => x!.Id == compileJobId)
                    .Include(x => x.Job!)
                        .ThenInclude(x => x.StartedBy)
                    .Include(x => x.Job!)
                        .ThenInclude(x => x.Instance)
                    .Include(x => x.RevisionInformation!)
                        .ThenInclude(x => x.PrimaryTestMerge!)
                            .ThenInclude(x => x.MergedBy)
                    .Include(x => x.RevisionInformation!)
                        .ThenInclude(x => x.ActiveTestMerges!)
                            .ThenInclude(x => x.TestMerge!)
                                .ThenInclude(x => x.MergedBy)
                    .FirstAsync(cancellationToken)); // can't wait to see that query
 
            EngineVersion engineVersion;
            if (!EngineVersion.TryParse(compileJob.EngineVersion, out var engineVersionNullable))
            {
                logger.LogWarning("Error loading compile job, bad engine version: {engineVersion}", compileJob.EngineVersion);
                return null; // omae wa mou shinderu
            }
            else
                engineVersion = engineVersionNullable!;
 
            if (!compileJob.Job.StoppedAt.HasValue)
            {
                // This happens when we're told to load the compile job that is currently finished up
                // It constitutes an API violation if it's returned by the DreamDaemonController so just set it here
                // Bit of a hack, but it works out to be nearly if not the same value that's put in the DB
                logger.LogTrace("Setting missing StoppedAt for CompileJob.Job #{id}...", compileJob.Job.Id);
                compileJob.Job.StoppedAt = DateTimeOffset.UtcNow;
            }
 
            var providerSubmitted = false;
 
            void CleanupAction()
            {
                if (providerSubmitted)
                    CleanRegisteredCompileJob(compileJob);
            }
 
            var newProvider = new DmbProvider(compileJob, engineVersion, ioManager, new DisposeInvoker(CleanupAction));
            try
            {
                const string LegacyADirectoryName = "A";
                const string LegacyBDirectoryName = "B";
 
                var dmbExistsAtRoot = await ioManager.FileExists(
                    ioManager.ConcatPath(
                        newProvider.Directory,
                        newProvider.DmbName),
                    cancellationToken);
 
                if (!dmbExistsAtRoot)
                {
                    logger.LogTrace("Didn't find .dmb at game directory root, checking A/B dirs...");
                    var primaryCheckTask = ioManager.FileExists(
                        ioManager.ConcatPath(
                            newProvider.Directory,
                            LegacyADirectoryName,
                            newProvider.DmbName),
                        cancellationToken);
                    var secondaryCheckTask = ioManager.FileExists(
                        ioManager.ConcatPath(
                            newProvider.Directory,
                            LegacyBDirectoryName,
                            newProvider.DmbName),
                        cancellationToken);
 
                    if (!(await primaryCheckTask && await secondaryCheckTask))
                    {
                        logger.LogWarning("Error loading compile job, .dmb missing!");
                        return null; // omae wa mou shinderu
                    }
 
                    // rebuild the provider because it's using the legacy style directories
                    // Don't dispose it
                    logger.LogDebug("Creating legacy two folder .dmb provider targeting {aDirName} directory...", LegacyADirectoryName);
                    newProvider = new DmbProvider(compileJob, engineVersion, ioManager, new DisposeInvoker(CleanupAction), Path.DirectorySeparatorChar + LegacyADirectoryName);
                }
 
                lock (jobLockCounts)
                {
                    if (!jobLockCounts.TryGetValue(compileJobId, out int value))
                    {
                        value = 1;
                        jobLockCounts.Add(compileJobId, 1);
                    }
                    else
                        jobLockCounts[compileJobId] = ++value;
 
                    providerSubmitted = true;
 
                    logger.LogTrace("Compile job {id} lock count now: {lockCount}", compileJobId, value);
                    return newProvider;
                }
            }
            finally
            {
                if (!providerSubmitted)
                    await newProvider.DisposeAsync();
            }
        }
#pragma warning restore CA1506
 
#pragma warning disable CA1506 // TODO: Decomplexify
        public async ValueTask CleanUnusedCompileJobs(CancellationToken cancellationToken)
        {
            List<long> jobIdsToSkip;
 
            // don't clean locked directories
            lock (jobLockCounts)
                jobIdsToSkip = jobLockCounts.Keys.ToList();
 
            List<string>? jobUidsToNotErase = null;
 
            // find the uids of locked directories
            if (jobIdsToSkip.Count > 0)
            {
                await databaseContextFactory.UseContext(async db =>
                {
                    jobUidsToNotErase = (await db
                            .CompileJobs
                            .AsQueryable()
                            .Where(
                                x => x.Job.Instance!.Id == metadata.Id
                                && jobIdsToSkip.Contains(x.Id!.Value))
                            .Select(x => x.DirectoryName!.Value)
                            .ToListAsync(cancellationToken))
                        .Select(x => x.ToString())
                        .ToList();
                });
            }
            else
                jobUidsToNotErase = new List<string>();
 
            jobUidsToNotErase!.Add(SwappableDmbProvider.LiveGameDirectory);
 
            logger.LogTrace("We will not clean the following directories: {directoriesToNotClean}", String.Join(", ", jobUidsToNotErase));
 
            // cleanup
            var gameDirectory = ioManager.ResolvePath();
            await ioManager.CreateDirectory(gameDirectory, cancellationToken);
            var directories = await ioManager.GetDirectories(gameDirectory, cancellationToken);
            int deleting = 0;
            var tasks = directories.Select(async x =>
            {
                var nameOnly = ioManager.GetFileName(x);
                if (jobUidsToNotErase.Contains(nameOnly))
                    return;
                logger.LogDebug("Cleaning unused game folder: {dirName}...", nameOnly);
                try
                {
                    ++deleting;
                    await DeleteCompileJobContent(x, cancellationToken);
                }
                catch (OperationCanceledException)
                {
                    throw;
                }
                catch (Exception e)
                {
                    logger.LogWarning(e, "Error deleting directory {dirName}!", x);
                }
            }).ToList();
            if (deleting > 0)
                await Task.WhenAll(tasks);
        }
#pragma warning restore CA1506
 
        public CompileJob? LatestCompileJob()
        {
            if (!DmbAvailable)
                return null;
            return LockNextDmb(0)?.CompileJob;
        }
 
        void CleanRegisteredCompileJob(CompileJob job)
        {
            async Task HandleCleanup()
            {
                // First kill the GitHub deployment
                var remoteDeploymentManager = remoteDeploymentManagerFactory.CreateRemoteDeploymentManager(metadata, job);
 
                // DCT: None available
                var deploymentJob = remoteDeploymentManager.MarkInactive(job, CancellationToken.None);
 
                var deleteTask = DeleteCompileJobContent(job.DirectoryName!.Value.ToString(), cleanupCts.Token);
                var otherTask = cleanupTask;
 
                async Task WrapThrowableTasks()
                {
                    try
                    {
                        await ValueTaskExtensions.WhenAll(deleteTask, deploymentJob);
                    }
                    catch (Exception ex)
                    {
                        logger.LogWarning(ex, "Error cleaning up compile job {jobGuid}!", job.DirectoryName);
                    }
                }
 
                await Task.WhenAll(otherTask, WrapThrowableTasks());
            }
 
            lock (jobLockCounts)
            {
                var jobId = job.Require(x => x.Id);
                if (jobLockCounts.TryGetValue(jobId, out var currentVal))
                    if (currentVal == 1)
                    {
                        jobLockCounts.Remove(jobId);
                        logger.LogDebug("Cleaning lock-free compile job {id} => {dirName}", jobId, job.DirectoryName);
                        cleanupTask = HandleCleanup();
                    }
                    else
                    {
                        var decremented = --jobLockCounts[jobId];
                        logger.LogTrace("Compile job {id} lock count now: {lockCount}", jobId, decremented);
                    }
                else
                    logger.LogError("Extra Dispose of DmbProvider for CompileJob {compileJobId}!", jobId);
            }
        }
 
        async ValueTask DeleteCompileJobContent(string directory, CancellationToken cancellationToken)
        {
            // Then call the cleanup event, waiting here first
            await eventConsumer.HandleEvent(EventType.DeploymentCleanup, new List<string> { ioManager.ResolvePath(directory) }, true, cancellationToken);
            await ioManager.DeleteDirectory(directory, cancellationToken);
        }
    }
}