ProcessExecutor.cs

Go to the documentation of this file.

using System;
using System.IO;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Host.Extensions;
using Tgstation.Server.Host.IO;
 
namespace Tgstation.Server.Host.System
{
    sealed class ProcessExecutor : IProcessExecutor
    {
        readonly IProcessFeatures processFeatures;
 
        readonly IIOManager ioManager;
 
        readonly ILogger<ProcessExecutor> logger;
 
        readonly ILoggerFactory loggerFactory;
 
        public ProcessExecutor(
            IProcessFeatures processFeatures,
            IIOManager ioManager,
            ILogger<ProcessExecutor> logger,
            ILoggerFactory loggerFactory)
        {
            this.processFeatures = processFeatures ?? throw new ArgumentNullException(nameof(processFeatures));
            this.ioManager = ioManager ?? throw new ArgumentNullException(nameof(ioManager));
            this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
            this.loggerFactory = loggerFactory ?? throw new ArgumentNullException(nameof(loggerFactory));
        }
 
        public IProcess GetProcess(int id)
        {
            logger.LogDebug("Attaching to process {pid}...", id);
            global::System.Diagnostics.Process handle;
            try
            {
                handle = global::System.Diagnostics.Process.GetProcessById(id);
            }
            catch (Exception e)
            {
                logger.LogDebug(e, "Unable to get process {pid}!", id);
                return null;
            }
 
            return CreateFromExistingHandle(handle);
        }
 
        public IProcess GetCurrentProcess()
        {
            logger.LogTrace("Getting current process...");
            var handle = global::System.Diagnostics.Process.GetCurrentProcess();
            return CreateFromExistingHandle(handle);
        }
 
        public async Task<IProcess> LaunchProcess(
            string fileName,
            string workingDirectory,
            string arguments,
            string fileRedirect,
            bool readStandardHandles,
            bool noShellExecute)
        {
            ArgumentNullException.ThrowIfNull(fileName);
            ArgumentNullException.ThrowIfNull(workingDirectory);
            ArgumentNullException.ThrowIfNull(arguments);
 
            if (!noShellExecute && readStandardHandles)
                throw new InvalidOperationException("Requesting output/error reading requires noShellExecute to be true!");
 
            logger.LogDebug(
                "{launchType}aunching process in {workingDirectory}: {exe} {arguments}",
                noShellExecute ? "L" : "Shell l",
                workingDirectory,
                fileName,
                arguments);
            var handle = new global::System.Diagnostics.Process();
            try
            {
                handle.StartInfo.FileName = fileName;
                handle.StartInfo.Arguments = arguments;
                handle.StartInfo.WorkingDirectory = workingDirectory;
 
                handle.StartInfo.UseShellExecute = !noShellExecute;
 
                var processStartTcs = new TaskCompletionSource();
                var lifetimeTaskTask = AttachExitHandlerBeforeLaunch(handle, processStartTcs.Task);
 
                Task<string> readTask = null;
                CancellationTokenSource disposeCts = null;
                if (readStandardHandles)
                {
                    handle.StartInfo.RedirectStandardOutput = true;
                    handle.StartInfo.RedirectStandardError = true;
 
                    disposeCts = new CancellationTokenSource();
                    readTask = ConsumeReaders(handle, processStartTcs.Task, fileRedirect, disposeCts.Token);
                }
 
                try
                {
                    handle.Start();
 
                    processStartTcs.SetResult();
                }
                catch (Exception ex)
                {
                    processStartTcs.SetException(ex);
                    throw;
                }
 
                var process = new Process(
                    processFeatures,
                    handle,
                    disposeCts,
                    await lifetimeTaskTask, // won't block
                    readTask,
                    loggerFactory.CreateLogger<Process>(),
                    false);
 
                return process;
            }
            catch
            {
                handle.Dispose();
                throw;
            }
        }
 
        public IProcess GetProcessByName(string name)
        {
            logger.LogTrace("GetProcessByName: {processName}...", name ?? throw new ArgumentNullException(nameof(name)));
            var procs = global::System.Diagnostics.Process.GetProcessesByName(name);
            global::System.Diagnostics.Process handle = null;
            foreach (var proc in procs)
                if (handle == null)
                    handle = proc;
                else
                {
                    logger.LogTrace("Disposing extra found PID: {pid}...", proc.Id);
                    proc.Dispose();
                }
 
            if (handle == null)
                return null;
 
            return CreateFromExistingHandle(handle);
        }
 
        async Task<Task<int>> AttachExitHandlerBeforeLaunch(global::System.Diagnostics.Process handle, Task startupTask)
        {
            var id = -1;
            var result = AttachExitHandler(handle, () => id);
            await startupTask;
            id = handle.Id;
            return result;
        }
 
        Task<int> AttachExitHandler(global::System.Diagnostics.Process handle, Func<int> idProvider)
        {
            handle.EnableRaisingEvents = true;
 
            var tcs = new TaskCompletionSource<int>();
            void ExitHandler(object sender, EventArgs args)
            {
                var id = idProvider();
                try
                {
                    try
                    {
                        var exitCode = handle.ExitCode;
 
                        // Try because this can be invoked twice for weird reasons
                        if (tcs.TrySetResult(exitCode))
                            logger.LogTrace("PID {pid} termination event completed", id);
                        else
                            logger.LogTrace("Ignoring duplicate PID {pid} termination event", id);
                    }
                    catch (InvalidOperationException ex)
                    {
                        if (!tcs.Task.IsCompleted)
                            throw;
 
                        logger.LogTrace(ex, "Ignoring expected PID {pid} exit handler exception!", id);
                    }
                }
                catch (Exception ex)
                {
                    logger.LogError(ex, "PID {pid} exit handler exception!", id);
                }
            }
 
            handle.Exited += ExitHandler;
 
            return tcs.Task;
        }
 
        async Task<string> ConsumeReaders(global::System.Diagnostics.Process handle, Task startTask, string fileRedirect, CancellationToken disposeToken)
        {
            await startTask;
 
            var pid = handle.Id;
            logger.LogTrace("Starting read for PID {pid}...", pid);
 
            var stdOutHandle = handle.StandardOutput;
            var stdErrHandle = handle.StandardError;
            Task<string> outputReadTask = null, errorReadTask = null;
            bool outputOpen = true, errorOpen = true;
            async Task<string> GetNextLine()
            {
                if (outputOpen && outputReadTask == null)
                    outputReadTask = stdOutHandle.ReadLineAsync();
 
                if (errorOpen && errorReadTask == null)
                    errorReadTask = stdErrHandle.ReadLineAsync();
 
                var completedTask = await Task.WhenAny(outputReadTask ?? errorReadTask, errorReadTask ?? outputReadTask).WithToken(disposeToken);
                var line = await completedTask;
                if (completedTask == outputReadTask)
                {
                    outputReadTask = null;
                    if (line == null)
                        outputOpen = false;
                }
                else
                {
                    errorReadTask = null;
                    if (line == null)
                        errorOpen = false;
                }
 
                if (line == null && (errorOpen || outputOpen))
                    return await GetNextLine();
 
                return line;
            }
 
            await using var fileStream = fileRedirect != null ? ioManager.CreateAsyncSequentialWriteStream(fileRedirect) : null;
            await using var writer = fileStream != null ? new StreamWriter(fileStream) : null;
 
            string text;
            var stringBuilder = fileStream == null ? new StringBuilder() : null;
            try
            {
                while ((text = await GetNextLine()) != null)
                {
                    if (fileStream != null)
                    {
                        await writer.WriteLineAsync(text);
                        await writer.FlushAsync();
                    }
                    else
                        stringBuilder.AppendLine(text);
                }
 
                logger.LogTrace("Finished read for PID {pid}", pid);
            }
            catch (OperationCanceledException ex)
            {
                logger.LogWarning(ex, "PID {pid} stream reading interrupted!", pid);
                if (fileStream != null)
                    await writer.WriteLineAsync("-- Process detached, log truncated. This is likely due a to TGS restart --");
            }
 
            return stringBuilder?.ToString();
        }
 
        IProcess CreateFromExistingHandle(global::System.Diagnostics.Process handle)
        {
            try
            {
                var pid = handle.Id;
                return new Process(
                    processFeatures,
                    handle,
                    null,
                    AttachExitHandler(handle, () => pid),
                    null,
                    loggerFactory.CreateLogger<Process>(),
                    true);
            }
            catch
            {
                handle.Dispose();
                throw;
            }
        }
    }
}