Re-worked RegistryWatcher more stable now

src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.ProvisioningStatusTracking.cs

@@ -2,7 +2,7 @@ using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text.Json;
-using System.Threading;
+using AutopilotMonitor.Agent.Core.Monitoring.Core;
 using AutopilotMonitor.Agent.Core.Monitoring.Interop;
 using AutopilotMonitor.Shared;
 using AutopilotMonitor.Shared.Models;
@@ -30,7 +30,7 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors
     ///   - Emit once when categorySucceeded resolves to true or false (final outcome)
     ///   - Fire EspFailureDetected on subcategory failure even if categorySucceeded is still null
     ///
-    /// Uses RegNotifyChangeKeyValue for instant registry change detection (no polling timer).
+    /// Uses RegistryWatcher (RegNotifyChangeKeyValue) for instant registry change detection.
     /// </summary>
     public partial class EspAndHelloTracker
     {
@@ -65,159 +65,80 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors
             _provisioningCategoriesResolved = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
             _lastSubcategoryStates = new Dictionary<string, Dictionary<string, string>>(StringComparer.OrdinalIgnoreCase);
 
-            _registryWatcherStopEvent = RegistryWatcherNativeMethods.CreateEvent(IntPtr.Zero, true, false, null);
-            if (_registryWatcherStopEvent == IntPtr.Zero)
+            // Try to start immediately; if key doesn't exist yet, retry every 2s
+            if (!TryStartProvisioningWatcher())
             {
-                _logger.Warning("Failed to create registry watcher stop event — provisioning status tracking disabled");
-                return;
-            }
-
-            _registryWatcherThread = new Thread(RegistryWatcherLoop)
-            {
-                IsBackground = true,
-                Name = "ProvisioningStatusWatcher"
-            };
-            _registryWatcherThread.Start();
-
-            _logger.Info("Provisioning status registry watcher started");
-        }
-
-        private void StopProvisioningStatusWatcher(string reason)
-        {
-            if (_registryWatcherStopEvent == IntPtr.Zero)
-                return;
-
-            try
-            {
-                // Signal stop event to wake up the watcher thread
-                RegistryWatcherNativeMethods.SetEvent(_registryWatcherStopEvent);
-
-                // Wait for thread to exit
-                if (_registryWatcherThread != null && _registryWatcherThread.IsAlive)
-                {
-                    _registryWatcherThread.Join(TimeSpan.FromSeconds(3));
-                }
-
-                RegistryWatcherNativeMethods.CloseHandle(_registryWatcherStopEvent);
-                _registryWatcherStopEvent = IntPtr.Zero;
-                _registryWatcherThread = null;
-
-                _logger.Info($"Provisioning status watcher stopped: {reason}");
-            }
-            catch (Exception ex)
-            {
-                _logger.Error("Error stopping provisioning status watcher", ex);
+                _logger.Info("Provisioning status registry key not yet present — retrying every 2s");
+                _provisioningWatcherRetryTimer = new System.Threading.Timer(
+                    _ =>
+                    {
+                        if (TryStartProvisioningWatcher())
+                        {
+                            _provisioningWatcherRetryTimer?.Dispose();
+                            _provisioningWatcherRetryTimer = null;
+                        }
+                    },
+                    null,
+                    TimeSpan.FromSeconds(2),
+                    TimeSpan.FromSeconds(2));
             }
         }
 
         /// <summary>
-        /// Background thread: opens the provisioning registry key and waits for value changes
-        /// via RegNotifyChangeKeyValue. Falls back to retry loop if key doesn't exist yet.
+        /// Attempts to create and start the RegistryWatcher for the provisioning key.
+        /// Returns false if the key doesn't exist yet.
         /// </summary>
-        private void RegistryWatcherLoop()
+        private bool TryStartProvisioningWatcher()
         {
-            IntPtr hKey = IntPtr.Zero;
-            IntPtr hNotifyEvent = IntPtr.Zero;
-
             try
             {
-                // Wait for registry key to appear (enrollment may not have started yet)
-                while (true)
+                using (var key = Registry.LocalMachine.OpenSubKey(ProvisioningStatusRegistryPath))
                 {
-                    int result = RegistryWatcherNativeMethods.RegOpenKeyEx(
-                        RegistryWatcherNativeMethods.HKEY_LOCAL_MACHINE,
-                        ProvisioningStatusRegistryPath,
-                        0,
-                        RegistryWatcherNativeMethods.KEY_NOTIFY | RegistryWatcherNativeMethods.KEY_READ,
-                        out hKey);
-
-                    if (result == 0)
-                        break;
-
-                    // Key doesn't exist yet — wait 2 seconds or until stop signaled
-                    uint waitResult = RegistryWatcherNativeMethods.WaitForSingleObject(_registryWatcherStopEvent, 2000);
-                    if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0)
-                        return; // Stop requested
+                    if (key == null)
+                        return false;
                 }
 
-                // Capture initial state before notification loop.
-                // Any changes between this read and the first RegNotifyChangeKeyValue registration
-                // are covered by the 15s heartbeat fallback.
+                // Capture initial state before watcher starts
                 CheckProvisioningStatus();
 
-                // Create notification event
-                hNotifyEvent = RegistryWatcherNativeMethods.CreateEvent(IntPtr.Zero, false, false, null);
-                if (hNotifyEvent == IntPtr.Zero)
+                // Create and start watcher
+                _provisioningWatcher = new RegistryWatcher(
+                    RegistryHive.LocalMachine,
+                    ProvisioningStatusRegistryPath,
+                    watchSubtree: true,
+                    filter: RegistryNativeMethods.RegChangeNotifyFilter.LastSet);
+
+                _provisioningWatcher.Changed += (s, e) => CheckProvisioningStatus();
+                _provisioningWatcher.Error += (s, ex) => _logger.Warning($"Provisioning watcher handler error: {ex.Message}");
+
+                _provisioningWatcher.Start();
+                _logger.Info("Provisioning status registry watcher started");
+                return true;
+            }
+            catch (Exception ex)
+            {
+                _logger.Warning($"Failed to start provisioning watcher: {ex.Message}");
+                return false;
+            }
+        }
+
+        private void StopProvisioningStatusWatcher(string reason)
+        {
+            try
+            {
+                _provisioningWatcherRetryTimer?.Dispose();
+                _provisioningWatcherRetryTimer = null;
+
+                if (_provisioningWatcher != null)
                 {
-                    _logger.Warning("Failed to create registry notify event — watcher thread exiting");
-                    return;
-                }
-
-                var waitHandles = new[] { hNotifyEvent, _registryWatcherStopEvent };
-
-                // Main watch loop: register → wait → read
-                // IMPORTANT: CheckProvisioningStatus() must run AFTER WaitForMultipleObjects, not between
-                // RegNotifyChangeKeyValue and Wait. The managed Registry.OpenSubKey() call interferes with
-                // pending RegNotifyChangeKeyValue notifications when called before Wait (proven in Build 416).
-                // The 15s heartbeat ensures changes are detected even if a notification is missed.
-                const uint heartbeatMs = 15_000;
-
-                while (true)
-                {
-                    // 1. Register for value change notifications
-                    int regResult = RegistryWatcherNativeMethods.RegNotifyChangeKeyValue(
-                        hKey,
-                        true,  // Watch subtree
-                        RegistryWatcherNativeMethods.REG_NOTIFY_CHANGE_LAST_SET,
-                        hNotifyEvent,
-                        true); // Asynchronous
-
-                    if (regResult != 0)
-                    {
-                        _logger.Warning($"RegNotifyChangeKeyValue failed (error {regResult}) — watcher thread exiting");
-                        return;
-                    }
-
-                    // 2. Wait for registry change notification or heartbeat timeout
-                    uint waitResult = RegistryWatcherNativeMethods.WaitForMultipleObjects(
-                        (uint)waitHandles.Length,
-                        waitHandles,
-                        false,
-                        heartbeatMs);
-
-                    if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0)
-                    {
-                        // Registry notification fired — instant detection
-                        _logger.Debug("Provisioning status watcher: registry change notification received");
-                    }
-                    else if (waitResult == RegistryWatcherNativeMethods.WAIT_OBJECT_0 + 1)
-                    {
-                        return; // Stop requested
-                    }
-                    else if (waitResult == RegistryWatcherNativeMethods.WAIT_TIMEOUT)
-                    {
-                        // Heartbeat — re-read to catch any missed notifications
-                    }
-                    else
-                    {
-                        _logger.Warning($"Registry watcher unexpected wait result: {waitResult} — exiting");
-                        return;
-                    }
-
-                    // 3. Read current state AFTER wait (not between register and wait)
-                    CheckProvisioningStatus();
+                    _provisioningWatcher.Dispose();
+                    _provisioningWatcher = null;
+                    _logger.Info($"Provisioning status watcher stopped: {reason}");
                 }
             }
             catch (Exception ex)
             {
-                _logger.Warning($"Registry watcher thread failed: {ex.Message}");
-            }
-            finally
-            {
-                if (hNotifyEvent != IntPtr.Zero)
-                    RegistryWatcherNativeMethods.CloseHandle(hNotifyEvent);
-                if (hKey != IntPtr.Zero)
-                    RegistryWatcherNativeMethods.RegCloseKey(hKey);
+                _logger.Error("Error stopping provisioning status watcher", ex);
             }
         }
 
@@ -255,16 +176,15 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors
                             }
                         }
 
-                        // Self-termination: only auto-stop when all categories resolved with success (no failures)
-                        // Signal the stop event directly (we're on the watcher thread — can't call StopProvisioningStatusWatcher
-                        // which would Join() on ourselves)
+                        // Self-termination: only auto-stop when all categories resolved with success (no failures).
+                        // Uses RequestStop() which is safe to call from within the Changed handler (non-blocking).
                         if (_provisioningCategoriesResolved.Count > 0
                             && _lastProvisioningJson.Count > 0
                             && _provisioningCategoriesResolved.Count >= _lastProvisioningJson.Count
                             && !_provisioningFailureFired.Any())
                         {
-                            _logger.Info("Provisioning status watcher: all categories resolved with success — signaling stop");
-                            RegistryWatcherNativeMethods.SetEvent(_registryWatcherStopEvent);
+                            _logger.Info("Provisioning status watcher: all categories resolved with success — requesting stop");
+                            _provisioningWatcher?.RequestStop();
                         }
 
                         // Fire EspFailureDetected AFTER event emission

src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Collectors/EspAndHelloTracker.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Diagnostics.Eventing.Reader;
 using System.Threading;
 using AutopilotMonitor.Agent.Core.Logging;
+using AutopilotMonitor.Agent.Core.Monitoring.Core;
 using AutopilotMonitor.Shared.Models;
 using Microsoft.Win32;
 
@@ -36,8 +37,8 @@ namespace AutopilotMonitor.Agent.Core.Monitoring.Collectors
         private System.Threading.Timer _policyCheckTimer;
         private System.Threading.Timer _helloWaitTimer;
         private System.Threading.Timer _helloCompletionTimer;
-        private Thread _registryWatcherThread;
-        private IntPtr _registryWatcherStopEvent = IntPtr.Zero;
+        private RegistryWatcher _provisioningWatcher;
+        private System.Threading.Timer _provisioningWatcherRetryTimer;
 
         private bool _isPolicyConfigured = false;
         private bool _isHelloPolicyEnabled = false; // true when Hello policy is explicitly detected as enabled

src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Core/RegistryWatcher.cs

@@ -0,0 +1,202 @@
+using System;
+using System.ComponentModel;
+using System.Threading;
+using System.Threading.Tasks;
+using AutopilotMonitor.Agent.Core.Monitoring.Interop;
+using Microsoft.Win32;
+using Microsoft.Win32.SafeHandles;
+
+namespace AutopilotMonitor.Agent.Core.Monitoring.Core
+{
+    /// <summary>
+    /// Watches a registry key for changes using RegNotifyChangeKeyValue.
+    /// Fires the Changed event on each detected change.
+    ///
+    /// Usage:
+    ///   var watcher = new RegistryWatcher(RegistryHive.LocalMachine, @"SOFTWARE\...", watchSubtree: true);
+    ///   watcher.Changed += (s, e) => { /* read values */ };
+    ///   watcher.Start();
+    ///   // later:
+    ///   await watcher.StopAsync();
+    ///   watcher.Dispose();
+    ///
+    /// Important: one change per registration — the watcher re-registers automatically after each notification.
+    /// The watched scope is a key (including values); to track a specific value, compare in the Changed handler.
+    /// </summary>
+    internal sealed class RegistryWatcher : IDisposable
+    {
+        private readonly UIntPtr _hive;
+        private readonly string _subKey;
+        private readonly bool _watchSubtree;
+        private readonly RegistryNativeMethods.RegChangeNotifyFilter _filter;
+
+        private CancellationTokenSource _cts;
+        private Task _watchTask;
+        private readonly object _sync = new object();
+
+        /// <summary>
+        /// Fired when the watched key (or subtree) changes. Handler runs on a thread pool thread.
+        /// </summary>
+        public event EventHandler Changed;
+
+        /// <summary>
+        /// Fired when the Changed handler throws. Allows the consumer to log without crashing the watcher.
+        /// </summary>
+        public event EventHandler<Exception> Error;
+
+        public RegistryWatcher(
+            RegistryHive hive,
+            string subKey,
+            bool watchSubtree = false,
+            RegistryNativeMethods.RegChangeNotifyFilter filter =
+                RegistryNativeMethods.RegChangeNotifyFilter.LastSet)
+        {
+            _hive = HiveToPointer(hive);
+            _subKey = subKey ?? throw new ArgumentNullException(nameof(subKey));
+            _watchSubtree = watchSubtree;
+            _filter = filter;
+        }
+
+        /// <summary>
+        /// Starts watching on a background thread pool thread.
+        /// Throws Win32Exception if the registry key cannot be opened.
+        /// </summary>
+        public void Start()
+        {
+            lock (_sync)
+            {
+                if (_watchTask != null)
+                    throw new InvalidOperationException("Watcher is already running.");
+
+                _cts = new CancellationTokenSource();
+                _watchTask = Task.Run(() => WatchLoopAsync(_cts.Token));
+            }
+        }
+
+        /// <summary>
+        /// Non-blocking stop request. Safe to call from within the Changed handler
+        /// (does not wait for the loop to exit, so no deadlock).
+        /// </summary>
+        public void RequestStop()
+        {
+            lock (_sync)
+            {
+                _cts?.Cancel();
+            }
+        }
+
+        /// <summary>
+        /// Cancels the watcher and waits for the background loop to exit.
+        /// Do NOT call from within the Changed handler — use RequestStop() instead.
+        /// </summary>
+        public async Task StopAsync()
+        {
+            Task taskToWait;
+
+            lock (_sync)
+            {
+                if (_cts == null)
+                    return;
+
+                _cts.Cancel();
+                taskToWait = _watchTask;
+            }
+
+            if (taskToWait != null)
+            {
+                try
+                {
+                    await taskToWait.ConfigureAwait(false);
+                }
+                catch (OperationCanceledException) { }
+                catch (Win32Exception) { }
+            }
+
+            lock (_sync)
+            {
+                _cts?.Dispose();
+                _cts = null;
+                _watchTask = null;
+            }
+        }
+
+        private async Task WatchLoopAsync(CancellationToken cancellationToken)
+        {
+            int openResult = RegistryNativeMethods.RegOpenKeyEx(
+                _hive, _subKey, 0, RegistryNativeMethods.KEY_NOTIFY,
+                out SafeRegistryHandle keyHandle);
+
+            if (openResult != 0)
+                throw new Win32Exception(openResult, $"RegOpenKeyEx failed for '{_subKey}'");
+
+            using (keyHandle)
+            using (var changedEvent = new AutoResetEvent(false))
+            using (var cancelEvent = new ManualResetEvent(false))
+            using (cancellationToken.Register(() => cancelEvent.Set()))
+            {
+                var waitHandles = new WaitHandle[] { changedEvent, cancelEvent };
+
+                while (!cancellationToken.IsCancellationRequested)
+                {
+                    // Register for exactly ONE change notification
+                    int notifyResult = RegistryNativeMethods.RegNotifyChangeKeyValue(
+                        keyHandle,
+                        _watchSubtree,
+                        _filter,
+                        changedEvent.SafeWaitHandle,
+                        true);
+
+                    if (notifyResult != 0)
+                        throw new Win32Exception(notifyResult, "RegNotifyChangeKeyValue failed");
+
+                    // Wait for change or cancellation
+                    int signaled = WaitHandle.WaitAny(waitHandles);
+
+                    if (signaled == 1)
+                        break; // cancel requested
+
+                    // Notify consumer
+                    try
+                    {
+                        Changed?.Invoke(this, EventArgs.Empty);
+                    }
+                    catch (Exception ex)
+                    {
+                        Error?.Invoke(this, ex);
+                    }
+
+                    await Task.Yield();
+                }
+
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+        }
+
+        public void Dispose()
+        {
+            try
+            {
+                StopAsync().GetAwaiter().GetResult();
+            }
+            catch
+            {
+                // Dispose must not throw
+            }
+        }
+
+        // ===== Helpers =====
+
+        private static UIntPtr HiveToPointer(RegistryHive hive)
+        {
+            switch (hive)
+            {
+                case RegistryHive.ClassesRoot: return (UIntPtr)0x80000000u;
+                case RegistryHive.CurrentUser: return (UIntPtr)0x80000001u;
+                case RegistryHive.LocalMachine: return (UIntPtr)0x80000002u;
+                case RegistryHive.Users: return (UIntPtr)0x80000003u;
+                case RegistryHive.CurrentConfig: return (UIntPtr)0x80000005u;
+                default: throw new ArgumentOutOfRangeException(nameof(hive));
+            }
+        }
+    }
+}

src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryNativeMethods.cs

@@ -0,0 +1,36 @@
+using System;
+using System.Runtime.InteropServices;
+using Microsoft.Win32.SafeHandles;
+
+namespace AutopilotMonitor.Agent.Core.Monitoring.Interop
+{
+    /// <summary>
+    /// P/Invoke declarations for registry change notification via RegNotifyChangeKeyValue.
+    /// Uses SafeRegistryHandle and SafeWaitHandle for proper handle lifecycle management.
+    /// Used by RegistryWatcher to detect registry changes instantly.
+    /// </summary>
+    internal static class RegistryNativeMethods
+    {
+        public const int KEY_NOTIFY = 0x0010;
+
+        [Flags]
+        public enum RegChangeNotifyFilter : uint
+        {
+            Name = 1,
+            Attributes = 2,
+            LastSet = 4,
+            Security = 8
+        }
+
+        [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+        public static extern int RegOpenKeyEx(
+            UIntPtr hKey, string lpSubKey, uint ulOptions,
+            int samDesired, out SafeRegistryHandle phkResult);
+
+        [DllImport("advapi32.dll", SetLastError = true)]
+        public static extern int RegNotifyChangeKeyValue(
+            SafeRegistryHandle hKey, bool bWatchSubtree,
+            RegChangeNotifyFilter dwNotifyFilter,
+            SafeWaitHandle hEvent, bool fAsynchronous);
+    }
+}

src/Agent/AutopilotMonitor.Agent.Core/Monitoring/Interop/RegistryWatcherNativeMethods.cs

@@ -1,51 +0,0 @@
-using System;
-using System.Runtime.InteropServices;
-
-namespace AutopilotMonitor.Agent.Core.Monitoring.Interop
-{
-    /// <summary>
-    /// P/Invoke declarations for registry change notification via RegNotifyChangeKeyValue.
-    /// Used by EspAndHelloTracker.ProvisioningStatusTracking to detect provisioning status changes instantly.
-    /// </summary>
-    internal static class RegistryWatcherNativeMethods
-    {
-        public static readonly IntPtr HKEY_LOCAL_MACHINE = new IntPtr(unchecked((int)0x80000002));
-
-        public const uint KEY_NOTIFY = 0x0010;
-        public const uint KEY_READ = 0x20019;
-        public const uint REG_NOTIFY_CHANGE_LAST_SET = 0x00000004;
-        public const uint WAIT_OBJECT_0 = 0;
-        public const uint WAIT_TIMEOUT = 258;
-        public const uint INFINITE = 0xFFFFFFFF;
-
-        [DllImport("advapi32.dll", SetLastError = true)]
-        public static extern int RegOpenKeyEx(
-            IntPtr hKey, string subKey, uint options, uint samDesired, out IntPtr result);
-
-        [DllImport("advapi32.dll", SetLastError = true)]
-        public static extern int RegNotifyChangeKeyValue(
-            IntPtr hKey, bool watchSubtree, uint notifyFilter, IntPtr hEvent, bool asynchronous);
-
-        [DllImport("advapi32.dll", SetLastError = true)]
-        public static extern int RegCloseKey(IntPtr hKey);
-
-        [DllImport("kernel32.dll", SetLastError = true)]
-        public static extern IntPtr CreateEvent(
-            IntPtr lpEventAttributes, bool bManualReset, bool bInitialState, string lpName);
-
-        [DllImport("kernel32.dll", SetLastError = true)]
-        [return: MarshalAs(UnmanagedType.Bool)]
-        public static extern bool SetEvent(IntPtr hEvent);
-
-        [DllImport("kernel32.dll", SetLastError = true)]
-        public static extern uint WaitForSingleObject(IntPtr hHandle, uint dwMilliseconds);
-
-        [DllImport("kernel32.dll", SetLastError = true)]
-        public static extern uint WaitForMultipleObjects(
-            uint nCount, IntPtr[] lpHandles, bool bWaitAll, uint dwMilliseconds);
-
-        [DllImport("kernel32.dll", SetLastError = true)]
-        [return: MarshalAs(UnmanagedType.Bool)]
-        public static extern bool CloseHandle(IntPtr hObject);
-    }
-}

src/Agent/AutopilotMonitor.Agent/buildcounter.txt

@@ -1 +1 @@
-421
+424