DmbFactory.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
 
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;
 
namespace Tgstation.Server.Host.Components.Deployment
{
    sealed class DmbFactory : IDmbFactory, ICompileJobSink
    {
        public Task OnNewerDmb
        {
            get
            {
                lock (jobLockCounts)
                    return newerDmbTcs.Task;
            }
        }
 
        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 IDictionary<long, int> jobLockCounts;
 
        Task cleanupTask;
 
        TaskCompletionSource newerDmbTcs;
 
        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 Task LoadCompileJob(CompileJob job, 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,
                        cancellationToken);
            }
 
            lock (jobLockCounts)
            {
                nextDmbProvider?.Dispose();
                nextDmbProvider = newProvider;
 
                // Oh god dammit
                var temp = newerDmbTcs;
                newerDmbTcs = new TaskCompletionSource();
                temp.SetResult();
            }
        }
 
        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.Id;
                var incremented = jobLockCounts[jobId.Value] += lockCount;
                logger.LogTrace("Compile job {0} lock count now: {1}", 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);
            });
 
            if (cj == default(CompileJob))
                return;
            await LoadCompileJob(cj, cancellationToken);
            started = true;
 
            // we dont do CleanUnusedCompileJobs here because the watchdog may have plans for them yet
        }
 
        public async Task StopAsync(CancellationToken cancellationToken)
        {
            try
            {
                using (cancellationToken.Register(() => cleanupCts.Cancel()))
                    await cleanupTask;
            }
            finally
            {
                started = false;
            }
        }
 
#pragma warning disable CA1506 // TODO: Decomplexify
        public async Task<IDmbProvider> FromCompileJob(CompileJob compileJob, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(compileJob);
 
            // ensure we have the entire metadata tree
            logger.LogTrace("Loading compile job {0}...", compileJob.Id);
            await databaseContextFactory.UseContext(
                async db => compileJob = await db
                    .CompileJobs
                    .AsQueryable()
                    .Where(x => x.Id == compileJob.Id)
                    .Include(x => x.Job)
                        .ThenInclude(x => x.StartedBy)
                    .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
 
            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 #{0}...", compileJob.Job.Id);
                compileJob.Job.StoppedAt = DateTimeOffset.UtcNow;
            }
 
            var providerSubmitted = false;
 
            void CleanupAction()
            {
                if (providerSubmitted)
                    CleanRegisteredCompileJob(compileJob);
            }
 
            var newProvider = new DmbProvider(compileJob, ioManager, 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 {0} directory...", LegacyADirectoryName);
                    newProvider = new DmbProvider(compileJob, ioManager, CleanupAction, Path.DirectorySeparatorChar + LegacyADirectoryName);
                }
 
                lock (jobLockCounts)
                {
                    if (!jobLockCounts.TryGetValue(compileJob.Id.Value, out int value))
                    {
                        value = 1;
                        jobLockCounts.Add(compileJob.Id.Value, 1);
                    }
                    else
                        jobLockCounts[compileJob.Id.Value] = ++value;
 
                    providerSubmitted = true;
 
                    logger.LogTrace("Compile job {0} lock count now: {1}", compileJob.Id, value);
                    return newProvider;
                }
            }
            finally
            {
                if (!providerSubmitted)
                    newProvider.Dispose();
            }
        }
#pragma warning restore CA1506
 
#pragma warning disable CA1506 // TODO: Decomplexify
        public async Task CleanUnusedCompileJobs(CancellationToken cancellationToken)
        {
            List<long> jobIdsToSkip;
 
            // don't clean locked directories
            lock (jobLockCounts)
                jobIdsToSkip = jobLockCounts.Select(x => x.Key).ToList();
 
            List<string> jobUidsToNotErase = null;
 
            // find the uids of locked directories
            if (jobIdsToSkip.Any())
            {
                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: {0}", 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: {0}...", nameOnly);
                try
                {
                    ++deleting;
                    await DeleteCompileJobContent(x, cancellationToken);
                }
                catch (OperationCanceledException)
                {
                    throw;
                }
                catch (Exception e)
                {
                    logger.LogWarning(e, "Error deleting directory {0}!", 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.ToString(), cleanupCts.Token);
                var otherTask = cleanupTask;
 
                async Task WrapThrowableTasks()
                {
                    try
                    {
                        await Task.WhenAll(deleteTask, deploymentJob);
                    }
                    catch (Exception ex)
                    {
                        logger.LogWarning(ex, "Error cleaning up compile job {jobGuid}!", job.DirectoryName);
                    }
                }
 
                await Task.WhenAll(otherTask, WrapThrowableTasks());
            }
 
            lock (jobLockCounts)
                if (jobLockCounts.TryGetValue(job.Id.Value, out var currentVal))
                    if (currentVal == 1)
                    {
                        jobLockCounts.Remove(job.Id.Value);
                        logger.LogDebug("Cleaning lock-free compile job {0} => {1}", job.Id, job.DirectoryName);
                        cleanupTask = HandleCleanup();
                    }
                    else
                    {
                        var decremented = --jobLockCounts[job.Id.Value];
                        logger.LogTrace("Compile job {0} lock count now: {1}", job.Id, decremented);
                    }
                else
                    logger.LogError("Extra Dispose of DmbProvider for CompileJob {compileJobId}!", job.Id);
        }
 
        async Task 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);
        }
    }
}