ProcessExecutor.cs

Go to the documentation of this file.

using System;
using System.Diagnostics;
using System.IO;
using System.Text;
using System.Threading;
using System.Threading.Channels;
using System.Threading.Tasks;
 
using Microsoft.Extensions.Logging;
 
using Tgstation.Server.Host.IO;
 
namespace Tgstation.Server.Host.System
{
    sealed class ProcessExecutor : IProcessExecutor
    {
        static readonly ReaderWriterLockSlim ExclusiveProcessLaunchLock = new();
 
        readonly IProcessFeatures processFeatures;
 
        readonly IIOManager ioManager;
 
        readonly ILogger<ProcessExecutor> logger;
 
        readonly ILoggerFactory loggerFactory;
 
        public static void WithProcessLaunchExclusivity(Action action)
        {
            ExclusiveProcessLaunchLock.EnterWriteLock();
            try
            {
                action();
            }
            finally
            {
                ExclusiveProcessLaunchLock.ExitWriteLock();
            }
        }
 
        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 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)
                logger.LogDebug(
                "Launching process in {workingDirectory}: {exe} {arguments}",
                workingDirectory,
                fileName,
                arguments);
            else
                logger.LogDebug(
                "Shell launching process in {workingDirectory}: {exe} {arguments}",
                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;
 
                Task<string?>? readTask = null;
                CancellationTokenSource? disposeCts = null;
                try
                {
                    TaskCompletionSource<int>? processStartTcs = null;
                    if (readStandardHandles)
                    {
                        processStartTcs = new TaskCompletionSource<int>();
                        disposeCts = new CancellationTokenSource();
                        readTask = ConsumeReaders(handle, processStartTcs.Task, fileRedirect, disposeCts.Token);
                    }
 
                    int pid;
                    try
                    {
                        ExclusiveProcessLaunchLock.EnterReadLock();
                        try
                        {
                            handle.Start();
                        }
                        finally
                        {
                            ExclusiveProcessLaunchLock.ExitReadLock();
                        }
 
                        pid = handle.Id;
                        processStartTcs?.SetResult(pid);
                    }
                    catch (Exception ex)
                    {
                        processStartTcs?.SetException(ex);
                        throw;
                    }
 
                    var process = new Process(
                        processFeatures,
                        handle,
                        disposeCts,
                        readTask,
                        loggerFactory.CreateLogger<Process>(),
                        false);
 
                    return process;
                }
                catch
                {
                    disposeCts?.Dispose();
                    throw;
                }
            }
            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<string?> ConsumeReaders(global::System.Diagnostics.Process handle, Task<int> startupAndPid, string? fileRedirect, CancellationToken cancellationToken)
        {
            handle.StartInfo.RedirectStandardOutput = true;
            handle.StartInfo.RedirectStandardError = true;
 
            bool writingToFile;
            await using var fileStream = (writingToFile = fileRedirect != null) ? ioManager.CreateAsyncSequentialWriteStream(fileRedirect!) : null;
            await using var fileWriter = fileStream != null ? new StreamWriter(fileStream) : null;
 
            var stringBuilder = fileStream == null ? new StringBuilder() : null;
 
            var dataChannel = Channel.CreateUnbounded<string>(
                new UnboundedChannelOptions
                {
                    AllowSynchronousContinuations = !writingToFile,
                    SingleReader = true,
                    SingleWriter = false,
                });
 
            var handlesOpen = 2;
            async void DataReceivedHandler(object sender, DataReceivedEventArgs eventArgs)
            {
                var line = eventArgs.Data;
                if (line == null)
                {
                    var handlesRemaining = Interlocked.Decrement(ref handlesOpen);
                    if (handlesRemaining == 0)
                        dataChannel.Writer.Complete();
 
                    return;
                }
 
                try
                {
                    await dataChannel.Writer.WriteAsync(line, cancellationToken);
                }
                catch (OperationCanceledException ex)
                {
                    logger.LogWarning(ex, "Handle channel write interrupted!");
                }
            }
 
            handle.OutputDataReceived += DataReceivedHandler;
            handle.ErrorDataReceived += DataReceivedHandler;
 
            async ValueTask OutputWriter()
            {
                var enumerable = dataChannel.Reader.ReadAllAsync(cancellationToken);
                if (writingToFile)
                {
                    var enumerator = enumerable.GetAsyncEnumerator(cancellationToken);
                    var nextEnumeration = enumerator.MoveNextAsync();
                    while (await nextEnumeration)
                    {
                        var text = enumerator.Current;
                        nextEnumeration = enumerator.MoveNextAsync();
                        await fileWriter!.WriteLineAsync(text.AsMemory(), cancellationToken);
 
                        if (!nextEnumeration.IsCompleted)
                            await fileWriter.FlushAsync(cancellationToken);
                    }
                }
                else
                    await foreach (var text in enumerable)
                        stringBuilder!.AppendLine(text);
            }
 
            var pid = await startupAndPid;
            logger.LogTrace("Starting read for PID {pid}...", pid);
 
            using (cancellationToken.Register(() => dataChannel.Writer.TryComplete()))
            {
                handle.BeginOutputReadLine();
                using (cancellationToken.Register(handle.CancelOutputRead))
                {
                    handle.BeginErrorReadLine();
                    using (cancellationToken.Register(handle.CancelErrorRead))
                    {
                        try
                        {
                            await OutputWriter();
 
                            logger.LogTrace("Finished read for PID {pid}", pid);
                        }
                        catch (OperationCanceledException ex)
                        {
                            logger.LogWarning(ex, "PID {pid} stream reading interrupted!", pid);
                            if (writingToFile)
                                await fileWriter!.WriteLineAsync("-- Process detached, log truncated. This is likely due a to TGS restart --");
                        }
                    }
                }
            }
 
            return stringBuilder?.ToString();
        }
 
        Process CreateFromExistingHandle(global::System.Diagnostics.Process handle)
        {
            try
            {
                var pid = handle.Id;
                return new Process(
                    processFeatures,
                    handle,
                    null,
                    null,
                    loggerFactory.CreateLogger<Process>(),
                    true);
            }
            catch
            {
                handle.Dispose();
                throw;
            }
        }
    }
}