lmxopcua/src/Core/ZB.MOM.WW.OtOpcUa.Core.Scripting.OutOfProcess/WorkerShard.cs

using System.Diagnostics;
using System.Diagnostics.Metrics;
using Serilog;
using ZB.MOM.WW.OtOpcUa.Commons.Observability;
using ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol.V1;

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

/// <summary>A point-in-time view of one shard (diagnostics, the readiness check, the runner-state gauge).</summary>
/// <param name="Label">The shard label (<c>0</c>…<c>K-1</c>, or <c>iso</c>).</param>
/// <param name="IsIsolation">Whether it is the isolation worker.</param>
/// <param name="Up">A worker is live and serving.</param>
/// <param name="DownFor">How long it has been down after a failure; null when up or never needed.</param>
/// <param name="Pid">The live worker's pid.</param>
/// <param name="ConsecutiveFailures">Start failures and crashes since the last healthy period.</param>
/// <param name="InFlight">Evaluations outstanding on the live worker.</param>
/// <param name="Queued">Evaluations queued in the worker (its last heartbeat).</param>
/// <param name="HeapBytes">The worker's managed heap (its last heartbeat).</param>
public sealed record ScriptWorkerShardStatus(
    string Label,
    bool IsIsolation,
    bool Up,
    TimeSpan? DownFor,
    int? Pid,
    int ConsecutiveFailures,
    int InFlight,
    int Queued,
    long HeapBytes);

/// <summary>
///     Supervises one shard's worker process (design §3.3, §5.3.5): starts it lazily on first use, reports its end to
///     the runner (for attribution, synchronously — see <see cref="WorkerConnection.StartAsync"/>), and restarts it after
///     an exponential backoff — 250 ms doubling to a 30 s cap with ±20 % jitter, reset once a worker has stayed up for
///     60 s. It never runs a script in the host when the worker cannot start (fail closed, §14 #4):
///     <see cref="GetReadyAsync"/> returns null and the caller reports <c>BadResourceUnavailable</c>.
/// </summary>
/// <remarks>
///     A caller that got <c>Unavailable</c> waits on <see cref="WhenRecoveredAsync"/>; that both tells it when to
///     re-evaluate and makes the shard restart proactively (with backoff) instead of waiting for the next trigger, so
///     victims of a crash are re-evaluated on the shard's return (design §9).
/// </remarks>
internal sealed class WorkerShard : IAsyncDisposable
{
    internal static readonly TimeSpan InitialBackoff = TimeSpan.FromMilliseconds(250);
    internal static readonly TimeSpan MaxBackoff = TimeSpan.FromSeconds(30);
    internal static readonly TimeSpan HealthyResetAfter = TimeSpan.FromSeconds(60);

    /// <summary>
    ///     While a shard keeps failing to start, at most one Warning per this interval (lmxopcua#788); the other retries
    ///     log at Debug. The first failure of a down period is the one Error (with the exception).
    /// </summary>
    internal static readonly TimeSpan DefaultStartFailureWarningInterval = TimeSpan.FromSeconds(60);

    private static readonly Counter<long> StartsCounter = OtOpcUaTelemetry.Meter.CreateCounter<long>(
        "otopcua.script.worker.starts", unit: "{process}", description: "ScriptWorker processes started (tag: shard).");

    private static readonly Counter<long> ExitsCounter = OtOpcUaTelemetry.Meter.CreateCounter<long>(
        "otopcua.script.worker.exits", unit: "{process}",
        description: "ScriptWorker processes ended (tags: shard, reason = clean | killed_deadline | killed_heartbeat | killed_load | crash | protocol | idle | disposed).");

    private readonly Func<Action<WorkerExit>, CancellationToken, Task<WorkerConnection>> _launch;
    private readonly Action<WorkerShard, WorkerExit> _onExit;
    private readonly ILogger _log;
    private readonly object _gate = new();
    private readonly Random _jitter = new();
    private readonly CancellationTokenSource _disposing = new();
    private readonly TimeSpan _startFailureWarningInterval;
    private WorkerConnection? _current;
    private Task<WorkerConnection?>? _starting;
    private TaskCompletionSource? _recovered;
    private bool _recovering;
    private long _backoffUntil;
    private long _downSince;
    private int _consecutiveFailures;
    private bool _disposed;

    // The current run of failed starts (lmxopcua#788): how many, since when, and when the last Warning went out.
    private int _startFailures;
    private long _startFailingSince;
    private long _lastStartFailureWarning;

    public WorkerShard(string label, bool isolation, Func<Action<WorkerExit>, CancellationToken, Task<WorkerConnection>> launch,
        Action<WorkerShard, WorkerExit> onExit, ILogger log, TimeSpan? startFailureWarningInterval = null)
    {
        _startFailureWarningInterval = startFailureWarningInterval ?? DefaultStartFailureWarningInterval;
        Label = label;
        IsIsolation = isolation;
        _launch = launch;
        _onExit = onExit;
        _log = log;
        // The isolation worker runs one evaluation at a time, so a crash there names exactly one key (design §5.3.3).
        IsolationGate = isolation ? new SemaphoreSlim(1, 1) : null;
    }

    /// <summary>The shard label (<c>0</c>…<c>K-1</c>, or <c>iso</c>).</summary>
    public string Label { get; }

    /// <summary>Whether this is the isolation ("penalty box") worker.</summary>
    public bool IsIsolation { get; }

    /// <summary>Serialises evaluations on the isolation worker; null on a normal shard.</summary>
    public SemaphoreSlim? IsolationGate { get; }

    /// <summary>The live worker, if any (diagnostics / tests).</summary>
    public WorkerConnection? Current
    {
        get { lock (_gate) return _current is { IsAlive: true } c ? c : null; }
    }

    /// <summary>Consecutive start failures and crashes since the last healthy period (diagnostics / tests).</summary>
    public int ConsecutiveFailures
    {
        get { lock (_gate) return _consecutiveFailures; }
    }

    /// <summary>How long the shard has been down after a failure; null when it is up or was never needed.</summary>
    public TimeSpan? DownFor
    {
        get
        {
            lock (_gate)
            {
                if (_current is { IsAlive: true } || _downSince == 0) return null;
                return Stopwatch.GetElapsedTime(_downSince);
            }
        }
    }

    /// <summary>A snapshot for diagnostics and health.</summary>
    /// <returns>The status.</returns>
    public ScriptWorkerShardStatus Status()
    {
        WorkerConnection? live;
        int failures;
        TimeSpan? downFor;
        lock (_gate)
        {
            live = _current is { IsAlive: true } c ? c : null;
            failures = _consecutiveFailures;
            downFor = live is null && _downSince != 0 ? Stopwatch.GetElapsedTime(_downSince) : null;
        }
        var heartbeat = live?.LastHeartbeat;
        return new ScriptWorkerShardStatus(Label, IsIsolation, live is not null, downFor, live?.Pid, failures,
            live?.InFlightCount ?? 0, (int)(heartbeat?.Queued ?? 0), heartbeat?.HeapBytes ?? 0);
    }

    /// <summary>
    ///     The ready worker, starting one if needed. Waits at most <paramref name="grace"/> — through a start in progress
    ///     or a backoff — and returns null if no worker is ready by then. A zero grace still kicks a start that is due.
    /// </summary>
    /// <param name="grace">How long to wait.</param>
    /// <param name="ct">Cancellation.</param>
    /// <returns>The worker, or null.</returns>
    public async Task<WorkerConnection?> GetReadyAsync(TimeSpan grace, CancellationToken ct)
    {
        var deadline = Stopwatch.GetTimestamp() + (long)(grace.TotalSeconds * Stopwatch.Frequency);
        while (true)
        {
            Task<WorkerConnection?> starting;
            TimeSpan wait;
            lock (_gate)
            {
                if (_disposed) return null;
                if (_current is { IsAlive: true } live) return live;
                var now = Stopwatch.GetTimestamp();
                if (_starting is null && now >= _backoffUntil)
                {
                    _starting = StartCoreAsync();
                }
                if (_starting is not null)
                {
                    starting = _starting;
                    wait = TimeSpan.Zero;
                }
                else
                {
                    starting = null!;
                    wait = Stopwatch.GetElapsedTime(now, _backoffUntil);
                }
            }

            var remaining = Stopwatch.GetElapsedTime(Stopwatch.GetTimestamp(), deadline);
            if (remaining <= TimeSpan.Zero) return null;
            if (starting is not null)
            {
                try
                {
                    var started = await starting.WaitAsync(remaining, ct).ConfigureAwait(false);
                    if (started is { IsAlive: true }) return started;
                }
                catch (TimeoutException)
                {
                    return null;
                }
                continue;
            }
            if (wait > remaining) return null;
            await Task.Delay(wait, ct).ConfigureAwait(false);
        }
    }

    /// <summary>
    ///     Completes when this shard has a live worker other than <paramref name="failed"/>. Starts a background
    ///     recovery (start with backoff, retried at the cap) if none is live, so the caller does not depend on a later
    ///     trigger to bring the shard back.
    /// </summary>
    /// <param name="failed">The worker the caller's evaluation died on, if any.</param>
    /// <returns>A task that completes on recovery (or when the shard is disposed).</returns>
    public Task WhenRecoveredAsync(WorkerConnection? failed)
    {
        lock (_gate)
        {
            if (_disposed) return Task.CompletedTask;
            if (_current is { IsAlive: true } live)
            {
                // Still the same worker (it answered NOT_LOADED, or the caller lost a race with a restart that has
                // already happened): a short pause, so a seam cannot spin on a key the worker keeps refusing.
                return ReferenceEquals(live, failed) ? Task.Delay(InitialBackoff) : Task.CompletedTask;
            }
            _recovered ??= new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            var task = _recovered.Task;
            EnsureRecoveringLocked();
            return task;
        }
    }

    private void EnsureRecoveringLocked()
    {
        if (_recovering || _disposed) return;
        _recovering = true;
        _ = Task.Run(RecoverAsync);
    }

    /// <summary>Stops the live worker if it has had nothing to do for <paramref name="idle"/> (the isolation worker).</summary>
    /// <param name="idle">The idle threshold.</param>
    /// <returns>True when a worker was stopped.</returns>
    public async Task<bool> StopIfIdleAsync(TimeSpan idle)
    {
        WorkerConnection? victim;
        // Holding the isolation gate means no evaluation can pick this worker up while it is being stopped.
        var gate = IsolationGate;
        try
        {
            if (gate is not null && !gate.Wait(0)) return false;
        }
        catch (ObjectDisposedException)
        {
            return false;
        }
        try
        {
            lock (_gate)
            {
                if (_disposed || _current is not { IsAlive: true } live) return false;
                if (live.InFlightCount > 0 || live.LoadingCount > 0) return false;
                if (Stopwatch.GetElapsedTime(live.LastUsedTimestamp) < idle) return false;
                victim = live;
                _current = null;
            }
            _log.Information("ScriptWorker[{Shard}] pid {Pid} idle for {Idle}; stopping it (it restarts on demand)", Label, victim.Pid, idle);
            await victim.StopAsync(WorkerExitReason.Idle).ConfigureAwait(false);
            return true;
        }
        finally
        {
            gate?.Release();
        }
    }

    private async Task RecoverAsync()
    {
        try
        {
            while (true)
            {
                // A long grace: GetReadyAsync waits out the backoff itself and returns null only on a failed start.
                var live = await GetReadyAsync(MaxBackoff * 2, _disposing.Token).ConfigureAwait(false);
                lock (_gate)
                {
                    if (_disposed || live is not null)
                    {
                        _recovering = false;
                        TakeRecoveredLocked()?.TrySetResult();
                        return;
                    }
                }
            }
        }
        catch (OperationCanceledException)
        {
            lock (_gate)
            {
                _recovering = false;
                TakeRecoveredLocked()?.TrySetResult();
            }
        }
    }

    private TaskCompletionSource? TakeRecoveredLocked()
    {
        var tcs = _recovered;
        _recovered = null;
        return tcs;
    }

    private async Task<WorkerConnection?> StartCoreAsync()
    {
        await Task.Yield();
        WorkerConnection? connection = null;
        try
        {
            connection = await _launch(Attribute, CancellationToken.None).ConfigureAwait(false);
            StartsCounter.Add(1, new KeyValuePair<string, object?>("shard", Label));
            int failedBefore;
            long failingSince;
            lock (_gate)
            {
                failedBefore = _startFailures;
                failingSince = _startFailingSince;
                _startFailures = 0;
                _startFailingSince = 0;
                _lastStartFailureWarning = 0;
            }
            if (failedBefore == 0)
                _log.Information("ScriptWorker[{Shard}] started, pid {Pid}", Label, connection.Pid);
            else
                _log.Information("ScriptWorker[{Shard}] started, pid {Pid}, after {Attempts} failed start attempt(s) over {Outage}",
                    Label, connection.Pid, failedBefore, Stopwatch.GetElapsedTime(failingSince));
            var started = connection;
            _ = started.Exited.ContinueWith(t => OnExited(started, t.Result), CancellationToken.None,
                TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.OnlyOnRanToCompletion, TaskScheduler.Default);
            TaskCompletionSource? recovered;
            lock (_gate)
            {
                _starting = null;
                if (_disposed)
                {
                    _ = started.DisposeAsync().AsTask();
                    return null;
                }
                _current = started;
                _downSince = 0;
                recovered = TakeRecoveredLocked();
            }
            recovered?.TrySetResult();
            return started;
        }
        catch (Exception ex)
        {
            int attempt;
            TimeSpan failingFor;
            bool warn;
            lock (_gate)
            {
                _starting = null;
                MarkDownLocked();
                ScheduleBackoffLocked();
                attempt = ++_startFailures;
                var now = Stopwatch.GetTimestamp();
                if (attempt == 1) _startFailingSince = now;
                failingFor = Stopwatch.GetElapsedTime(_startFailingSince, now);
                warn = attempt > 1 && (_lastStartFailureWarning == 0
                                       || Stopwatch.GetElapsedTime(_lastStartFailureWarning, now) >= _startFailureWarningInterval);
                if (attempt == 1 || warn) _lastStartFailureWarning = now;
            }
            // lmxopcua#788: one Error with the stack per down period; the retries (every backoff, capped at 30 s, per
            // shard) are a rate-limited Warning without it, and the rest Debug. Recovery is one Information line above.
            if (attempt == 1)
                _log.Error(ex, "ScriptWorker[{Shard}] failed to start; scripts on it report runner unavailable until it does (retrying with backoff)", Label);
            else if (warn)
                _log.Warning("ScriptWorker[{Shard}] still failing to start: {Attempts} failed attempt(s) over {Outage}; last error: {Error}",
                    Label, attempt, failingFor, ex.Message);
            else
                _log.Debug("ScriptWorker[{Shard}] start attempt {Attempts} failed: {Error}", Label, attempt, ex.Message);
            if (connection is not null) await connection.DisposeAsync().ConfigureAwait(false);
            return null;
        }
    }

    /// <summary>Synchronous attribution, from <see cref="WorkerConnection"/>'s exit path, before any evaluation completes.</summary>
    private void Attribute(WorkerExit exit)
    {
        try { _onExit(this, exit); }
        catch (Exception ex) { _log.Error(ex, "ScriptWorker[{Shard}] exit handler threw", Label); }
    }

    private void OnExited(WorkerConnection connection, WorkerExit exit)
    {
        ExitsCounter.Add(1, new KeyValuePair<string, object?>("shard", Label),
            new KeyValuePair<string, object?>("reason", MetricReason(exit.Reason)));
        var expected = exit.Reason is WorkerExitReason.Disposed or WorkerExitReason.Idle;
        lock (_gate)
        {
            if (ReferenceEquals(_current, connection)) _current = null;
            if (!_disposed && !expected)
            {
                if (Stopwatch.GetElapsedTime(connection.ReadyTimestamp) >= HealthyResetAfter) _consecutiveFailures = 0;
                MarkDownLocked();
                ScheduleBackoffLocked();
                // A worker that was serving and died comes back by itself (after the backoff), so the scripts placed
                // on it find it warm and its victims are re-evaluated on its return (design §9).
                EnsureRecoveringLocked();
            }
        }
        if (!expected)
        {
            _log.Warning("ScriptWorker[{Shard}] pid {Pid} ended: {Reason} (exit code {ExitCode}, culprit {Culprit}, {InFlight} evaluation(s) and {Loading} compile(s) in flight)",
                Label, connection.Pid, exit.Reason, exit.ExitCode, exit.CulpritKey ?? "none", exit.InFlightKeys.Count, exit.LoadingKeys.Count);
        }
        _ = connection.DisposeAsync().AsTask();
    }

    private void MarkDownLocked()
    {
        if (_downSince == 0) _downSince = Stopwatch.GetTimestamp();
    }

    private void ScheduleBackoffLocked()
    {
        _consecutiveFailures++;
        var delay = NextBackoff(_consecutiveFailures, _jitter.NextDouble());
        _backoffUntil = Stopwatch.GetTimestamp() + (long)(delay.TotalSeconds * Stopwatch.Frequency);
    }

    /// <summary>The restart delay after <paramref name="failures"/> consecutive failures (design §5.3.5).</summary>
    /// <param name="failures">1 for the first failure.</param>
    /// <param name="unitJitter">A sample in [0, 1) mapped to ±20 %.</param>
    /// <returns>The delay.</returns>
    internal static TimeSpan NextBackoff(int failures, double unitJitter)
    {
        var exponent = Math.Clamp(failures - 1, 0, 30);
        var baseMs = Math.Min(MaxBackoff.TotalMilliseconds, InitialBackoff.TotalMilliseconds * Math.Pow(2, exponent));
        var factor = 0.8 + (0.4 * Math.Clamp(unitJitter, 0, 1));
        return TimeSpan.FromMilliseconds(Math.Min(MaxBackoff.TotalMilliseconds * 1.2, baseMs * factor));
    }

    private static string MetricReason(WorkerExitReason reason) => reason switch
    {
        WorkerExitReason.Clean => "clean",
        WorkerExitReason.KilledDeadline => "killed_deadline",
        WorkerExitReason.KilledHeartbeat => "killed_heartbeat",
        WorkerExitReason.KilledLoadTimeout => "killed_load",
        WorkerExitReason.ProtocolFault => "protocol",
        WorkerExitReason.Disposed => "disposed",
        WorkerExitReason.Idle => "idle",
        _ => "crash",
    };

    /// <summary>Stops the shard's worker.</summary>
    /// <returns>A task that completes when the worker has ended.</returns>
    public async ValueTask DisposeAsync()
    {
        WorkerConnection? current;
        Task<WorkerConnection?>? starting;
        TaskCompletionSource? recovered;
        lock (_gate)
        {
            if (_disposed) return;
            _disposed = true;
            current = _current;
            starting = _starting;
            _current = null;
            recovered = TakeRecoveredLocked();
        }
        recovered?.TrySetResult();
        _disposing.Cancel();
        if (starting is not null)
        {
            try { current ??= await starting.WaitAsync(TimeSpan.FromSeconds(15)).ConfigureAwait(false); } catch { /* shutting down */ }
        }
        if (current is not null) await current.DisposeAsync().ConfigureAwait(false);
        IsolationGate?.Dispose();
    }

    /// <summary>Maps a context kind to the wire enum.</summary>
    /// <param name="kind">The kind.</param>
    /// <returns>The wire kind.</returns>
    internal static ContextKind ToWire(ScriptContextKind kind) => kind switch
    {
        ScriptContextKind.VirtualTag => ContextKind.VirtualTag,
        ScriptContextKind.Calculation => ContextKind.Calculation,
        ScriptContextKind.AlarmPredicate => ContextKind.AlarmPredicate,
        _ => ContextKind.Unspecified,
    };
}