fix(mqtt): dispose an unpublished connection and stamp lifecycle generations (R06:M-3, M-10)

A refused initial connect (or a failed Sparkplug establish) now disposes the MqttConnection it built. Every lifecycle operation takes a generation (initialize inside the gate; shutdown/dispose before waiting for it): a wedged initialize that completes after a timed-out ShutdownAsync's ungated teardown releases its session instead of publishing it, and a superseded operation's health/host-state writes are refused. TeardownAsync captures every field before disposing so an ungated teardown never reaches into a later initialize's state. Refs #744 Claude-Session: https://claude.ai/code/session_01Ub9hGmYLQvn9XETzw7CMsT
2026-10-22 15:26:55 -04:00
parent 01efd39177
commit 43f45bb978
2 changed files with 186 additions and 22 deletions
@@ -208,6 +208,15 @@ public sealed class MqttDriver
    private SparkplugPrimaryHost? _primaryHost;

    private DriverHealth _health = new(DriverState.Unknown, null, null);
+
+    /// <summary>Guards the generation bump against the generation-checked health write.</summary>
+    private readonly object _healthLock = new();
+
+    /// <summary>
+    ///     Bumped by every lifecycle operation (initialize inside the gate; shutdown and dispose before
+    ///     waiting for it). The stamp that keeps a superseded operation from publishing state.
+    /// </summary>
+    private long _lifecycleGeneration;
    private HostState _hostState = HostState.Unknown;
    private DateTime _hostStateChangedUtc = DateTime.UtcNow;
    private int _disposed;
@@ -289,6 +298,9 @@ public sealed class MqttDriver
    /// </summary>
    internal Func<CancellationToken, Task>? BeforeConnectHookForTests { get; set; }

+    /// <summary>Test seam: observes every <see cref="MqttConnection"/> an initialize builds. Null in production.</summary>
+    internal Action<MqttConnection>? ConnectionCreatedForTests { get; set; }
+
    // ---- IDriver: lifecycle ----

    /// <inheritdoc />
@@ -387,6 +399,11 @@ public sealed class MqttDriver
        // watchdog's reconnect) is measured against, reporting a driver that was delivering a moment ago
        // as one that has never delivered.
        var lastMessage = GetHealth().LastSuccessfulRead;
+
+        // Taken BEFORE the gate wait: a wedged initialize holding the gate is superseded by this call, so
+        // when it finally completes it disposes the session it built instead of publishing it behind the
+        // ungated teardown below (R06:M-10), and its late health write is refused.
+        var generation = BeginLifecycleOperation();
        if (await _lifecycleGate.WaitAsync(TimeSpan.FromSeconds(_options.ConnectTimeoutSeconds), cancellationToken)
                .ConfigureAwait(false))
        {
@@ -408,8 +425,10 @@ public sealed class MqttDriver
            await TeardownAsync().ConfigureAwait(false);
        }

-        WriteHealth(new DriverHealth(DriverState.Unknown, lastMessage, null));
-        TransitionHost(HostState.Unknown);
+        if (WriteHealthIfCurrent(generation, new DriverHealth(DriverState.Unknown, lastMessage, null)))
+        {
+            TransitionHost(HostState.Unknown);
+        }
    }

    /// <inheritdoc />
@@ -911,6 +930,9 @@ public sealed class MqttDriver
            return;
        }

+        // Supersedes any in-flight lifecycle operation, exactly as ShutdownAsync does.
+        BeginLifecycleOperation();
+
        if (await _lifecycleGate.WaitAsync(TimeSpan.FromSeconds(_options.ConnectTimeoutSeconds))
                .ConfigureAwait(false))
        {
@@ -940,7 +962,12 @@ public sealed class MqttDriver
    /// <summary>The gate-held body of <see cref="InitializeAsync"/>.</summary>
    private async Task InitializeCoreAsync(string driverConfigJson, CancellationToken cancellationToken)
    {
-        WriteHealth(new DriverHealth(DriverState.Initializing, ReadHealth().LastSuccessfulRead, null));
+        // Taken INSIDE the gate (the caller holds it), so a second initialize queued behind this one does
+        // not supersede it; only a ShutdownAsync / DisposeAsync — which take theirs before waiting — can.
+        var generation = BeginLifecycleOperation();
+        WriteHealthIfCurrent(generation, new DriverHealth(DriverState.Initializing, ReadHealth().LastSuccessfulRead, null));
+        MqttConnection? built = null;
+        var published = false;
        try
        {
            // Tear down whatever this driver already holds BEFORE building a new session. Both fields
@@ -965,6 +992,8 @@ public sealed class MqttDriver
            RegisterAuthoredTags(_options);

            var connection = new MqttConnection(_options, _driverInstanceId, _logger);
+            built = connection;
+            ConnectionCreatedForTests?.Invoke(connection);

            // MUST precede the connect: this wires message delivery AND the reconnect re-subscribe.
            // The handler is passed through unwrapped on purpose — its throw-on-total-failure is what
@@ -992,10 +1021,13 @@ public sealed class MqttDriver
            // disposing what we just built, never leaking it.
            if (Volatile.Read(ref _disposed) != 0)
            {
-                await connection.DisposeAsync().ConfigureAwait(false);
                throw new ObjectDisposedException(nameof(MqttDriver));
            }

+            // The same race against a ShutdownAsync whose bounded gate wait expired: it has already torn
+            // down ungated, so publishing now would leave a live session no later shutdown can see.
+            ThrowIfSuperseded(generation);
+
            // Sparkplug subscribes ONCE, here, rather than on the OPC UA server's first subscribe: the
            // driver must see birth certificates whether or not anything is subscribed yet, because a
            // birth is where every alias and datatype comes from. This also issues the late-join
@@ -1008,28 +1040,44 @@ public sealed class MqttDriver
            // MqttConnection's own remarks describe having already fixed once.
            if (_sparkplug is { } sparkplug)
            {
-                try
-                {
-                    await sparkplug.EstablishAsync(cancellationToken).ConfigureAwait(false);
-                }
-                catch
-                {
-                    await connection.DisposeAsync().ConfigureAwait(false);
-                    throw;
-                }
+                await sparkplug.EstablishAsync(cancellationToken).ConfigureAwait(false);
            }

+            ThrowIfSuperseded(generation);
            _connection = connection;
+            published = true;
+
+            if (WriteHealthIfCurrent(generation, new DriverHealth(DriverState.Healthy, connection.LastMessageUtc, null)))
+            {
+                TransitionHost(HostState.Running);
+            }

-            WriteHealth(new DriverHealth(DriverState.Healthy, connection.LastMessageUtc, null));
-            TransitionHost(HostState.Running);
            StartHostProbe();
            await StartPrimaryHostAsync(cancellationToken).ConfigureAwait(false);
+
+            // Re-checked after publishing: a shutdown whose gate wait expired between the check above and
+            // here captured whatever it found and left these fields for us. We still hold the gate, so no
+            // newer initialize can have published anything — everything in them is ours to release.
+            if (!IsCurrentLifecycle(generation))
+            {
+                await TeardownAsync().ConfigureAwait(false);
+                ThrowIfSuperseded(generation);
+            }
        }
        catch (Exception ex)
        {
-            WriteHealth(new DriverHealth(DriverState.Faulted, ReadHealth().LastSuccessfulRead, ex.Message));
-            TransitionHost(HostState.Faulted);
+            // R06:M-3: a connection built but never published is disposed here — a refused connect, a
+            // failed Sparkplug establish, a supersede. The published one belongs to TeardownAsync.
+            if (!published && built is not null)
+            {
+                await built.DisposeAsync().ConfigureAwait(false);
+            }
+
+            if (WriteHealthIfCurrent(generation, new DriverHealth(DriverState.Faulted, ReadHealth().LastSuccessfulRead, ex.Message)))
+            {
+                TransitionHost(HostState.Faulted);
+            }
+
            throw;
        }
    }
@@ -1380,24 +1428,28 @@ public sealed class MqttDriver
    /// <summary>Shared teardown for <see cref="ShutdownAsync"/>, <see cref="DisposeAsync"/> and the session rebuild.</summary>
    private async Task TeardownAsync()
    {
+        // Every field is CAPTURED before anything is disposed. A shutdown that timed out on the gate runs
+        // this ungated; capturing up front means it only ever releases what existed when it started, and
+        // never reaches into fields a later initialize publishes while the slow disposals below run.
+        var primaryHost = Interlocked.Exchange(ref _primaryHost, null);
+        var probe = Interlocked.Exchange(ref _hostProbeCts, null);
+        var connection = Interlocked.Exchange(ref _connection, null);
+
        // Stand the host claim down FIRST, while the broker is still reachable and before the ingest
        // session goes. Its explicit retained OFFLINE is what makes the transition immediate for edge
        // nodes; leaving it to the Will costs them up to 1.5 keep-alive periods of believing a host
        // that has already stopped is still online.
-        var primaryHost = Interlocked.Exchange(ref _primaryHost, null);
        if (primaryHost is not null)
        {
            await primaryHost.DisposeAsync().ConfigureAwait(false);
        }

-        var probe = Interlocked.Exchange(ref _hostProbeCts, null);
        if (probe is not null)
        {
            await probe.CancelAsync().ConfigureAwait(false);
            probe.Dispose();
        }

-        var connection = Interlocked.Exchange(ref _connection, null);
        if (connection is not null)
        {
            await connection.DisposeAsync().ConfigureAwait(false);
@@ -1479,6 +1531,49 @@ public sealed class MqttDriver
    /// <summary>Barrier-protected read of the multi-threaded health field.</summary>
    private DriverHealth ReadHealth() => Volatile.Read(ref _health);

-    /// <summary>Barrier-protected publish of a new health snapshot.</summary>
-    private void WriteHealth(DriverHealth value) => Volatile.Write(ref _health, value);
+    /// <summary>
+    ///     Starts a lifecycle operation and returns its generation. Every later operation supersedes it:
+    ///     its health / host-state writes are refused and an initialize stops before publishing a session.
+    /// </summary>
+    private long BeginLifecycleOperation()
+    {
+        lock (_healthLock)
+        {
+            return ++_lifecycleGeneration;
+        }
+    }
+
+    /// <summary>Whether <paramref name="generation"/> is still the newest lifecycle operation.</summary>
+    private bool IsCurrentLifecycle(long generation) => Interlocked.Read(ref _lifecycleGeneration) == generation;
+
+    /// <summary>Throws when a newer lifecycle operation has superseded <paramref name="generation"/>.</summary>
+    private void ThrowIfSuperseded(long generation)
+    {
+        if (!IsCurrentLifecycle(generation))
+        {
+            throw new OperationCanceledException(
+                $"MQTT driver '{_driverInstanceId}': this initialize was superseded by a later shutdown/dispose; "
+                + "the session it built has been released.");
+        }
+    }
+
+    /// <summary>
+    ///     Publishes a health snapshot only while <paramref name="generation"/> is the newest lifecycle
+    ///     operation — the check and the write are atomic against <see cref="BeginLifecycleOperation"/>, so a
+    ///     superseded initialize or shutdown can never overwrite a newer one's state.
+    /// </summary>
+    /// <returns><see langword="true"/> when the write was applied.</returns>
+    private bool WriteHealthIfCurrent(long generation, DriverHealth value)
+    {
+        lock (_healthLock)
+        {
+            if (_lifecycleGeneration != generation)
+            {
+                return false;
+            }
+
+            Volatile.Write(ref _health, value);
+            return true;
+        }
+    }
 }
@@ -219,6 +219,75 @@ public sealed class MqttDriverSessionLifecycleTests
        }
    }

+    /// <summary>
+    ///     <b>R06:M-3.</b> A refused initial connect used to rethrow without disposing the
+    ///     <see cref="MqttConnection"/> (and its <c>IMqttClient</c>) it had just built. Control: a
+    ///     successful initialize leaves its connection live, so the seam is not reporting "disposed" for
+    ///     everything.
+    /// </summary>
+    [Fact]
+    public async Task InitializeAsync_RefusedConnect_DisposesTheConnectionItBuilt()
+    {
+        using var broker = new MiniBroker { ConnAckReturnCode = MiniBroker.ReturnCodeNotAuthorized };
+        var options = BrokerOptions(broker.Port);
+        await using var driver = new MqttDriver(options, "d", null);
+        var built = new List<MqttConnection>();
+        driver.ConnectionCreatedForTests = c => built.Add(c);
+
+        await Should.ThrowAsync<Exception>(() => driver.InitializeAsync(Json(options), TestContext.Current.CancellationToken));
+
+        built.ShouldHaveSingleItem().State.ShouldBe(MqttConnectionState.Disposed);
+
+        broker.ConnAckReturnCode = MiniBroker.ReturnCodeAccepted;
+        await driver.InitializeAsync(Json(options), TestContext.Current.CancellationToken);
+        built.Count.ShouldBe(2);
+        built[1].State.ShouldBe(MqttConnectionState.Connected);
+    }
+
+    /// <summary>
+    ///     <b>R06:M-10 + the lifecycle-generation rule.</b> A shutdown that times out on the lifecycle gate
+    ///     tears down ungated; the wedged initialize that held the gate must then NOT publish its freshly
+    ///     connected session behind that teardown, and its late health write must not overwrite the
+    ///     shutdown's. The positive controls: the wedged init really did reach the broker (AcceptedCount),
+    ///     and a later initialize connects and reports Healthy.
+    /// </summary>
+    [Trait("Category", "SlowUnit")]
+    [Fact]
+    public async Task AWedgedInit_CompletingAfterATimedOutShutdown_PublishesNothing()
+    {
+        using var broker = new MiniBroker();
+        var options = BrokerOptions(broker.Port) with { ConnectTimeoutSeconds = 1 };
+        await using var driver = new MqttDriver(options, "d", null);
+
+        var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        driver.BeforeConnectHookForTests = async _ =>
+        {
+            entered.TrySetResult();
+            await release.Task;
+        };
+
+        var initialize = Task.Run(() => driver.InitializeAsync(Json(options), CancellationToken.None));
+        await entered.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
+
+        // Times out on the gate (1 s) and tears down ungated.
+        await driver.ShutdownAsync(TestContext.Current.CancellationToken);
+        driver.GetHealth().State.ShouldBe(DriverState.Unknown);
+
+        driver.BeforeConnectHookForTests = null;
+        release.SetResult();
+        await Should.ThrowAsync<OperationCanceledException>(() => initialize);
+
+        broker.AcceptedCount.ShouldBe(1, "the wedged initialize did connect — the control for the absence below");
+        await WaitForAsync(() => broker.LiveConnections == 0, "the superseded session to be torn down");
+        driver.GetHealth().State.ShouldBe(DriverState.Unknown, "a superseded init must not overwrite newer health");
+        driver.GetHostStatuses().ShouldHaveSingleItem().State.ShouldBe(HostState.Unknown);
+
+        await driver.InitializeAsync(Json(options), TestContext.Current.CancellationToken);
+        await WaitForAsync(() => broker.LiveConnections == 1, "the fresh session");
+        driver.GetHealth().State.ShouldBe(DriverState.Healthy);
+    }
+
    private static async Task WaitForAsync(Func<bool> condition, string what)
    {
        var deadline = DateTime.UtcNow + TimeSpan.FromSeconds(10);