SessionPersistor.cs

Go to the documentation of this file.

using System;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Host.Components.Deployment;
using Tgstation.Server.Host.Database;
using Tgstation.Server.Host.System;
 
using Z.EntityFramework.Plus;
 
namespace Tgstation.Server.Host.Components.Session
{
    sealed class SessionPersistor : ISessionPersistor
    {
        readonly IDatabaseContextFactory databaseContextFactory;
 
        readonly IDmbFactory dmbFactory;
 
        readonly IProcessExecutor processExecutor;
 
        readonly ILogger<SessionPersistor> logger;
 
        readonly Api.Models.Instance metadata;
 
        public SessionPersistor(
            IDatabaseContextFactory databaseContextFactory,
            IDmbFactory dmbFactory,
            IProcessExecutor processExecutor,
            ILogger<SessionPersistor> logger,
            Api.Models.Instance metadata)
        {
            this.databaseContextFactory = databaseContextFactory ?? throw new ArgumentNullException(nameof(databaseContextFactory));
            this.dmbFactory = dmbFactory ?? throw new ArgumentNullException(nameof(dmbFactory));
            this.processExecutor = processExecutor ?? throw new ArgumentNullException(nameof(processExecutor));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
            this.metadata = metadata ?? throw new ArgumentNullException(nameof(metadata));
        }
 
        public Task Save(ReattachInformation reattachInformation, CancellationToken cancellationToken) => databaseContextFactory.UseContext(async (db) =>
        {
            ArgumentNullException.ThrowIfNull(reattachInformation);
 
            logger.LogDebug("Saving reattach information: {info}...", reattachInformation);
 
            await ClearImpl(db, false, cancellationToken);
 
            var dbReattachInfo = new Models.ReattachInformation
            {
                AccessIdentifier = reattachInformation.AccessIdentifier,
                CompileJobId = reattachInformation.Dmb.CompileJob.Id.Value,
                InitialCompileJobId = reattachInformation.InitialDmb?.CompileJob.Id.Value,
                Port = reattachInformation.Port,
                ProcessId = reattachInformation.ProcessId,
                RebootState = reattachInformation.RebootState,
                LaunchSecurityLevel = reattachInformation.LaunchSecurityLevel,
                LaunchVisibility = reattachInformation.LaunchVisibility,
            };
 
            db.ReattachInformations.Add(dbReattachInfo);
            await db.Save(cancellationToken);
        });
 
        public async Task<ReattachInformation> Load(CancellationToken cancellationToken)
        {
            Models.ReattachInformation result = null;
            TimeSpan? topicTimeout = null;
 
            async Task KillProcess(Models.ReattachInformation reattachInfo)
            {
                try
                {
                    await using var process = processExecutor.GetProcess(reattachInfo.ProcessId);
                    if (process != null)
                    {
                        if (reattachInfo == result)
                        {
                            logger.LogWarning("Killing PID {pid} associated with CompileJob-less reattach information...", reattachInfo.ProcessId);
                        }
                        else
                        {
                            logger.LogWarning("Killing PID {pid} associated with extra reattach information...", reattachInfo.ProcessId);
                        }
 
                        process.Terminate();
                        await process.Lifetime;
                    }
                }
                catch (Exception ex)
                {
                    logger.LogWarning(ex, "Failed to kill process!");
                }
            }
 
            await databaseContextFactory.UseContext(async (db) =>
            {
                var dbReattachInfos = await db
                    .ReattachInformations
                    .AsQueryable()
                    .Where(x => x.CompileJob.Job.Instance.Id == metadata.Id)
                    .Include(x => x.CompileJob)
                    .Include(x => x.InitialCompileJob)
                    .ToListAsync(cancellationToken);
                result = dbReattachInfos.FirstOrDefault();
                if (result == default)
                    return;
 
                var timeoutMilliseconds = await db
                    .Instances
                    .AsQueryable()
                    .Where(x => x.Id == metadata.Id)
                    .Select(x => x.DreamDaemonSettings.TopicRequestTimeout)
                    .FirstOrDefaultAsync(cancellationToken);
 
                if (timeoutMilliseconds == default)
                {
                    logger.LogCritical("Missing TopicRequestTimeout!");
                    return;
                }
 
                topicTimeout = TimeSpan.FromMilliseconds(timeoutMilliseconds.Value);
 
                bool first = true;
                foreach (var reattachInfo in dbReattachInfos)
                {
                    if (first)
                    {
                        first = false;
                        continue;
                    }
 
                    await KillProcess(reattachInfo);
 
                    db.ReattachInformations.Remove(reattachInfo);
                    logger.LogTrace("Deleting ReattachInformation {id}...", reattachInfo.Id);
                }
 
                await db.Save(cancellationToken);
            });
 
            if (!topicTimeout.HasValue)
            {
                logger.LogDebug("Reattach information not found!");
                return null;
            }
 
            var dmb = await dmbFactory.FromCompileJob(result.CompileJob, cancellationToken);
            if (dmb == null)
            {
                logger.LogError("Unable to reattach! Could not load .dmb!");
                await KillProcess(result);
 
                await databaseContextFactory.UseContext(async db =>
                {
                    logger.LogTrace("Deleting ReattachInformation {id}...", result.Id);
                    await db
                        .ReattachInformations
                        .AsQueryable()
                        .Where(x => x.Id == result.Id)
                        .DeleteAsync(cancellationToken);
                });
                return null;
            }
 
            IDmbProvider initialDmb = null;
            if (result.InitialCompileJob != null)
            {
                logger.LogTrace("Loading initial compile job...");
                initialDmb = await dmbFactory.FromCompileJob(result.InitialCompileJob, cancellationToken);
            }
 
            logger.LogTrace("Retrieved ReattachInformation");
 
            var info = new ReattachInformation(
                result,
                dmb,
                initialDmb,
                topicTimeout.Value);
 
            logger.LogDebug("Reattach information loaded: {info}", info);
 
            return info;
        }
 
        public Task Clear(CancellationToken cancellationToken) => databaseContextFactory
            .UseContext(
                db =>
                {
                    logger.LogDebug("Clearing reattach information");
                    return ClearImpl(db, true, cancellationToken);
                });
 
        async Task ClearImpl(IDatabaseContext databaseContext, bool instant, CancellationToken cancellationToken)
        {
            var baseQuery = databaseContext
                .ReattachInformations
                .AsQueryable()
                .Where(x => x.CompileJob.Job.Instance.Id == metadata.Id);
 
            if (instant)
                await baseQuery
                    .DeleteAsync(cancellationToken);
            else
            {
                var results = await baseQuery.ToListAsync(cancellationToken);
                foreach (var result in results)
                    databaseContext.ReattachInformations.Remove(result);
            }
        }
    }
}