Server.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
 
using HotChocolate.Execution;
 
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
 
using Tgstation.Server.Common.Extensions;
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Core;
using Tgstation.Server.Host.IO;
using Tgstation.Server.Host.Utils;
 
namespace Tgstation.Server.Host
{
    sealed class Server : IServer, IServerControl
    {
        public bool RestartRequested { get; private set; }
 
        public bool UpdateInProgress { get; private set; }
 
        public bool WatchdogPresent =>
#if WATCHDOG_FREE_RESTART
            true;
#else
            updatePath != null;
#endif
 
        internal IHost? Host { get; private set; }
 
        readonly IHostBuilder hostBuilder;
 
        readonly List<IRestartHandler> restartHandlers;
 
        readonly string? updatePath;
 
        readonly object restartLock;
 
        ILogger<Server>? logger;
 
        GeneralConfiguration? generalConfiguration;
 
        CancellationTokenSource? cancellationTokenSource;
 
        Exception? propagatedException;
 
        Task? updateTask;
 
        bool shutdownInProgress;
 
        bool terminateIfUpdateFails;
 
        public Server(IHostBuilder hostBuilder, string? updatePath)
        {
            this.hostBuilder = hostBuilder ?? throw new ArgumentNullException(nameof(hostBuilder));
            this.updatePath = updatePath;
 
            hostBuilder.ConfigureServices(serviceCollection => serviceCollection.AddSingleton<IServerControl>(this));
 
            restartHandlers = new List<IRestartHandler>();
            restartLock = new object();
            logger = null;
        }
 
        public async ValueTask Run(CancellationToken cancellationToken)
        {
            var updateDirectory = updatePath != null ? Path.GetDirectoryName(updatePath) : null;
            using (cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
            using (var fsWatcher = updateDirectory != null ? new FileSystemWatcher(updateDirectory) : null)
            {
                if (fsWatcher != null)
                {
                    // If ever there is a NECESSARY update to the Host Watchdog, change this to use a pipe
                    // I don't know why I'm only realizing this in 2023 when this is 2019 code
                    // As it stands, FSWatchers use async I/O on Windows and block a new thread on Linux
                    // That's an acceptable, if saddening, resource loss for now
                    fsWatcher.Created += WatchForShutdownFileCreation;
                    fsWatcher.EnableRaisingEvents = true;
                }
 
                try
                {
                    using (Host = hostBuilder.Build())
                    {
                        logger = Host.Services.GetRequiredService<ILogger<Server>>();
                        try
                        {
                            using (cancellationToken.Register(() => logger.LogInformation("Server termination requested!")))
                            {
                                if (await DumpGraphQLSchemaIfRequested(Host.Services, cancellationToken))
                                    return;
 
                                var generalConfigurationOptions = Host.Services.GetRequiredService<IOptions<GeneralConfiguration>>();
                                generalConfiguration = generalConfigurationOptions.Value;
                                await Host.RunAsync(cancellationTokenSource.Token);
                            }
 
                            if (updateTask != null)
                                await updateTask;
                        }
                        catch (OperationCanceledException ex)
                        {
                            logger.LogDebug(ex, "Server run cancelled!");
                        }
                        catch (Exception ex)
                        {
                            CheckExceptionPropagation(ex);
                            throw;
                        }
                        finally
                        {
                            logger = null;
                        }
                    }
                }
                finally
                {
                    Host = null;
                }
            }
 
            CheckExceptionPropagation(null);
        }
 
        public bool TryStartUpdate(IServerUpdateExecutor updateExecutor, Version newVersion)
        {
            ArgumentNullException.ThrowIfNull(updateExecutor);
            ArgumentNullException.ThrowIfNull(newVersion);
 
            CheckSanity(true);
 
            if (updatePath == null)
                throw new InvalidOperationException("Tried to start update when server was initialized without an updatePath set!");
 
            var logger = this.logger!;
            logger.LogTrace("Begin ApplyUpdate...");
 
            CancellationToken criticalCancellationToken;
            lock (restartLock)
            {
                if (UpdateInProgress || shutdownInProgress)
                {
                    logger.LogDebug("Aborted update due to concurrency conflict!");
                    return false;
                }
 
                if (cancellationTokenSource == null)
                    throw new InvalidOperationException("Tried to update a non-running Server!");
 
                criticalCancellationToken = cancellationTokenSource.Token;
                UpdateInProgress = true;
            }
 
            async Task RunUpdate()
            {
                if (await updateExecutor.ExecuteUpdate(updatePath, criticalCancellationToken, criticalCancellationToken))
                {
                    logger.LogTrace("Update complete!");
                    await RestartImpl(newVersion, null, true, true);
                }
                else if (terminateIfUpdateFails)
                {
                    logger.LogTrace("Stopping host due to termination request...");
                    cancellationTokenSource.Cancel();
                }
                else
                {
                    logger.LogTrace("Update failed!");
                    UpdateInProgress = false;
                }
            }
 
            updateTask = RunUpdate();
            return true;
        }
 
        public IRestartRegistration RegisterForRestart(IRestartHandler handler)
        {
            ArgumentNullException.ThrowIfNull(handler);
 
            CheckSanity(false);
 
            var logger = this.logger!;
            lock (restartLock)
                if (!shutdownInProgress)
                {
                    logger.LogTrace("Registering restart handler {handlerImplementationName}...", handler);
                    restartHandlers.Add(handler);
                    return new RestartRegistration(
                        new DisposeInvoker(() =>
                        {
                            lock (restartLock)
                                if (!shutdownInProgress)
                                    restartHandlers.Remove(handler);
                        }));
                }
 
            logger.LogWarning("Restart handler {handlerImplementationName} register after a shutdown had begun!", handler);
            return new RestartRegistration(null);
        }
 
        public ValueTask Restart() => RestartImpl(null, null, true, true);
 
        public ValueTask GracefulShutdown(bool detach) => RestartImpl(null, null, false, detach);
 
        public ValueTask Die(Exception? exception)
        {
            if (exception != null)
                return RestartImpl(null, exception, false, true);
 
            StopServerImmediate();
            return ValueTask.CompletedTask;
        }
 
        async ValueTask<bool> DumpGraphQLSchemaIfRequested(IServiceProvider services, CancellationToken cancellationToken)
        {
            var internalConfigurationOptions = services.GetRequiredService<IOptions<InternalConfiguration>>();
            var apiDumpPath = internalConfigurationOptions.Value.DumpGraphQLApiPath;
            if (String.IsNullOrWhiteSpace(apiDumpPath))
                return false;
 
            logger!.LogInformation("Dumping GraphQL API spec to {path} and exiting...", apiDumpPath);
 
            // https://github.com/ChilliCream/graphql-platform/discussions/5885
            var resolver = services.GetRequiredService<IRequestExecutorResolver>();
            var executor = await resolver.GetRequestExecutorAsync(cancellationToken: cancellationToken);
            var sdl = executor.Schema.Print();
 
            var ioManager = services.GetRequiredService<IIOManager>();
            await ioManager.WriteAllBytes(apiDumpPath, Encoding.UTF8.GetBytes(sdl), cancellationToken);
            return true;
        }
 
        void CheckSanity(bool checkWatchdog)
        {
            if (checkWatchdog && !WatchdogPresent && propagatedException == null)
                throw new InvalidOperationException("Server restarts are not supported");
 
            if (cancellationTokenSource == null || logger == null)
                throw new InvalidOperationException("Tried to control a non-running Server!");
        }
 
        void CheckExceptionPropagation(Exception? otherException)
        {
            if (propagatedException == null)
                return;
 
            if (otherException != null)
                throw new AggregateException(propagatedException, otherException);
 
            throw propagatedException;
        }
 
        async ValueTask RestartImpl(Version? newVersion, Exception? exception, bool requireWatchdog, bool completeAsap)
        {
            CheckSanity(requireWatchdog);
 
            // if the watchdog isn't required and there's no issue, this is just a graceful shutdown
            bool isGracefulShutdown = !requireWatchdog && exception == null;
            var logger = this.logger!;
            logger.LogTrace(
                "Begin {restartType}...",
                isGracefulShutdown
                    ? completeAsap
                        ? "semi-graceful shutdown"
                        : "graceful shutdown"
                    : "restart");
 
            lock (restartLock)
            {
                if ((UpdateInProgress && newVersion == null) || shutdownInProgress)
                {
                    logger.LogTrace("Aborted restart due to concurrency conflict!");
                    return;
                }
 
                RestartRequested = !isGracefulShutdown;
                propagatedException ??= exception;
            }
 
            if (exception == null)
            {
                var giveHandlersTimeToWaitAround = isGracefulShutdown && !completeAsap;
                logger.LogInformation("Stopping server...");
                using var cts = new CancellationTokenSource(
                    TimeSpan.FromMinutes(
                        giveHandlersTimeToWaitAround
                            ? generalConfiguration!.ShutdownTimeoutMinutes
                            : generalConfiguration!.RestartTimeoutMinutes));
                var cancellationToken = cts.Token;
                try
                {
                    ValueTask eventsTask;
                    lock (restartLock)
                        eventsTask = ValueTaskExtensions.WhenAll(
                            restartHandlers
                                .Select(
                                    x => x.HandleRestart(newVersion, giveHandlersTimeToWaitAround, cancellationToken))
                                .ToList());
 
                    logger.LogTrace("Joining restart handlers...");
                    await eventsTask;
                }
                catch (OperationCanceledException ex)
                {
                    if (isGracefulShutdown)
                        logger.LogWarning(ex, "Graceful shutdown timeout hit! Existing DreamDaemon processes will be terminated!");
                    else
                        logger.LogError(
                            ex,
                            "Restart timeout hit! Existing DreamDaemon processes will be lost and must be killed manually before being restarted with TGS!");
                }
                catch (Exception e)
                {
                    logger.LogError(e, "Restart handlers error!");
                }
            }
 
            StopServerImmediate();
        }
 
        void WatchForShutdownFileCreation(object sender, FileSystemEventArgs eventArgs)
        {
            logger?.LogTrace("FileSystemWatcher triggered.");
 
            // TODO: Refactor this to not use System.IO function here.
            if (eventArgs.FullPath == Path.GetFullPath(updatePath!) && File.Exists(eventArgs.FullPath))
            {
                logger?.LogInformation("Host watchdog appears to be requesting server termination!");
                lock (restartLock)
                {
                    if (!UpdateInProgress)
                    {
                        StopServerImmediate();
                        return;
                    }
 
                    terminateIfUpdateFails = true;
                }
 
                logger?.LogInformation("An update is in progress, we will wait for that to complete...");
            }
        }
 
        void StopServerImmediate()
        {
            shutdownInProgress = true;
            logger!.LogDebug("Stopping host...");
            cancellationTokenSource!.Cancel();
        }
    }
}