lmxopcua/src/Core/ZB.MOM.WW.OtOpcUa.Core.Scripting/ScriptRunGuard.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using System.Diagnostics.Metrics;
using Serilog;
using ZB.MOM.WW.OtOpcUa.Commons.Observability;

namespace ZB.MOM.WW.OtOpcUa.Core.Scripting;

/// <summary>
///     Per-script run quarantine shared by every script seam on the node. Wraps each evaluation of a
///     <see cref="ScriptEvaluator{TContext, TResult}"/> with the wall-clock timeout
///     <see cref="TimedScriptEvaluator{TContext, TResult}"/> used to provide, and adds what that wrapper
///     could not: knowledge of whether the script it abandoned is <em>still running</em>.
/// </summary>
/// <remarks>
///     <para>
///         <b>The leak this closes.</b> A timed-out script is abandoned, not stopped — a CPU-bound
///         <c>while(true)</c> keeps its thread-pool thread forever. Each seam re-evaluated on every trigger,
///         so a runaway script on a 1 Hz input leaked one pool thread per second until the pool starved and
///         unrelated work (Akka dispatchers, the OPC UA publish thread, gRPC) stalled. The guard keys every
///         run by script instance and:
///     </para>
///     <list type="number">
///         <item><b>Refuses to start a run while that script's previous run is still executing</b> —
///         <see cref="ScriptRunRefusedException"/> (<see cref="PreviousRunStillRunningReason"/>), counted on
///         <see cref="RefusedWhileRunningCount"/>. One orphan per script, never one per trigger.</item>
///         <item><b>Quarantines after <see cref="ScriptingOptions.QuarantineAfterConsecutiveTimeouts"/>
///         consecutive timeouts</b> (only a timeout is a strike; a refusal is a Bad publish plus its counter
///         and leaves the streak untouched — on a 1 Hz trigger with a 2 s budget one 4 s overrun is a
///         timeout plus two refusals and must not quarantine; a run that returns — a value or a throw —
///         resets the streak), <b>or on a refusal behind a stuck orphan</b> — one that has been executing for
///         longer than <see cref="ScriptingOptions.StuckOrphanQuarantineFactor"/> × the timeout (20 s at the
///         2 s default), which is how a never-returning <c>while(true)</c> (one timeout, then only refusals)
///         gets quarantined: no further dispatch, every attempt throws
///         <see cref="ScriptQuarantinedException"/>, one Error log, one <c>otopcua.script.quarantined</c>
///         increment, one <see cref="Quarantined"/> event. Lifted by a source change (the operator fixed the
///         script — a new source identity under the same key starts clean and the old orphan no longer blocks
///         it), by <see cref="ResetAll"/> on redeploy, or by <see cref="Reset"/>.</item>
///         <item><b>Quarantines on an allocation-budget breach</b>: the run's own-thread allocation is sampled
///         with <see cref="GC.GetAllocatedBytesForCurrentThread"/> around the script delegate and compared
///         with <see cref="ScriptingOptions.MaxAllocBytesPerRun"/>. This is observation after the fact — a
///         single run that allocates enough to OOM the process is not stopped; that needs an out-of-process
///         runner.</item>
///     </list>
///     <para>
///         <b>The timeout clock starts when the script starts, not when it is queued.</b> The run is pushed to
///         the pool (a CPU-bound script would otherwise execute inline and return an already-completed task
///         that no timeout can fire on), and the wait only begins once the pool has actually picked the run
///         up — so thread-pool starvation is charged to nobody's script. A caller-supplied cancellation wins
///         over the timeout (<see cref="OperationCanceledException"/>, never
///         <see cref="ScriptTimeoutException"/>) so shutdown paths are not misread as timeouts.
///     </para>
///     <para>
///         <b>Cancellation reaches the evaluator, not (yet) the script.</b> On timeout the token handed to
///         <see cref="ScriptEvaluator{TContext, TResult}.RunAsync"/> is cancelled; a script can only observe
///         it once <c>ScriptContext</c> exposes it (follow-up — the context type lives in the Abstractions
///         project).
///     </para>
/// </remarks>
public sealed class ScriptRunGuard
{
    /// <summary>Failure reason a seam publishes when a run was refused because the previous one is still executing.</summary>
    public const string PreviousRunStillRunningReason = "previous evaluation still running";

    /// <summary>Prefix of every quarantine reason / <see cref="ScriptQuarantinedException"/> message.</summary>
    public const string QuarantinedReasonPrefix = "script quarantined: ";

    /// <summary>
    ///     Counter <c>otopcua.script.quarantined</c> (Prometheus: <c>otopcua_script_quarantined</c>): one
    ///     increment per script quarantined, tagged <c>reason</c> = <c>timeouts</c> | <c>stuck</c> | <c>allocation</c>.
    /// </summary>
    public static readonly Counter<long> QuarantinedCounter =
        OtOpcUaTelemetry.Meter.CreateCounter<long>("otopcua.script.quarantined", unit: "{script}",
            description: "Scripts quarantined by the run guard (reason: timeouts | stuck | allocation).");

    /// <summary>
    ///     Counter <c>otopcua.script.run_refused</c>: runs refused because the script's previous run was
    ///     still executing.
    /// </summary>
    public static readonly Counter<long> RefusedCounter =
        OtOpcUaTelemetry.Meter.CreateCounter<long>("otopcua.script.run_refused", unit: "{run}",
            description: "Script runs refused because the previous run of the same script was still executing.");

    private readonly ConcurrentDictionary<string, KeyState> _keys = new(StringComparer.Ordinal);
    private readonly long _maxAllocBytesPerRun;
    private readonly int _quarantineThreshold;
    private readonly double _stuckFactor;
    private readonly ILogger _logger;
    private int _inFlight;
    private long _started;
    private long _refused;
    private long _quarantined;

    /// <summary>Initializes a new guard.</summary>
    /// <param name="options">Budget + threshold; null ⇒ <see cref="ScriptingOptions"/> defaults.</param>
    /// <param name="logger">Where the quarantine Error line goes; null ⇒ Serilog's global logger.</param>
    public ScriptRunGuard(ScriptingOptions? options = null, ILogger? logger = null)
    {
        // Clamped here as well as at the bind site, because options also arrive from a Calculation driver's own
        // config and from tests. An unclamped factor made the refusal path's timeout * factor overflow.
        options = (options ?? new ScriptingOptions()).Sanitized(out _);
        _maxAllocBytesPerRun = options.MaxAllocBytesPerRun;
        _quarantineThreshold = options.QuarantineAfterConsecutiveTimeouts;
        _stuckFactor = options.StuckOrphanQuarantineFactor;
        _logger = logger ?? Log.Logger;
    }

    /// <summary>Raised (on the thread that detected it) when a script is quarantined.</summary>
    public event EventHandler<ScriptQuarantinedEventArgs>? Quarantined;

    /// <summary>Runs currently executing on a pool thread — abandoned (timed-out) orphans included.</summary>
    public int InFlightCount => Volatile.Read(ref _inFlight);

    /// <summary>Runs actually dispatched to the pool, lifetime.</summary>
    public long StartedRunCount => Interlocked.Read(ref _started);

    /// <summary>Runs refused because the same script's previous run was still executing, lifetime.</summary>
    public long RefusedWhileRunningCount => Interlocked.Read(ref _refused);

    /// <summary>Scripts quarantined, lifetime (a reset and re-quarantine counts twice).</summary>
    public long QuarantinedCount => Interlocked.Read(ref _quarantined);

    /// <summary>The keys currently quarantined.</summary>
    public IReadOnlyCollection<string> QuarantinedKeys =>
        _keys.Where(kv => kv.Value.QuarantineReason is not null).Select(kv => kv.Key).ToArray();

    /// <summary>Whether <paramref name="key"/> is quarantined, and why.</summary>
    /// <param name="key">The script instance key.</param>
    /// <param name="reason">The quarantine reason (prefixed <see cref="QuarantinedReasonPrefix"/>), or null.</param>
    /// <returns><c>true</c> when quarantined.</returns>
    public bool IsQuarantined(string key, out string? reason)
    {
        reason = _keys.TryGetValue(key, out var state) ? state.QuarantineReason : null;
        return reason is not null;
    }

    /// <summary>
    ///     Forget everything about <paramref name="key"/>: clears its quarantine and its consecutive count, and
    ///     abandons its in-flight slot (an orphan still running keeps counting on <see cref="InFlightCount"/>
    ///     until it finishes, but no longer blocks the key).
    /// </summary>
    /// <param name="key">The script instance key.</param>
    public void Reset(string key) => _keys.TryRemove(key, out _);

    /// <summary>
    ///     <see cref="Reset"/> for every key — the redeploy / compile-cache-clear boundary. Orphans keep
    ///     counting on <see cref="InFlightCount"/> until they finish.
    /// </summary>
    public void ResetAll() => _keys.Clear();

    /// <summary>
    ///     Run <paramref name="evaluator"/> against <paramref name="context"/> under the guard.
    /// </summary>
    /// <typeparam name="TContext">The script context type.</typeparam>
    /// <typeparam name="TResult">The script result type.</typeparam>
    /// <param name="key">The script instance key (virtual-tag id, calc-tag RawPath, alarm id).</param>
    /// <param name="sourceIdentity">
    ///     The script source (or its hash). A change under the same key clears the key's quarantine and
    ///     abandons its orphan slot — the operator fixed the script. Null ⇒ not tracked.
    /// </param>
    /// <param name="evaluator">The compiled script.</param>
    /// <param name="context">The script's <c>ctx</c>.</param>
    /// <param name="timeout">Wall-clock budget, measured from the moment the pool starts the run.</param>
    /// <param name="ct">Caller cancellation; wins over the timeout.</param>
    /// <returns>The script's result.</returns>
    /// <exception cref="ScriptQuarantinedException">The key is (or just became) quarantined.</exception>
    /// <exception cref="ScriptRunRefusedException">The key's previous run is still executing.</exception>
    /// <exception cref="ScriptTimeoutException">The run exceeded <paramref name="timeout"/>.</exception>
    /// <exception cref="OperationCanceledException"><paramref name="ct"/> was cancelled.</exception>
    public async Task<TResult> RunAsync<TContext, TResult>(
        string key,
        string? sourceIdentity,
        ScriptEvaluator<TContext, TResult> evaluator,
        TContext context,
        TimeSpan timeout,
        CancellationToken ct = default)
        where TContext : ScriptContext
    {
        ArgumentNullException.ThrowIfNull(key);
        ArgumentNullException.ThrowIfNull(evaluator);
        ArgumentNullException.ThrowIfNull(context);
        if (timeout <= TimeSpan.Zero) throw new ArgumentOutOfRangeException(nameof(timeout), "Timeout must be positive.");

        var state = Admit(key, sourceIdentity);

        // Slot acquisition is the whole mechanism: a second trigger landing while the previous run (orphan or
        // not) still holds the slot is refused here, before anything is dispatched to the pool.
        if (Interlocked.CompareExchange(ref state.InFlight, 1, 0) != 0)
        {
            // A refusal is a Bad publish plus its counter — NOT a strike. Only a timeout (and an allocation
            // breach) counts toward quarantine: on a 1 Hz trigger with a 2 s budget a single 4 s overrun is
            // one timeout followed by two refusals, and counting those would quarantine a script that merely
            // ran long once. The refusal also leaves the timeout streak untouched.
            Interlocked.Increment(ref _refused);
            RefusedCounter.Add(1, new KeyValuePair<string, object?>("script", key));
            // …unless the orphan is STUCK: a never-returning script times out once and is then only ever
            // refused, so the timeout streak alone would never quarantine it. Past the stuck threshold the
            // refusal itself quarantines (same path: Error log, metric, Bad, lifted on redeploy/source change).
            if (_stuckFactor > 0 && Volatile.Read(ref state.RunStartedTimestamp) is var startedAt and not 0)
            {
                var running = Stopwatch.GetElapsedTime(startedAt);
                // Saturating: even the clamped maximum factor times a huge timeout must not overflow TimeSpan.
                var thresholdTicks = timeout.Ticks * _stuckFactor;
                var threshold = thresholdTicks >= TimeSpan.MaxValue.Ticks
                    ? TimeSpan.MaxValue
                    : TimeSpan.FromTicks((long)thresholdTicks);
                if (running > threshold)
                {
                    Quarantine(key, state,
                        $"its previous evaluation has been running for {running.TotalSeconds:F1} s, over " +
                        $"{_stuckFactor:0.##}x its {timeout.TotalMilliseconds:F0} ms budget " +
                        "(Scripting:StuckOrphanQuarantineFactor) — it is stuck; the script is not evaluated again " +
                        "until its source changes or it is redeployed",
                        "stuck");
                    throw new ScriptQuarantinedException(key, state.QuarantineReason!);
                }
            }
            throw new ScriptRunRefusedException(key);
        }

        Interlocked.Increment(ref _inFlight);
        Interlocked.Increment(ref _started);

        var slot = new RunSlot();
        var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var runCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        var runToken = runCts.Token;

        // Task.Run deliberately takes NO token: a pre-cancelled token would skip the lambda, `started` would
        // never complete, and the slot would never be released. The evaluator observes the token instead.
        var runTask = Task.Run(async () =>
        {
            Volatile.Write(ref state.RunStartedTimestamp, Stopwatch.GetTimestamp());
            started.TrySetResult();
            var before = GC.GetAllocatedBytesForCurrentThread();
            try
            {
                // ScriptEvaluator.RunAsync executes the compiled delegate synchronously and returns a completed
                // task, so the after-sample is taken on the same thread, right after the script returned.
                var pending = evaluator.RunAsync(context, runToken);
                slot.AllocatedBytes = GC.GetAllocatedBytesForCurrentThread() - before;
                return await pending.ConfigureAwait(false);
            }
            catch
            {
                slot.AllocatedBytes = GC.GetAllocatedBytesForCurrentThread() - before;
                throw;
            }
        });

        // Release the slot when the run REALLY finishes — which for an abandoned run is long after this
        // method has thrown ScriptTimeoutException. Registered before the wait so it fires regardless of
        // which way the wait below comes out. A late-finishing orphan that breached the allocation budget
        // still quarantines its key here.
        _ = runTask.ContinueWith(t =>
        {
            Volatile.Write(ref state.RunStartedTimestamp, 0);
            Interlocked.Exchange(ref state.InFlight, 0);
            Interlocked.Decrement(ref _inFlight);
            runCts.Dispose();
            if (t.IsFaulted) _ = t.Exception; // observed below by the awaiter when it is still listening
            if (t.IsCompletedSuccessfully && BreachedBudget(slot.AllocatedBytes))
                Quarantine(key, state, AllocationReason(slot.AllocatedBytes), "allocation");
        }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);

        // The budget clock starts when the pool starts the run, not when we queued it (P5): pool starvation
        // is a host problem and must not be reported as "script timed out".
        await started.Task.WaitAsync(ct).ConfigureAwait(false);

        TResult result;
        try
        {
            result = await runTask.WaitAsync(timeout, ct).ConfigureAwait(false);
        }
        catch (TimeoutException)
        {
            // The caller's cancel wins even when it lands a few µs after the timeout fired.
            if (ct.IsCancellationRequested) throw new OperationCanceledException(ct);
            runCts.Cancel();
            if (NoteTimeout(key, state, timeout))
                throw new ScriptQuarantinedException(key, state.QuarantineReason!);
            throw new ScriptTimeoutException(timeout);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception)
        {
            // The script returned control (by throwing) — that is a completion for quarantine purposes.
            Interlocked.Exchange(ref state.ConsecutiveTimeouts, 0);
            throw;
        }

        Interlocked.Exchange(ref state.ConsecutiveTimeouts, 0);
        if (BreachedBudget(slot.AllocatedBytes))
        {
            // The continuation above has already (or will have) quarantined the key; make the reason visible
            // to THIS caller too, so the breaching run itself publishes Bad rather than a suspect value.
            Quarantine(key, state, AllocationReason(slot.AllocatedBytes), "allocation");
            throw new ScriptQuarantinedException(key, state.QuarantineReason!);
        }
        return result;
    }

    /// <summary>Resolve the key's state, rolling it over when the source identity changed, and refuse a quarantined key.</summary>
    private KeyState Admit(string key, string? sourceIdentity)
    {
        while (true)
        {
            var state = _keys.GetOrAdd(key, static (_, src) => new KeyState(src), sourceIdentity);
            if (sourceIdentity is not null && state.SourceIdentity is not null
                && !string.Equals(state.SourceIdentity, sourceIdentity, StringComparison.Ordinal))
            {
                // The script changed under this key: start from a clean slate. The old state's orphan (if any)
                // releases the GLOBAL in-flight count through its own continuation; nothing else references it.
                _keys.TryUpdate(key, new KeyState(sourceIdentity), state);
                continue;
            }
            if (sourceIdentity is not null && state.SourceIdentity is null)
            {
                state.SourceIdentity = sourceIdentity;
            }
            if (state.QuarantineReason is { } reason)
                throw new ScriptQuarantinedException(key, reason);
            return state;
        }
    }

    /// <summary>Record a timeout strike; returns true when it tipped the key into quarantine. A completed run
    /// (value or throw) resets the streak; a refusal neither counts nor resets.</summary>
    private bool NoteTimeout(string key, KeyState state, TimeSpan timeout)
    {
        var count = Interlocked.Increment(ref state.ConsecutiveTimeouts);
        if (count < _quarantineThreshold) return false;
        Quarantine(key, state,
            $"{count} consecutive evaluation(s) timed out (budget {timeout.TotalMilliseconds:F0} ms); " +
            "the script is not evaluated again until its source changes or it is redeployed",
            "timeouts");
        return true;
    }

    private bool BreachedBudget(long allocatedBytes) =>
        _maxAllocBytesPerRun > 0 && allocatedBytes > _maxAllocBytesPerRun;

    private string AllocationReason(long allocatedBytes) =>
        $"one evaluation allocated {allocatedBytes / (1024.0 * 1024.0):F1} MB, over the " +
        $"{_maxAllocBytesPerRun / (1024.0 * 1024.0):F0} MB budget (Scripting:MaxAllocBytesPerRun); " +
        "the script is not evaluated again until its source changes or it is redeployed";

    private void Quarantine(string key, KeyState state, string why, string metricReason)
    {
        var reason = QuarantinedReasonPrefix + why;
        // First writer wins: a late orphan and the awaiter can both reach here for the same run.
        if (Interlocked.CompareExchange(ref state.QuarantineReason, reason, null) is not null) return;

        Interlocked.Increment(ref _quarantined);
        QuarantinedCounter.Add(1,
            new KeyValuePair<string, object?>("script", key),
            new KeyValuePair<string, object?>("reason", metricReason));
        _logger.Error("Script {ScriptKey} quarantined: {Reason}", key, why);
        try { Quarantined?.Invoke(this, new ScriptQuarantinedEventArgs(key, reason)); }
        catch (Exception ex)
        {
            _logger.Warning(ex, "ScriptRunGuard Quarantined subscriber threw for {ScriptKey}", key);
        }
    }

    private sealed class KeyState(string? sourceIdentity)
    {
        public string? SourceIdentity = sourceIdentity;
        public int InFlight;
        public int ConsecutiveTimeouts;
        /// <summary><see cref="Stopwatch.GetTimestamp"/> when the pool started the in-flight run; 0 when idle.</summary>
        public long RunStartedTimestamp;
        public string? QuarantineReason;
    }

    private sealed class RunSlot
    {
        public long AllocatedBytes;
    }
}

/// <summary>Carried by <see cref="ScriptRunGuard.Quarantined"/>.</summary>
/// <param name="Key">The quarantined script instance key.</param>
/// <param name="Reason">Why, prefixed <see cref="ScriptRunGuard.QuarantinedReasonPrefix"/>.</param>
public sealed record ScriptQuarantinedEventArgs(string Key, string Reason);

/// <summary>
///     Thrown by <see cref="ScriptRunGuard.RunAsync{TContext, TResult}"/> when the script's previous run is
///     still executing (an abandoned, timed-out run that never returned). Nothing was dispatched; the seam
///     publishes Bad with <see cref="ScriptRunGuard.PreviousRunStillRunningReason"/>.
/// </summary>
public sealed class ScriptRunRefusedException : Exception
{
    /// <summary>The script instance key.</summary>
    public string Key { get; }

    /// <summary>Initializes the exception for <paramref name="key"/>.</summary>
    /// <param name="key">The script instance key.</param>
    public ScriptRunRefusedException(string key) : base(ScriptRunGuard.PreviousRunStillRunningReason)
    {
        Key = key;
    }
}

/// <summary>
///     Thrown by <see cref="ScriptRunGuard.RunAsync{TContext, TResult}"/> for a quarantined script: repeated
///     timeouts, a stuck orphan or an allocation-budget breach. Nothing is dispatched while quarantined; the seam
///     publishes Bad carrying the reason, which an operator clears by fixing the script (source change) or
///     redeploying.
/// </summary>
public sealed class ScriptQuarantinedException : Exception
{
    /// <summary>The script instance key.</summary>
    public string Key { get; }

    /// <summary>Initializes the exception with the quarantine reason (already prefixed).</summary>
    /// <param name="key">The script instance key.</param>
    /// <param name="reason">The quarantine reason, prefixed <see cref="ScriptRunGuard.QuarantinedReasonPrefix"/>.</param>
    public ScriptQuarantinedException(string key, string reason) : base(reason)
    {
        Key = key;
    }
}