From 98d8df4adf2cecd40256374dede057befc69e1ec Mon Sep 17 00:00:00 2001
From: Joseph Doherty <dohejw01@gmail.com>
Date: Fri, 22 May 2026 10:43:35 -0400
Subject: [PATCH] fix(driver-twincat): resolve Medium code-review finding
 (Driver.TwinCAT-009)

Swap _devices and _tagsByName to ConcurrentDictionary so ShutdownAsync
Clear() no longer races concurrent TryGetValue calls; store ProbeTask
on DeviceState and await it in ShutdownAsync before disposing the client
and gate, eliminating the probe-disposal race.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 code-reviews/Driver.TwinCAT/findings.md       |  2 +-
 .../TwinCATDriver.cs                          | 25 ++++++++++++++++---
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/code-reviews/Driver.TwinCAT/findings.md b/code-reviews/Driver.TwinCAT/findings.md
index 2afc0ff..04de94f 100644
--- a/code-reviews/Driver.TwinCAT/findings.md
+++ b/code-reviews/Driver.TwinCAT/findings.md
@@ -248,7 +248,7 @@ them on rebuild, or introduce a lifecycle lock / `volatile` running guard so rea
 with `BadServerHalted`/`BadNodeIdUnknown` once shutdown begins. Cancel and await the probe
 tasks before disposing `DeviceState`s.
 
-**Resolution:** _(open)_
+**Resolution:** Resolved 2026-05-22 — swapped `_devices` and `_tagsByName` to `ConcurrentDictionary` so concurrent `TryGetValue` / `Clear` calls are safe; added `DeviceState.ProbeTask` and updated `ShutdownAsync` to cancel then `await` each probe task before disposing the client and gate, eliminating the disposal race.
 
 ### Driver.TwinCAT-010
 
diff --git a/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.TwinCAT/TwinCATDriver.cs b/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.TwinCAT/TwinCATDriver.cs
index b4035c9..2236bfc 100644
--- a/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.TwinCAT/TwinCATDriver.cs
+++ b/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.TwinCAT/TwinCATDriver.cs
@@ -20,8 +20,13 @@ public sealed class TwinCATDriver : IDriver, IReadable, IWritable, ITagDiscovery
     private readonly ITwinCATClientFactory _clientFactory;
     private readonly ILogger<TwinCATDriver> _logger;
     private readonly PollGroupEngine _poll;
-    private readonly Dictionary<string, DeviceState> _devices = new(StringComparer.OrdinalIgnoreCase);
-    private readonly Dictionary<string, TwinCATTagDefinition> _tagsByName = new(StringComparer.OrdinalIgnoreCase);
+    // ConcurrentDictionary so ShutdownAsync (Clear) and ReadAsync/WriteAsync/SubscribeAsync
+    // (TryGetValue) don't race — plain Dictionary is not safe for concurrent read+write
+    // (Driver.TwinCAT-009).
+    private readonly ConcurrentDictionary<string, DeviceState> _devices =
+        new(StringComparer.OrdinalIgnoreCase);
+    private readonly ConcurrentDictionary<string, TwinCATTagDefinition> _tagsByName =
+        new(StringComparer.OrdinalIgnoreCase);
     private DriverHealth _health = new(DriverState.Unknown, null, null);
 
     public event EventHandler<DataChangeEventArgs>? OnDataChange;
@@ -76,7 +81,8 @@ public sealed class TwinCATDriver : IDriver, IReadable, IWritable, ITagDiscovery
                 {
                     state.ProbeCts = new CancellationTokenSource();
                     var ct = state.ProbeCts.Token;
-                    _ = Task.Run(() => ProbeLoopAsync(state, ct), ct);
+                    // Store the task so ShutdownAsync can await it (Driver.TwinCAT-009).
+                    state.ProbeTask = Task.Run(() => ProbeLoopAsync(state, ct), ct);
                 }
             }
             _health = new DriverHealth(DriverState.Healthy, DateTime.UtcNow, null);
@@ -104,9 +110,19 @@ public sealed class TwinCATDriver : IDriver, IReadable, IWritable, ITagDiscovery
         _nativeSubs.Clear();
 
         await _poll.DisposeAsync().ConfigureAwait(false);
+
+        // Cancel every probe loop and await its task before disposing the client + gate so the
+        // loop can never touch a disposed object. (Driver.TwinCAT-009: ShutdownAsync previously
+        // cancelled ProbeCts but did not await the task before calling DisposeClient.)
         foreach (var state in _devices.Values)
         {
             try { state.ProbeCts?.Cancel(); } catch { }
+            if (state.ProbeTask is Task pt)
+            {
+                try { await pt.ConfigureAwait(false); }
+                catch (OperationCanceledException) { /* expected — probe loop exits on cancel */ }
+                catch { /* other probe errors are not fatal to shutdown */ }
+            }
             state.ProbeCts?.Dispose();
             state.ProbeCts = null;
             state.DisposeClient();
@@ -526,6 +542,9 @@ public sealed class TwinCATDriver : IDriver, IReadable, IWritable, ITagDiscovery
         public HostState HostState { get; set; } = HostState.Unknown;
         public DateTime HostStateChangedUtc { get; set; } = DateTime.UtcNow;
         public CancellationTokenSource? ProbeCts { get; set; }
+        /// <summary>The running probe-loop task — awaited by <see cref="TwinCATDriver.ShutdownAsync"/>
+        ///     so the loop cannot touch a disposed client (Driver.TwinCAT-009).</summary>
+        public Task? ProbeTask { get; set; }
 
         public void DisposeClient()
         {