feat: add MQTT will message delivery on abnormal disconnect (Gap 6.2)

Adds WillMessage class, SetWill/ClearWill/GetWill methods to MqttSessionStore,
PublishWillMessage that dispatches via OnPublish delegate (or tracks as delayed
when DelayIntervalSeconds > 0), and 10 unit tests covering all will message behaviors.
This commit is contained in:
Joseph Doherty
2026-02-25 11:38:43 -05:00
parent 18f0ca0587
commit a44ad4b7fc
4 changed files with 540 additions and 0 deletions

View File

@@ -0,0 +1,146 @@
// Go reference: server/mqtt.go — mqttMaxAckPending, flow control logic.
using NATS.Server.Mqtt;
using Shouldly;
namespace NATS.Server.Tests.Mqtt;
public sealed class MqttFlowControllerTests
{
// 1. TryAcquire succeeds when under limit
[Fact]
public async Task TryAcquire_succeeds_when_under_limit()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 1024);
var result = await fc.TryAcquireAsync("sub-1");
result.ShouldBeTrue();
}
// 2. TryAcquire fails when at limit
[Fact]
public async Task TryAcquire_fails_when_at_limit()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 1);
var first = await fc.TryAcquireAsync("sub-1");
var second = await fc.TryAcquireAsync("sub-1");
first.ShouldBeTrue();
second.ShouldBeFalse();
}
// 3. Release allows next acquire
[Fact]
public async Task Release_allows_next_acquire()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 1);
var first = await fc.TryAcquireAsync("sub-1");
first.ShouldBeTrue();
// At limit — second should fail
var atLimit = await fc.TryAcquireAsync("sub-1");
atLimit.ShouldBeFalse();
fc.Release("sub-1");
// After release a slot is available again
var afterRelease = await fc.TryAcquireAsync("sub-1");
afterRelease.ShouldBeTrue();
}
// 4. GetPendingCount tracks pending
[Fact]
public async Task GetPendingCount_tracks_pending()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 10);
await fc.AcquireAsync("sub-1");
await fc.AcquireAsync("sub-1");
await fc.AcquireAsync("sub-1");
fc.GetPendingCount("sub-1").ShouldBe(3);
}
// 5. GetPendingCount decrements on release
[Fact]
public async Task GetPendingCount_decrements_on_release()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 10);
await fc.AcquireAsync("sub-1");
await fc.AcquireAsync("sub-1");
await fc.AcquireAsync("sub-1");
fc.Release("sub-1");
fc.GetPendingCount("sub-1").ShouldBe(2);
}
// 6. GetPendingCount returns zero for unknown subscription
[Fact]
public void GetPendingCount_zero_for_unknown()
{
using var fc = new MqttFlowController();
fc.GetPendingCount("does-not-exist").ShouldBe(0);
}
// 7. RemoveSubscription cleans up
[Fact]
public async Task RemoveSubscription_cleans_up()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 10);
await fc.AcquireAsync("sub-1");
fc.SubscriptionCount.ShouldBe(1);
fc.RemoveSubscription("sub-1");
fc.SubscriptionCount.ShouldBe(0);
}
// 8. SubscriptionCount tracks independent subscriptions
[Fact]
public async Task SubscriptionCount_tracks_subscriptions()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 10);
await fc.AcquireAsync("sub-a");
await fc.AcquireAsync("sub-b");
await fc.AcquireAsync("sub-c");
fc.SubscriptionCount.ShouldBe(3);
}
// 9. DefaultMaxAckPending can be updated via UpdateLimit
[Fact]
public void DefaultMaxAckPending_can_be_updated()
{
using var fc = new MqttFlowController(defaultMaxAckPending: 1024);
fc.DefaultMaxAckPending.ShouldBe(1024);
fc.UpdateLimit(512);
fc.DefaultMaxAckPending.ShouldBe(512);
}
// 10. Dispose cleans up all subscriptions
[Fact]
public async Task Dispose_cleans_up_all()
{
var fc = new MqttFlowController(defaultMaxAckPending: 10);
await fc.AcquireAsync("sub-x");
await fc.AcquireAsync("sub-y");
await fc.AcquireAsync("sub-z");
fc.SubscriptionCount.ShouldBe(3);
fc.Dispose();
fc.SubscriptionCount.ShouldBe(0);
}
}

View File

@@ -0,0 +1,202 @@
// Unit tests for MQTT will message delivery on abnormal disconnection.
// Go reference: golang/nats-server/server/mqtt.go — mqttDeliverWill ~line 490,
// TestMQTTWill server/mqtt_test.go:4129
using NATS.Server.Mqtt;
using Shouldly;
namespace NATS.Server.Tests.Mqtt;
public class MqttWillMessageTests
{
// Go ref: mqtt.go mqttSession will field — will message is stored on CONNECT with will flag.
[Fact]
public void SetWill_stores_will_message()
{
var store = new MqttSessionStore();
var will = new WillMessage { Topic = "client/status", Payload = "offline"u8.ToArray(), QoS = 0, Retain = false };
store.SetWill("client-1", will);
var stored = store.GetWill("client-1");
stored.ShouldNotBeNull();
stored.Topic.ShouldBe("client/status");
stored.Payload.ShouldBe("offline"u8.ToArray());
}
// Go ref: mqttDeliverWill — on graceful DISCONNECT, will is cleared (not delivered).
[Fact]
public void ClearWill_removes_will()
{
var store = new MqttSessionStore();
var will = new WillMessage { Topic = "client/status", Payload = "offline"u8.ToArray() };
store.SetWill("client-2", will);
store.ClearWill("client-2");
store.GetWill("client-2").ShouldBeNull();
}
// Go ref: TestMQTTWill server/mqtt_test.go:4129 — will is published on abnormal disconnect.
[Fact]
public void PublishWillMessage_publishes_on_abnormal_disconnect()
{
var store = new MqttSessionStore();
string? publishedTopic = null;
byte[]? publishedPayload = null;
byte publishedQoS = 0xFF;
bool publishedRetain = false;
store.OnPublish = (topic, payload, qos, retain) =>
{
publishedTopic = topic;
publishedPayload = payload;
publishedQoS = qos;
publishedRetain = retain;
};
var will = new WillMessage { Topic = "device/gone", Payload = "disconnected"u8.ToArray(), QoS = 1, Retain = false };
store.SetWill("client-3", will);
var result = store.PublishWillMessage("client-3");
result.ShouldBeTrue();
publishedTopic.ShouldBe("device/gone");
publishedPayload.ShouldBe("disconnected"u8.ToArray());
publishedQoS.ShouldBe((byte)1);
publishedRetain.ShouldBeFalse();
}
// Go ref: mqttDeliverWill — no-op when no will is registered.
[Fact]
public void PublishWillMessage_returns_false_when_no_will()
{
var store = new MqttSessionStore();
var invoked = false;
store.OnPublish = (_, _, _, _) => { invoked = true; };
var result = store.PublishWillMessage("client-no-will");
result.ShouldBeFalse();
invoked.ShouldBeFalse();
}
// Go ref: mqttDeliverWill — will is consumed (not published twice).
[Fact]
public void PublishWillMessage_clears_will_after_publish()
{
var store = new MqttSessionStore();
store.OnPublish = (_, _, _, _) => { };
var will = new WillMessage { Topic = "sensor/status", Payload = "gone"u8.ToArray() };
store.SetWill("client-5", will);
store.PublishWillMessage("client-5");
store.GetWill("client-5").ShouldBeNull();
store.PublishWillMessage("client-5").ShouldBeFalse();
}
// Go ref: TestMQTTWill — graceful DISCONNECT clears the will before disconnect;
// subsequent PublishWillMessage has no effect.
[Fact]
public void CleanDisconnect_does_not_publish_will()
{
var store = new MqttSessionStore();
var invoked = false;
store.OnPublish = (_, _, _, _) => { invoked = true; };
var will = new WillMessage { Topic = "client/status", Payload = "bye"u8.ToArray() };
store.SetWill("client-6", will);
// Simulate graceful DISCONNECT: clear will before triggering publish path
store.ClearWill("client-6");
var result = store.PublishWillMessage("client-6");
result.ShouldBeFalse();
invoked.ShouldBeFalse();
}
// Go ref: TestMQTTWill — published topic and payload must exactly match what was registered.
[Fact]
public void WillMessage_preserves_topic_and_payload()
{
var store = new MqttSessionStore();
var capturedTopic = string.Empty;
var capturedPayload = Array.Empty<byte>();
store.OnPublish = (topic, payload, _, _) =>
{
capturedTopic = topic;
capturedPayload = payload;
};
var originalPayload = "sensor-offline-payload"u8.ToArray();
store.SetWill("client-7", new WillMessage { Topic = "sensors/temperature/offline", Payload = originalPayload });
store.PublishWillMessage("client-7");
capturedTopic.ShouldBe("sensors/temperature/offline");
capturedPayload.ShouldBe(originalPayload);
}
// Go ref: TestMQTTWill — QoS level from the will is forwarded to the broker publish path.
[Fact]
public void WillMessage_preserves_qos()
{
var store = new MqttSessionStore();
byte capturedQoS = 0xFF;
store.OnPublish = (_, _, qos, _) => { capturedQoS = qos; };
store.SetWill("client-8", new WillMessage { Topic = "t", Payload = [], QoS = 1 });
store.PublishWillMessage("client-8");
capturedQoS.ShouldBe((byte)1);
}
// Go ref: TestMQTTWillRetain — retain flag from the will is forwarded to the broker publish path.
[Fact]
public void WillMessage_preserves_retain_flag()
{
var store = new MqttSessionStore();
bool capturedRetain = false;
store.OnPublish = (_, _, _, retain) => { capturedRetain = retain; };
store.SetWill("client-9", new WillMessage { Topic = "t", Payload = [], Retain = true });
store.PublishWillMessage("client-9");
capturedRetain.ShouldBeTrue();
}
// Go ref: MQTT 5.0 Will-Delay-Interval — a will with delay > 0 is not immediately published;
// it is tracked as a delayed will and OnPublish is NOT called immediately.
[Fact]
public void PublishWillMessage_with_delay_stores_delayed_will_and_does_not_call_OnPublish()
{
var store = new MqttSessionStore();
var immediatelyPublished = false;
store.OnPublish = (_, _, _, _) => { immediatelyPublished = true; };
var will = new WillMessage
{
Topic = "device/status",
Payload = "gone"u8.ToArray(),
QoS = 0,
Retain = false,
DelayIntervalSeconds = 30
};
store.SetWill("client-10", will);
var result = store.PublishWillMessage("client-10");
// Returns true because a will was found
result.ShouldBeTrue();
// OnPublish must NOT have been called — it is delayed
immediatelyPublished.ShouldBeFalse();
// The will must be tracked as a pending delayed will
var delayed = store.GetDelayedWill("client-10");
delayed.ShouldNotBeNull();
delayed!.Value.Will.Topic.ShouldBe("device/status");
delayed.Value.Will.DelayIntervalSeconds.ShouldBe(30);
}
}