ServiceLifetime.cs

Go to the documentation of this file.

using System;
using System.IO;
using System.IO.Pipes;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Host.Common;
using Tgstation.Server.Host.Watchdog;
 
namespace Tgstation.Server.Host.Service
{
    sealed class ServiceLifetime : ISignalChecker, IAsyncDisposable
    {
        readonly ILogger<ServiceLifetime> logger;
 
        readonly Task watchdogTask;
 
        readonly CancellationTokenSource cancellationTokenSource;
 
        AnonymousPipeServerStream? commandPipeServer;
 
        AnonymousPipeServerStream? readyPipeServer;
 
        public ServiceLifetime(Action stopService, Func<ISignalChecker, IWatchdog> watchdogFactory, ILogger<ServiceLifetime> logger, string[] args)
        {
            ArgumentNullException.ThrowIfNull(stopService);
            ArgumentNullException.ThrowIfNull(watchdogFactory);
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
            ArgumentNullException.ThrowIfNull(args);
 
            logger.LogDebug("Starting service lifetime as user: {username}", Environment.UserName);
 
            cancellationTokenSource = new CancellationTokenSource();
            watchdogTask = RunWatchdog(
                stopService,
                watchdogFactory(this),
                args,
                cancellationTokenSource.Token);
        }
 
        public async ValueTask DisposeAsync()
        {
            cancellationTokenSource.Cancel();
            await watchdogTask;
            cancellationTokenSource.Dispose();
 
            if (commandPipeServer != null)
                await commandPipeServer.DisposeAsync();
 
            if (readyPipeServer != null)
                await readyPipeServer.DisposeAsync();
        }
 
        public async ValueTask CheckSignals(Func<string, (int, Task)> startChildAndGetPid, CancellationToken cancellationToken)
        {
            try
            {
                await using (commandPipeServer = new AnonymousPipeServerStream(PipeDirection.Out, HandleInheritability.Inheritable))
                await using (readyPipeServer = new AnonymousPipeServerStream(PipeDirection.In, HandleInheritability.Inheritable))
                {
                    var (_, lifetimeTask) = startChildAndGetPid($"--Internal:CommandPipe={commandPipeServer.GetClientHandleAsString()} --Internal:ReadyPipe={readyPipeServer.GetClientHandleAsString()}");
                    commandPipeServer.DisposeLocalCopyOfClientHandle();
                    readyPipeServer.DisposeLocalCopyOfClientHandle();
                    await lifetimeTask;
                }
            }
            finally
            {
                readyPipeServer = null;
                commandPipeServer = null;
            }
        }
 
        public void HandleCustomCommand(int command)
        {
            var commandsToCheck = PipeCommands.AllCommands;
            foreach (var stringCommand in commandsToCheck)
            {
                var commandId = PipeCommands.GetServiceCommandId(stringCommand);
                if (command == commandId)
                {
                    SendCommandToHostThroughPipe(stringCommand);
                    return;
                }
            }
 
            logger.LogWarning("Received unknown service command: {command}", command);
        }
 
        async Task RunWatchdog(Action stopService, IWatchdog watchdog, string[] args, CancellationToken cancellationToken)
        {
            var localWatchdogTask = watchdog.RunAsync(false, args, cancellationToken);
 
            if (!localWatchdogTask.IsCompleted && (await watchdog.InitialHostVersion) >= new Version(5, 14, 0))
                if (readyPipeServer != null)
                {
                    logger.LogInformation("Waiting for host to finish starting...");
                    using var streamReader = new StreamReader(
                        readyPipeServer,
                        Encoding.UTF8,
                        leaveOpen: true);
 
                    var line = streamReader.ReadLine(); // Intentionally blocking service startup
                    logger.LogDebug("Pipe read: {line}", line);
 
                    // Maybe we'll use this pipe more in the future, but for now leaving it open is just a resource waste
                    readyPipeServer.Dispose();
                }
                else
                    logger.LogError("Watchdog started and ready pipe was not initialized!");
 
            await localWatchdogTask;
 
            async void StopServiceAsync()
            {
                try
                {
                    // This can call OnStop which waits on this task to complete, must be threaded off or it will deadlock
                    await Task.Run(stopService, cancellationToken);
                }
                catch (OperationCanceledException ex)
                {
                    logger.LogDebug(ex, "Stopping service cancelled!");
                }
                catch (Exception ex)
                {
                    logger.LogError(ex, "Error stopping service!");
                }
            }
 
            StopServiceAsync();
        }
 
        void SendCommandToHostThroughPipe(string command)
        {
            var localPipeServer = commandPipeServer;
            if (localPipeServer == null)
            {
                logger.LogWarning("Unable to send command \"{command}\" to main server process. Is the service running?", command);
                return;
            }
 
            logger.LogDebug("Send command: {command}", command);
            try
            {
                var encoding = Encoding.UTF8;
                using var streamWriter = new StreamWriter(
                    localPipeServer,
                    encoding,
                    PipeCommands
                        .AllCommands
                        .Select(
                            command => encoding.GetByteCount(
                                command + Environment.NewLine))
                        .Max(),
                    true);
                streamWriter.WriteLine(command);
            }
            catch (Exception ex)
            {
                logger.LogError(ex, "Error attempting to send command \"{command}\"", command);
            }
        }
    }
}