Server.cs

Go to the documentation of this file.

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

namespace Tgstation.Server.Host
{
#pragma warning disable CA1001 // Types that own disposable fields should be disposable
    sealed class Server : IServer, IServerControl
#pragma warning restore CA1001 // Types that own disposable fields should be disposable
    {
        public bool RestartRequested { get; private set; }

        public bool WatchdogPresent => updatePath != null;

        readonly IWebHostBuilder webHostBuilder;

        readonly List<IRestartHandler> restartHandlers;

        readonly string updatePath;

        ILogger<Server> logger;

        CancellationTokenSource cancellationTokenSource;

        Exception propagatedException;

        bool updating;

        public Server(IWebHostBuilder webHostBuilder, string updatePath)
        {
            this.webHostBuilder = webHostBuilder ?? throw new ArgumentNullException(nameof(webHostBuilder));
            this.updatePath = updatePath;

            webHostBuilder.ConfigureServices(serviceCollection => serviceCollection.AddSingleton<IServerControl>(this));

            restartHandlers = new List<IRestartHandler>();
        }

        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()
        {
            if (propagatedException != null)
                throw propagatedException;
        }

        public async Task RunAsync(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))
                            cancellationTokenSource.Cancel();
                    };
                    fsWatcher.EnableRaisingEvents = true;
                }
                
                using (var webHost = webHostBuilder.Build())
                    try
                    {
                        logger = webHost.Services.GetRequiredService<ILogger<Server>>();
                        await webHost.RunAsync(cancellationTokenSource.Token).ConfigureAwait(false);
                    }
                    catch (OperationCanceledException)
                    {
                        CheckExceptionPropagation();
                        throw;
                    }
            }
            CheckExceptionPropagation();
        }

        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 (this)
            {
                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).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 (this)
                if (!RestartRequested)
                {
                    logger.LogTrace("Registering restart handler {0}...", handler);
                    restartHandlers.Add(handler);
                    return new RestartRegistration(() =>
                    {
                        lock (this)
                            if (!RestartRequested)
                                restartHandlers.Remove(handler);
                    });
                }
            return new RestartRegistration(() => { });
        }

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

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

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

            lock (this)
            {
                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());
                    //YA GOT 10 SECONDS
                    var expiryTask = Task.Delay(TimeSpan.FromSeconds(10));
                    await Task.WhenAny(eventsTask, expiryTask).ConfigureAwait(false);
                    logger.LogTrace("Joining restart handlers...");
                    cts.Cancel();
                    try
                    {
                        await eventsTask.ConfigureAwait(false);
                    }
                    catch (OperationCanceledException) { }
                    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);
    }
}