Server.cs

Go to the documentation of this file.

using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Tgstation.Server.Host.Configuration;
using Tgstation.Server.Host.Core;
using Tgstation.Server.Host.IO;

namespace Tgstation.Server.Host
{
    sealed class Server : IServer, IServerControl
    {
        public bool RestartRequested { get; private set; }

        public bool WatchdogPresent => updatePath != null;

        readonly IHostBuilder hostBuilder;

        readonly List<IRestartHandler> restartHandlers;

        readonly string updatePath;

        readonly object restartLock;

        ILogger<Server> logger;

        GeneralConfiguration generalConfiguration;

        CancellationTokenSource cancellationTokenSource;

        Exception propagatedException;

        bool updating;

        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();
        }

        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;
        }

        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 += (a, b) =>
                    {
                        if (b.FullPath == updatePath && File.Exists(b.FullPath))
                        {
                            if (logger != null)
                                logger.LogInformation("Host watchdog appears to be requesting server termination!");
                            cancellationTokenSource.Cancel();
                        }
                    };
                    fsWatcher.EnableRaisingEvents = true;
                }

                using var 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).ConfigureAwait(false);
                    }
                }
                catch (Exception ex)
                {
                    CheckExceptionPropagation(ex);
                    throw;
                }
            }

            CheckExceptionPropagation(null);
        }

        public bool ApplyUpdate(Version version, Uri updateZipUrl, IIOManager ioManager)
        {
            if (version == null)
                throw new ArgumentNullException(nameof(version));
            if (updateZipUrl == null)
                throw new ArgumentNullException(nameof(updateZipUrl));
            if (ioManager == null)
                throw new ArgumentNullException(nameof(ioManager));

            CheckSanity(true);

            logger.LogTrace("Begin ApplyUpdate...");

            lock (restartLock)
            {
                if (updating || RestartRequested)
                {
                    logger.LogTrace("Aborted due to concurrency conflict!");
                    return false;
                }

                updating = true;
            }

            async void RunUpdate()
            {
                try
                {
                    logger.LogInformation("Updating server to version {0} ({1})...", version, updateZipUrl);

                    if (cancellationTokenSource == null)
                        throw new InvalidOperationException("Tried to update a non-running Server!");
                    var cancellationToken = cancellationTokenSource.Token;

                    logger.LogTrace("Downloading zip package...");
                    var updateZipData = await ioManager.DownloadFile(updateZipUrl, cancellationToken).ConfigureAwait(false);

                    try
                    {
                        logger.LogTrace("Exctracting zip package to {0}...", updatePath);
                        await ioManager.ZipToDirectory(updatePath, updateZipData, cancellationToken).ConfigureAwait(false);
                    }
                    catch (Exception e)
                    {
                        updating = false;
                        try
                        {
                            // important to not leave this directory around if possible
                            await ioManager.DeleteDirectory(updatePath, default).ConfigureAwait(false);
                        }
                        catch (Exception e2)
                        {
                            throw new AggregateException(e, e2);
                        }

                        throw;
                    }

                    await Restart(version, null, true).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    logger.LogInformation("Server update cancelled!");
                }
                catch (Exception e)
                {
                    logger.LogError("Error updating server! Exception: {0}", e);
                }
                finally
                {
                    updating = false;
                }
            }

            RunUpdate();
            return true;
        }

        public IRestartRegistration RegisterForRestart(IRestartHandler handler)
        {
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));

            CheckSanity(false);

            lock (restartLock)
                if (!RestartRequested)
                {
                    logger.LogTrace("Registering restart handler {0}...", handler);
                    restartHandlers.Add(handler);
                    return new RestartRegistration(() =>
                    {
                        lock (restartLock)
                            if (!RestartRequested)
                                restartHandlers.Remove(handler);
                    });
                }

            return new RestartRegistration(() => { });
        }

        public Task Restart() => Restart(null, null, true);

        async Task Restart(Version newVersion, Exception exception, bool requireWatchdog)
        {
            CheckSanity(requireWatchdog);

            logger.LogTrace("Begin Restart...");

            lock (restartLock)
            {
                if ((updating && newVersion == null) || RestartRequested)
                {
                    logger.LogTrace("Aborted due to concurrency conflict!");
                    return;
                }

                RestartRequested = true;
                propagatedException = exception;
            }

            if (exception == null)
                using (var cts = new CancellationTokenSource())
                {
                    logger.LogInformation("Restarting server...");
                    var cancellationToken = cts.Token;
                    var eventsTask = Task.WhenAll(restartHandlers.Select(x => x.HandleRestart(newVersion, cancellationToken)).ToList());

                    var expiryTask = Task.Delay(TimeSpan.FromMilliseconds(generalConfiguration.RestartTimeout));
                    await Task.WhenAny(eventsTask, expiryTask).ConfigureAwait(false);
                    logger.LogTrace("Joining restart handlers...");
                    cts.Cancel();
                    try
                    {
                        await eventsTask.ConfigureAwait(false);
                    }
                    catch (OperationCanceledException)
                    {
                        logger.LogError("Restart timeout hit! Existing DreamDaemon processes will be lost and must be killed manually before being restarted with TGS!");
                    }
                    catch (Exception e)
                    {
                        logger.LogError("Restart handlers error! Exception: {0}", e);
                    }
                }

            logger.LogTrace("Stopping host...");
            cancellationTokenSource.Cancel();
        }

        public Task Die(Exception exception) => Restart(null, exception, false);
    }
}