[BED-5972] Log timeout-wrapped execution time stats, add adaptive timeout logic

src/CommonLib/AdaptiveTimeout.cs

@@ -0,0 +1,226 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using SharpHoundRPC.NetAPINative;
+
+namespace SharpHoundCommonLib;
+
+public sealed class AdaptiveTimeout : IDisposable {
+    private readonly ExecutionTimeSampler _sampler;
+    private readonly ILogger _log;
+    private readonly TimeSpan _maxTimeout;
+    private readonly bool _useAdaptiveTimeout;
+    private readonly int _minSamplesForAdaptiveTimeout;
+    private int _clearSamplesDecay;
+    private const int TimeSpikePenalty = 2;
+    private const int TimeSpikeForgiveness = 1;
+    private const int ClearSamplesThreshold = 5;
+    private const int StdDevMultiplier = 5;
+
+    public AdaptiveTimeout(TimeSpan maxTimeout, ILogger log, int sampleCount = 100, int logFrequency = 1000, int minSamplesForAdaptiveTimeout = 30, bool useAdaptiveTimeout = true) {
+        if (maxTimeout <= TimeSpan.Zero)
+            throw new ArgumentException("maxTimeout must be positive", nameof(maxTimeout));
+        if (sampleCount <= 0)
+            throw new ArgumentException("sampleCount must be positive", nameof(sampleCount));
+        if (logFrequency <= 0)
+            throw new ArgumentException("logFrequency must be positive", nameof(logFrequency));
+        if (minSamplesForAdaptiveTimeout <= 0)
+            throw new ArgumentException("minSamplesForAdaptiveTimeout must be positive", nameof(minSamplesForAdaptiveTimeout));
+        if (log == null)
+            throw new ArgumentNullException(nameof(log));
+
+        _sampler = new ExecutionTimeSampler(log, sampleCount, logFrequency);
+        _log = log;
+        _maxTimeout = maxTimeout;
+        _useAdaptiveTimeout = useAdaptiveTimeout;
+        _minSamplesForAdaptiveTimeout = minSamplesForAdaptiveTimeout;
+    }
+
+    public void ClearSamples() {
+        _clearSamplesDecay = 0;
+        _sampler.ClearSamples();
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<Result<T>> ExecuteWithTimeout<T>(Func<CancellationToken, T> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<Result> ExecuteWithTimeout(Action<CancellationToken> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<Result<T>> ExecuteWithTimeout<T>(Func<CancellationToken, Task<T>> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<Result> ExecuteWithTimeout(Func<CancellationToken, Task> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<NetAPIResult<T>> ExecuteNetAPIWithTimeout<T>(Func<CancellationToken, NetAPIResult<T>> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteNetAPIWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<SharpHoundRPC.Result<T>> ExecuteRPCWithTimeout<T>(Func<CancellationToken, SharpHoundRPC.Result<T>> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteRPCWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    /// <summary>
+    /// Ignores the result of a function if it runs longer than a budgeted time, unblocking the caller.
+    /// Logs aggregate execution time data.
+    /// Manages its own timeout.
+    /// A cancellation token is passed to the executing function so it may exit cleanly if timeout is reached.
+    /// Please don't wrap a cached function in this timeout if adaptive timeouts enabled, normal distributions are better.
+    /// DO NOT use a single AdaptiveTimeout for multiple functions.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    /// <param name="func"></param>
+    /// <param name="parentToken"></param>
+    /// <returns>Returns a Fail result if a task runs longer than its budgeted time.</returns>
+    public async Task<SharpHoundRPC.Result<T>> ExecuteRPCWithTimeout<T>(Func<CancellationToken, Task<SharpHoundRPC.Result<T>>> func, CancellationToken parentToken = default) {
+        var result = await Timeout.ExecuteRPCWithTimeout(GetAdaptiveTimeout(), (timeoutToken) => _sampler.SampleExecutionTime(() => func(timeoutToken)), parentToken);
+        TimeSpikeSafetyValve(result.IsSuccess);
+        return result;
+    }
+
+    public void Dispose() {
+        _sampler.Dispose();
+    }
+
+    // Within 5 standard deviations will have a conservative lower bound of catching 96% of executions (1 - 1/5^2),
+    // regardless of sample shape
+    // so long as those samples are independent and identically distributed
+    // (and if they're not, our TimeSpikeSafetyValve should provide us with some adaptability)
+    // But the effective collection rate is probably closer to 98+%
+    // (in part because we don't need to filter out "too fast" outliers)
+    // But we'll cap at configured maximum timeout
+    // https://modelassist.epixanalytics.com/space/EA/26574957/Tchebysheffs+Rule
+    // https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables
+    public TimeSpan GetAdaptiveTimeout() {
+        if (!UseAdaptiveTimeout())
+            return _maxTimeout;
+
+        try {
+            var stdDev = _sampler.StandardDeviation();
+            var adaptiveTimeoutMs = _sampler.Average() + (stdDev * StdDevMultiplier);
+            var cappedTimeoutMS = Math.Min(adaptiveTimeoutMs, _maxTimeout.TotalMilliseconds);
+            return TimeSpan.FromMilliseconds(cappedTimeoutMS);
+        }
+        catch (Exception ex) {
+            _log.LogError(ex, "Error calculating adaptive timeout, defaulting to max timeout.");
+            return _maxTimeout;
+        }
+    }
+
+    // AdaptiveTimeout will not respond well to rapid spikes in execution time
+    // imagine the wrapped function very regularly executes in 10ms
+    // then suddenly starts taking a regular 100ms
+    // this is fine (if it fits in our max timeout budget), and we shouldn't block
+    // so we should create a safety valve in case this happens to reset our data samples
+    private void TimeSpikeSafetyValve(bool isSuccess) {
+        if (isSuccess) {
+            _clearSamplesDecay -= TimeSpikeForgiveness;
+            _clearSamplesDecay = Math.Max(0, _clearSamplesDecay);
+        }
+        else
+            _clearSamplesDecay += TimeSpikePenalty;
+
+
+        if (_clearSamplesDecay >= ClearSamplesThreshold) {
+            if (UseAdaptiveTimeout()) {
+                ClearSamples();
+                _log.LogTrace("Time spike safety valve event at timeout {CurrentTimeout}.", GetAdaptiveTimeout());
+            }
+            else {
+                _log.LogWarning("This call is frequently running over the maximum allowed timeout of {MaxTimeout}.", _maxTimeout);
+            }
+        }
+    }
+
+    private bool UseAdaptiveTimeout() {
+        return _useAdaptiveTimeout && _sampler.Count >= _minSamplesForAdaptiveTimeout;
+    }
+}

src/CommonLib/ExecutionTimeSampler.cs

@@ -0,0 +1,104 @@
+using System;
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+
+namespace SharpHoundCommonLib;
+
+/// <summary>
+/// Holds a rolling sample of execution times on a function, providing and logging data aggregates.
+/// </summary>
+public class ExecutionTimeSampler : IDisposable {
+    private readonly ILogger _log;
+    private readonly int _sampleCount;
+    private readonly int _logFrequency;
+    private int _samplesSinceLastLog;
+    private ConcurrentQueue<double> _samples;
+
+    public int Count => _samples.Count;
+
+    public ExecutionTimeSampler(ILogger log, int sampleCount, int logFrequency) {
+        _log = log;
+        _sampleCount = sampleCount;
+        _logFrequency = logFrequency;
+        _samplesSinceLastLog = 0;
+        _samples = new ConcurrentQueue<double>();
+    }
+
+    public void ClearSamples() {
+        Log(flush: true);
+        _samples = new ConcurrentQueue<double>();
+    }
+
+    public double StandardDeviation() {
+        double average = _samples.Average();
+        double sumOfSquaresOfDifferences = _samples.Select(val => (val - average) * (val - average)).Sum();
+        double stddiv = Math.Sqrt(sumOfSquaresOfDifferences / _samples.Count);
+
+        return stddiv;
+    }
+
+    public double Average() => _samples.Average();
+
+    public async Task<T> SampleExecutionTime<T>(Func<Task<T>> func) {
+        var stopwatch = Stopwatch.StartNew();
+        var result = await func.Invoke();
+        stopwatch.Stop();
+        AddTimeSample(stopwatch.Elapsed);
+
+        return result;
+    }
+
+    public async Task SampleExecutionTime(Func<Task> func) {
+        var stopwatch = Stopwatch.StartNew();
+        await func.Invoke();
+        stopwatch.Stop();
+        AddTimeSample(stopwatch.Elapsed);
+    }
+
+    public T SampleExecutionTime<T>(Func<T> func) {
+        var stopwatch = Stopwatch.StartNew();
+        var result = func.Invoke();
+        stopwatch.Stop();
+        AddTimeSample(stopwatch.Elapsed);
+
+        return result;
+    }
+
+    public void SampleExecutionTime(Action func) {
+        var stopwatch = Stopwatch.StartNew();
+        func.Invoke();
+        stopwatch.Stop();
+        AddTimeSample(stopwatch.Elapsed);
+    }
+
+    public void Dispose() {
+        Log(flush: true);
+    }
+
+    private void AddTimeSample(TimeSpan timeSpan) {
+        while (_samples.Count >= _sampleCount) {
+            _samples.TryDequeue(out _);
+        }
+
+        _samples.Enqueue(timeSpan.TotalMilliseconds);
+        _samplesSinceLastLog++;
+
+        Log();
+    }
+
+    private void Log(bool flush = false) {
+        if ((flush || _samplesSinceLastLog >= _logFrequency) && _samples.Count > 0) {
+            try {
+                _log.LogInformation("Execution time Average: {Average}ms, StdDiv: {StandardDeviation}ms", _samples.Average(), StandardDeviation());
+            }
+            catch (Exception ex) {
+                _log.LogWarning("Failed to calculate execution time statistics: {Error}", ex.Message);
+            }
+
+            _samplesSinceLastLog = 0;
+        }
+    }
+}

src/CommonLib/IRegistryKey.cs

@@ -10,6 +10,7 @@ public interface IRegistryKey {
 
     public class SHRegistryKey : IRegistryKey, IDisposable {
         private readonly RegistryKey _currentKey;
+        private static readonly AdaptiveTimeout _adaptiveTimeout = new AdaptiveTimeout(maxTimeout: TimeSpan.FromSeconds(10), Logging.LogProvider.CreateLogger(nameof(SHRegistryKey)));
 
         private SHRegistryKey(RegistryKey registryKey) {
             _currentKey = registryKey;
@@ -35,7 +36,7 @@ public object GetValue(string subkey, string name) {
         /// <exception cref="System.Security.SecurityException"></exception>
         /// <exception cref="UnauthorizedAccessException"></exception>
         public static async Task<SHRegistryKey> Connect(RegistryHive hive, string machineName) {
-            var remoteKey = await Timeout.ExecuteWithTimeout(TimeSpan.FromSeconds(10), (_) => RegistryKey.OpenRemoteBaseKey(hive, machineName));
+            var remoteKey = await _adaptiveTimeout.ExecuteWithTimeout((_) => RegistryKey.OpenRemoteBaseKey(hive, machineName));
             if (remoteKey.IsSuccess)
                 return new SHRegistryKey(remoteKey.Value);
             else