Server.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.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
 
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Core;
 
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();
        }
 
        public async Task Run(CancellationToken cancellationToken)
        {
            using (cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
            using (var fsWatcher = updatePath != null ? new FileSystemWatcher(Path.GetDirectoryName(updatePath)) : null)
            {
                if (fsWatcher != null)
                {
                    fsWatcher.Created += WatchForShutdownFileCreation;
                    fsWatcher.EnableRaisingEvents = true;
                }
 
                try
                {
                    using (Host = hostBuilder.Build())
                    {
                        try
                        {
                            logger = Host.Services.GetRequiredService<ILogger<Server>>();
                            using (cancellationToken.Register(() => logger.LogInformation("Server termination requested!")))
                            {
                                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);
 
            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);
 
            lock (restartLock)
                if (!shutdownInProgress)
                {
                    logger.LogTrace("Registering restart handler {handlerImplementationName}...", handler);
                    restartHandlers.Add(handler);
                    return new RestartRegistration(() =>
                    {
                        lock (restartLock)
                            if (!shutdownInProgress)
                                restartHandlers.Remove(handler);
                    });
                }
 
            logger.LogWarning("Restart handler {handlerImplementationName} register after a shutdown had begun!", handler);
            return new RestartRegistration(null);
        }
 
        public Task Restart() => RestartImpl(null, null, true, true);
 
        public Task GracefulShutdown(bool detach) => RestartImpl(null, null, false, detach);
 
        public Task Die(Exception exception) => RestartImpl(null, exception, false, 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 Task 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;
            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
                {
                    var eventsTask = Task.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.LogTrace("Stopping host...");
            cancellationTokenSource.Cancel();
        }
    }
}