InstanceManager.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
 
using Tgstation.Server.Api.Models;
using Tgstation.Server.Common;
using Tgstation.Server.Common.Extensions;
using Tgstation.Server.Host.Components.Interop;
using Tgstation.Server.Host.Components.Interop.Bridge;
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Controllers;
using Tgstation.Server.Host.Core;
using Tgstation.Server.Host.Database;
using Tgstation.Server.Host.Extensions;
using Tgstation.Server.Host.IO;
using Tgstation.Server.Host.Jobs;
using Tgstation.Server.Host.Models;
using Tgstation.Server.Host.Security;
using Tgstation.Server.Host.Swarm;
using Tgstation.Server.Host.System;
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host.Components
{
    sealed class InstanceManager :
        IInstanceManager,
        IInstanceCoreProvider,
        IHostedService,
        IBridgeRegistrar,
        IAsyncDisposable
    {
        public Task Ready => readyTcs.Task;
 
        readonly IInstanceFactory instanceFactory;
 
        readonly IIOManager ioManager;
 
        readonly IDatabaseContextFactory databaseContextFactory;
 
        readonly IAssemblyInformationProvider assemblyInformationProvider;
 
        readonly IJobService jobService;
 
        readonly IServerControl serverControl;
 
        readonly ISystemIdentityFactory systemIdentityFactory;
 
        readonly IAsyncDelayer asyncDelayer;
 
        readonly IServerPortProvider serverPortProvider;
 
        readonly ISwarmServiceController swarmServiceController;
 
        readonly IConsole console;
 
        readonly IPlatformIdentifier platformIdentifier;
 
        readonly ILogger<InstanceManager> logger;
 
        readonly Dictionary<long, ReferenceCountingContainer<IInstance, InstanceWrapper>> instances;
 
        readonly Dictionary<string, IBridgeHandler> bridgeHandlers;
 
        readonly SemaphoreSlim instanceStateChangeSemaphore;
 
        readonly GeneralConfiguration generalConfiguration;
 
        readonly SwarmConfiguration swarmConfiguration;
 
        readonly TaskCompletionSource readyTcs;
 
        readonly CancellationTokenSource startupCancellationTokenSource;
 
        readonly CancellationTokenSource shutdownCancellationTokenSource;
 
        readonly string? originalConsoleTitle;
 
        Task? startupTask;
 
        bool disposed;
 
        public InstanceManager(
            IInstanceFactory instanceFactory,
            IIOManager ioManager,
            IDatabaseContextFactory databaseContextFactory,
            IAssemblyInformationProvider assemblyInformationProvider,
            IJobService jobService,
            IServerControl serverControl,
            ISystemIdentityFactory systemIdentityFactory,
            IAsyncDelayer asyncDelayer,
            IServerPortProvider serverPortProvider,
            ISwarmServiceController swarmServiceController,
            IConsole console,
            IPlatformIdentifier platformIdentifier,
            IOptions<GeneralConfiguration> generalConfigurationOptions,
            IOptions<SwarmConfiguration> swarmConfigurationOptions,
            ILogger<InstanceManager> logger)
        {
            this.instanceFactory = instanceFactory ?? throw new ArgumentNullException(nameof(instanceFactory));
            this.ioManager = ioManager ?? throw new ArgumentNullException(nameof(ioManager));
            this.databaseContextFactory = databaseContextFactory ?? throw new ArgumentNullException(nameof(databaseContextFactory));
            this.assemblyInformationProvider = assemblyInformationProvider ?? throw new ArgumentNullException(nameof(assemblyInformationProvider));
            this.jobService = jobService ?? throw new ArgumentNullException(nameof(jobService));
            this.serverControl = serverControl ?? throw new ArgumentNullException(nameof(serverControl));
            this.systemIdentityFactory = systemIdentityFactory ?? throw new ArgumentNullException(nameof(systemIdentityFactory));
            this.asyncDelayer = asyncDelayer ?? throw new ArgumentNullException(nameof(asyncDelayer));
            this.serverPortProvider = serverPortProvider ?? throw new ArgumentNullException(nameof(serverPortProvider));
            this.swarmServiceController = swarmServiceController ?? throw new ArgumentNullException(nameof(swarmServiceController));
            this.console = console ?? throw new ArgumentNullException(nameof(console));
            this.platformIdentifier = platformIdentifier ?? throw new ArgumentNullException(nameof(platformIdentifier));
            generalConfiguration = generalConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(generalConfigurationOptions));
            swarmConfiguration = swarmConfigurationOptions?.Value ?? throw new ArgumentNullException(nameof(swarmConfigurationOptions));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
            originalConsoleTitle = console.Title;
 
            instances = new Dictionary<long, ReferenceCountingContainer<IInstance, InstanceWrapper>>();
            bridgeHandlers = new Dictionary<string, IBridgeHandler>();
            readyTcs = new TaskCompletionSource();
            instanceStateChangeSemaphore = new SemaphoreSlim(1);
            startupCancellationTokenSource = new CancellationTokenSource();
            shutdownCancellationTokenSource = new CancellationTokenSource();
        }
 
        public async ValueTask DisposeAsync()
        {
            lock (instances)
            {
                if (disposed)
                    return;
                disposed = true;
            }
 
            foreach (var instanceKvp in instances)
                await instanceKvp.Value.Instance.DisposeAsync();
 
            instanceStateChangeSemaphore.Dispose();
            startupCancellationTokenSource.Dispose();
            shutdownCancellationTokenSource.Dispose();
 
            logger.LogInformation("Server shutdown");
        }
 
        public IInstanceReference? GetInstanceReference(Api.Models.Instance metadata)
        {
            ArgumentNullException.ThrowIfNull(metadata);
 
            lock (instances)
            {
                if (!instances.TryGetValue(metadata.Require(x => x.Id), out var instance))
                    return null;
 
                return instance.AddReference();
            }
        }
 
        public async ValueTask MoveInstance(Models.Instance instance, string oldPath, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(oldPath);
 
            using var lockContext = await SemaphoreSlimContext.Lock(instanceStateChangeSemaphore, cancellationToken);
            using var instanceReferenceCheck = GetInstanceReference(instance);
            if (instanceReferenceCheck != null)
                throw new InvalidOperationException("Cannot move an online instance!");
            var newPath = instance.Path!;
            try
            {
                await ioManager.MoveDirectory(oldPath, newPath, cancellationToken);
 
                // Delete the Game directory to clear out broken symlinks
                var instanceGameIOManager = instanceFactory.CreateGameIOManager(instance);
                await instanceGameIOManager.DeleteDirectory(DefaultIOManager.CurrentDirectory, cancellationToken);
            }
            catch (Exception ex)
            {
                logger.LogError(
                    ex,
                    "Error moving instance {instanceId}!",
                    instance.Id);
                try
                {
                    logger.LogDebug("Reverting instance {instanceId}'s path to {oldPath} in the DB...", instance.Id, oldPath);
 
                    // DCT: Operation must always run
                    await databaseContextFactory.UseContextTaskReturn(db =>
                    {
                        var targetInstance = new Models.Instance
                        {
                            Id = instance.Id,
                        };
                        db.Instances.Attach(targetInstance);
                        targetInstance.Path = oldPath;
                        return db.Save(CancellationToken.None);
                    });
                }
                catch (Exception innerEx)
                {
                    logger.LogCritical(
                        innerEx,
                        "Error reverting database after failing to move instance {instanceId}! Attempting to detach...",
                        instance.Id);
 
                    try
                    {
                        // DCT: Operation must always run
                        await ioManager.WriteAllBytes(
                            ioManager.ConcatPath(oldPath, InstanceController.InstanceAttachFileName),
                            Array.Empty<byte>(),
                            CancellationToken.None);
                    }
                    catch (Exception tripleEx)
                    {
                        logger.LogCritical(
                            tripleEx,
                            "Okay, what gamma radiation are you under? Failed to write instance attach file!");
 
                        throw new AggregateException(tripleEx, innerEx, ex);
                    }
 
                    throw new AggregateException(ex, innerEx);
                }
 
                throw;
            }
        }
 
        public async ValueTask OfflineInstance(Models.Instance metadata, User user, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(metadata);
 
            using (await SemaphoreSlimContext.Lock(instanceStateChangeSemaphore, cancellationToken))
            {
                ReferenceCountingContainer<IInstance, InstanceWrapper>? container;
                var instanceId = metadata.Require(x => x.Id);
                lock (instances)
                {
                    if (!instances.TryGetValue(instanceId, out container))
                    {
                        logger.LogDebug("Not offlining removed instance {instanceId}", metadata.Id);
                        return;
                    }
 
                    instances.Remove(instanceId);
                }
 
                logger.LogInformation("Offlining instance ID {instanceId}", metadata.Id);
 
                try
                {
                    await container.OnZeroReferences.WaitAsync(cancellationToken);
 
                    // we are the one responsible for cancelling his jobs
                    ValueTask<Job?[]> groupedTask = default;
                    await databaseContextFactory.UseContext(
                        async db =>
                        {
                            var jobs = await db
                                .Jobs
                                .AsQueryable()
                                .Where(x => x.Instance!.Id == metadata.Id && !x.StoppedAt.HasValue)
                                .Select(x => new Job(x.Id!.Value))
                                .ToListAsync(cancellationToken);
 
                            groupedTask = ValueTaskExtensions.WhenAll(
                                jobs.Select(job => jobService.CancelJob(job, user, true, cancellationToken)),
                                jobs.Count);
                        });
 
                    await groupedTask;
                }
                catch
                {
                    // not too late to change your mind
                    lock (instances)
                        instances.Add(instanceId, container);
 
                    throw;
                }
 
                try
                {
                    // at this point we can't really stop offlining the instance just because the request was cancelled
                    await container.Instance.StopAsync(shutdownCancellationTokenSource.Token);
                }
                finally
                {
                    await container.Instance.DisposeAsync();
                }
            }
        }
 
        public async ValueTask OnlineInstance(Models.Instance metadata, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(metadata);
 
            var instanceId = metadata.Require(x => x.Id);
            using var lockContext = await SemaphoreSlimContext.Lock(instanceStateChangeSemaphore, cancellationToken);
            lock (instances)
                if (instances.ContainsKey(instanceId))
                {
                    logger.LogDebug("Aborting instance creation due to it seemingly already being online");
                    return;
                }
 
            logger.LogInformation("Onlining instance ID {instanceId} ({instanceName}) at {instancePath}", metadata.Id, metadata.Name, metadata.Path);
            var instance = await instanceFactory.CreateInstance(this, metadata);
            try
            {
                await instance.StartAsync(cancellationToken);
 
                try
                {
                    lock (instances)
                        instances.Add(
                            instanceId,
                            new ReferenceCountingContainer<IInstance, InstanceWrapper>(instance));
                }
                catch (Exception ex)
                {
                    logger.LogError("Unable to commit onlined instance {instanceId} into service, offlining!", metadata.Id);
                    try
                    {
                        // DCT: Must always run
                        await instance.StopAsync(CancellationToken.None);
                    }
                    catch (Exception innerEx)
                    {
                        throw new AggregateException(innerEx, ex);
                    }
 
                    throw;
                }
            }
            catch
            {
                await instance.DisposeAsync();
                throw;
            }
        }
 
        public Task StartAsync(CancellationToken cancellationToken)
        {
            cancellationToken.Register(startupCancellationTokenSource.Cancel);
            startupTask = Initialize(startupCancellationTokenSource.Token);
            return Task.CompletedTask;
        }
 
        public async Task StopAsync(CancellationToken cancellationToken)
        {
            try
            {
                using (cancellationToken.Register(shutdownCancellationTokenSource.Cancel))
                    try
                    {
                        if (startupTask == null)
                        {
                            logger.LogWarning("InstanceManager was never started!");
                            return;
                        }
 
                        logger.LogDebug("Stopping instance manager...");
 
                        if (!startupTask.IsCompleted)
                        {
                            logger.LogTrace("Interrupting startup task...");
                            startupCancellationTokenSource.Cancel();
                            await startupTask;
                        }
 
                        var instanceFactoryStopTask = instanceFactory.StopAsync(cancellationToken);
                        await jobService.StopAsync(cancellationToken);
 
                        async ValueTask OfflineInstanceImmediate(IInstance instance, CancellationToken cancellationToken)
                        {
                            try
                            {
                                await instance.StopAsync(cancellationToken);
                            }
                            catch (Exception ex)
                            {
                                logger.LogError(ex, "Instance shutdown exception!");
                            }
                        }
 
                        await ValueTaskExtensions.WhenAll(instances.Select(x => OfflineInstanceImmediate(x.Value.Instance, cancellationToken)));
                        await instanceFactoryStopTask;
 
                        await swarmServiceController.Shutdown(cancellationToken);
                    }
                    finally
                    {
                        if (originalConsoleTitle != null)
                            console.SetTitle(originalConsoleTitle);
                    }
            }
            catch (Exception ex)
            {
                logger.LogCritical(ex, "Instance manager stop exception!");
            }
        }
 
        public async ValueTask<BridgeResponse?> ProcessBridgeRequest(BridgeParameters parameters, CancellationToken cancellationToken)
        {
            ArgumentNullException.ThrowIfNull(parameters);
 
            var accessIdentifier = parameters.AccessIdentifier;
            if (accessIdentifier == null)
            {
                logger.LogWarning("Received invalid bridge request with null access identifier!");
                return null;
            }
 
            IBridgeHandler? bridgeHandler = null;
            for (var i = 0; bridgeHandler == null && i < 30; ++i)
            {
                // There's a miniscule time period where we could potentially receive a bridge request and not have the registration ready when we launch DD
                // This is a stopgap
                Task delayTask = Task.CompletedTask;
                lock (bridgeHandlers)
                    if (!bridgeHandlers.TryGetValue(accessIdentifier, out bridgeHandler))
                        delayTask = asyncDelayer.Delay(TimeSpan.FromMilliseconds(100), cancellationToken);
 
                await delayTask;
            }
 
            if (bridgeHandler == null)
                lock (bridgeHandlers)
                    if (!bridgeHandlers.TryGetValue(accessIdentifier, out bridgeHandler))
                    {
                        logger.LogWarning("Received invalid bridge request with access identifier: {accessIdentifier}", accessIdentifier);
                        return null;
                    }
 
            return await bridgeHandler.ProcessBridgeRequest(parameters, cancellationToken);
        }
 
        public IBridgeRegistration RegisterHandler(IBridgeHandler bridgeHandler)
        {
            ArgumentNullException.ThrowIfNull(bridgeHandler);
 
            var accessIdentifier = bridgeHandler.DMApiParameters.AccessIdentifier
                ?? throw new InvalidOperationException("Attempted bridge registration with null AccessIdentifier!");
            lock (bridgeHandlers)
            {
                bridgeHandlers.Add(accessIdentifier, bridgeHandler);
                logger.LogTrace("Registered bridge handler: {accessIdentifier}", accessIdentifier);
            }
 
            return new BridgeRegistration(() =>
            {
                lock (bridgeHandlers)
                {
                    bridgeHandlers.Remove(accessIdentifier);
                    logger.LogTrace("Unregistered bridge handler: {accessIdentifier}", accessIdentifier);
                }
            });
        }
 
        public IInstanceCore? GetInstance(Models.Instance metadata)
        {
            lock (instances)
            {
                instances.TryGetValue(metadata.Require(x => x.Id), out var container);
                return container?.Instance;
            }
        }
 
        async Task Initialize(CancellationToken cancellationToken)
        {
            try
            {
                logger.LogInformation("{versionString}", assemblyInformationProvider.VersionString);
                console.SetTitle(assemblyInformationProvider.VersionString);
 
                CheckSystemCompatibility();
 
                // To let the web server startup immediately before we do any intense work
                await Task.Yield();
 
                await InitializeSwarm(cancellationToken);
 
                List<Models.Instance>? dbInstances = null;
 
                async ValueTask EnumerateInstances(IDatabaseContext databaseContext)
                    => dbInstances = await databaseContext
                        .Instances
                        .AsQueryable()
                        .Where(x => x.Online!.Value && x.SwarmIdentifer == swarmConfiguration.Identifier)
                        .Include(x => x.RepositorySettings)
                        .Include(x => x.ChatSettings)
                            .ThenInclude(x => x.Channels)
                        .Include(x => x.DreamDaemonSettings)
                        .ToListAsync(cancellationToken);
 
                var instanceEnumeration = databaseContextFactory.UseContext(EnumerateInstances);
 
                var factoryStartup = instanceFactory.StartAsync(cancellationToken);
                var jobManagerStartup = jobService.StartAsync(cancellationToken);
 
                await Task.WhenAll(instanceEnumeration.AsTask(), factoryStartup, jobManagerStartup);
 
                var instanceOnliningTasks = dbInstances!.Select(
                    async metadata =>
                    {
                        try
                        {
                            await OnlineInstance(metadata, cancellationToken);
                        }
                        catch (Exception ex)
                        {
                            logger.LogError(ex, "Failed to online instance {instanceId}!", metadata.Id);
                        }
                    });
 
                await Task.WhenAll(instanceOnliningTasks);
 
                logger.LogInformation("Server ready!");
                readyTcs.SetResult();
 
                // this needs to happen after the HTTP API opens with readyTcs otherwise it can race and cause failed bridge requests with 503's
                jobService.Activate(this);
            }
            catch (OperationCanceledException ex)
            {
                logger.LogInformation(ex, "Cancelled instance manager initialization!");
            }
            catch (Exception e)
            {
                logger.LogCritical(e, "Instance manager startup error!");
                try
                {
                    await serverControl.Die(e);
                    return;
                }
                catch (Exception e2)
                {
                    logger.LogCritical(e2, "Failed to kill server!");
                }
            }
        }
 
        void CheckSystemCompatibility()
        {
            logger.LogDebug("Running as user: {username}", Environment.UserName);
 
            generalConfiguration.CheckCompatibility(logger);
 
            using (var systemIdentity = systemIdentityFactory.GetCurrent())
            {
                if (!systemIdentity.CanCreateSymlinks)
                    throw new InvalidOperationException($"The user running {Constants.CanonicalPackageName} cannot create symlinks! Please try running as an administrative user!");
            }
 
            // This runs before the real socket is opened, ensures we don't perform reattaches unless we're fairly certain the bind won't fail
            // If it does fail, DD will be killed.
            SocketExtensions.BindTest(platformIdentifier, serverPortProvider.HttpApiPort, true, false);
        }
 
        async ValueTask InitializeSwarm(CancellationToken cancellationToken)
        {
            SwarmRegistrationResult registrationResult;
            do
            {
                registrationResult = await swarmServiceController.Initialize(cancellationToken);
 
                if (registrationResult == SwarmRegistrationResult.Unauthorized)
                    throw new InvalidOperationException("Swarm private key does not match the swarm controller's!");
 
                if (registrationResult == SwarmRegistrationResult.VersionMismatch)
                    throw new InvalidOperationException("Swarm controller's TGS version does not match our own!");
 
                if (registrationResult != SwarmRegistrationResult.Success)
                    await asyncDelayer.Delay(TimeSpan.FromSeconds(5), cancellationToken);
            }
            while (registrationResult != SwarmRegistrationResult.Success && !cancellationToken.IsCancellationRequested);
        }
    }
}