ProcessTerminationHandler.cs

// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;

namespace System.CommandLine.Invocation;

internal sealed class ProcessTerminationHandler : IDisposable
{
    private const int SIGINT_EXIT_CODE = 130;
    private const int SIGTERM_EXIT_CODE = 143;
        
    internal readonly TaskCompletionSource<int> ProcessTerminationCompletionSource;
    private readonly CancellationTokenSource _handlerCancellationTokenSource;
    private Task<int>? _startedHandler;
    private readonly TimeSpan _processTerminationTimeout;
#if NET7_0_OR_GREATER
    private readonly IDisposable? _sigIntRegistration, _sigTermRegistration;
#endif

    internal Task<int> StartedHandler { set => Volatile.Write(ref _startedHandler, value); }

    internal ProcessTerminationHandler(
        CancellationTokenSource handlerCancellationTokenSource, 
        TimeSpan processTerminationTimeout)
    {
        ProcessTerminationCompletionSource = new ();
        _handlerCancellationTokenSource = handlerCancellationTokenSource;
        _processTerminationTimeout = processTerminationTimeout;

#if NET7_0_OR_GREATER // we prefer the new API as they allow for cancelling SIGTERM
        if (!OperatingSystem.IsAndroid() 
            && !OperatingSystem.IsIOS() 
            && !OperatingSystem.IsTvOS()
            && !OperatingSystem.IsBrowser())
        {
            _sigIntRegistration = PosixSignalRegistration.Create(PosixSignal.SIGINT, OnPosixSignal);
            _sigTermRegistration = PosixSignalRegistration.Create(PosixSignal.SIGTERM, OnPosixSignal);
            return;
        }
#endif

        Console.CancelKeyPress += OnCancelKeyPress;
        AppDomain.CurrentDomain.ProcessExit += OnProcessExit;
    }

    public void Dispose()
    {
#if NET7_0_OR_GREATER
        if (_sigIntRegistration is not null)
        {
            _sigIntRegistration.Dispose();
            _sigTermRegistration!.Dispose();
            return;
        }
#endif

        Console.CancelKeyPress -= OnCancelKeyPress;
        AppDomain.CurrentDomain.ProcessExit -= OnProcessExit;    
    }
        
#if NET7_0_OR_GREATER
    void OnPosixSignal(PosixSignalContext context)
    {
        context.Cancel = true;
            
        Cancel(context.Signal == PosixSignal.SIGINT ? SIGINT_EXIT_CODE : SIGTERM_EXIT_CODE);
    }
#endif

    void OnCancelKeyPress(object? sender, ConsoleCancelEventArgs e)
    {
        e.Cancel = true;

        Cancel(SIGINT_EXIT_CODE);
    }

    void OnProcessExit(object? sender, EventArgs e) => Cancel(SIGTERM_EXIT_CODE);

    void Cancel(int forcedTerminationExitCode)
    {
        // request cancellation
        _handlerCancellationTokenSource.Cancel();

        try
        {
            var startedHandler = Volatile.Read(ref _startedHandler);
            // wait for the configured interval
            if (startedHandler is null || !startedHandler.Wait(_processTerminationTimeout))
            {
                // if the handler does not finish within configured time,
                // use the completion source to signal forced completion (preserving native exit code)
                ProcessTerminationCompletionSource.SetResult(forcedTerminationExitCode);
            }
        }
        catch (AggregateException)
        {
            // The task was cancelled or an exception was thrown during the task execution.
        }
    }
}