MQTTnet/Source/MQTTnet.Tests/Server/Subscription_TopicHash_Tests.cs

using System.Diagnostics;
using System.Linq;
using System.Text;
using MQTTnet.Packets;
using MQTTnet.Protocol;
using MQTTnet.Server;
using MQTTnet.Server.Internal;
using MQTTnet.Tests.Mockups;

namespace MQTTnet.Tests.Server;

[TestClass]
public sealed class SubscriptionTopicHashTests : IDisposable
{
    MqttSession _clientSession;

    /// <summary>
    ///     Even fairly regularly named topics as generated by the topic generator should result in shallow hash buckets
    /// </summary>
    /// <returns></returns>
    [TestMethod]
    public void Check_Hash_Bucket_Depth()
    {
        const int numPublishers = 5000;
        const int numTopicsPerPublisher = 10;

        TopicGenerator.Generate(numPublishers, numTopicsPerPublisher, out var topicsByPublisher, out _, out _);

        // There will be many 'similar' topics ending with, i.e. "sensor100", "sensor101", ...
        // Hash bucket depths should remain low.
        var bucketDepths = new Dictionary<ulong, int>();
        var maxBucketDepth = 0;
        ulong maxBucketDepthHash = 0;

        var topicsByHash = new Dictionary<ulong, List<string>>();

        foreach (var t in topicsByPublisher)
        {
            var topics = t.Value;
            foreach (var topic in topics)
            {
                MqttTopicHash.Calculate(topic, out var topicHash, out _, out _);

                bucketDepths.TryGetValue(topicHash, out var currentValue);
                ++currentValue;
                bucketDepths[topicHash] = currentValue;

                if (currentValue > maxBucketDepth)
                {
                    maxBucketDepth = currentValue;
                    maxBucketDepthHash = topicHash;
                }

                if (!topicsByHash.TryGetValue(topicHash, out var topicList))
                {
                    topicList = [];
                    topicsByHash.Add(topicHash, topicList);
                }

                topicList.Add(topic);
            }
        }

        var maxDepthTopics = topicsByHash[maxBucketDepthHash];

        Console.Write("Max bucket depth is " + maxBucketDepth);

        // for the test case the bucket depth should be less than 100
        Assert.IsLessThan(100, maxBucketDepth, "Unexpected high topic hash bucket depth");
    }


    [TestMethod]
    public void Check_Selected_Topic_Hashes()
    {
        CheckTopicHash("client1/building1/level1/sensor1", 0x655D4AF100000000, 0xFFFFFFFFFFFFFFFF);
        CheckTopicHash("client1/building1/+/sensor1", 0x655D00F100000000, 0xFFFF00FFFFFFFFFF);
        CheckTopicHash("client1/+/level1/+", 0x65004A0000000000, 0xFF00FF00FFFFFFFF);
        CheckTopicHash("client1/building1/level1/#", 0x655D4A0000000000, 0xFFFFFF0000000000);
        CheckTopicHash("client1/+/level1/#", 0x65004A0000000000, 0xFF00FF0000000000);
    }

    public void Dispose()
    {
        _clientSession?.Dispose();
    }

    /// <summary>
    ///     Test long topic name with last level being #
    /// </summary>
    [TestMethod]
    public void Match_Hash_Test_LongTopic_DetectMultiLevelWildcard()
    {
        var topic = "asdfasdf/asdfasdf/asdfasdf/asdfasdf/asdfas/dfaf/assfdgsdfgdf/#";

        MqttTopicHash.Calculate(topic, out _, out _, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");
    }


    /// <summary>
    ///     Test long topic name with last level being +
    /// </summary>
    [TestMethod]
    public void Match_Hash_Test_LongTopic_DetectSingleLevelWildcard()
    {
        var topic = "asdfasdf/asdfasdf/asdfasdf/asdfasdf/asdfas/dfaf/assfdgsdfgdf/+";

        MqttTopicHash.Calculate(topic, out _, out _, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");
    }

    /// <summary>
    ///     Test long topic name with last level being #
    /// </summary>
    [TestMethod]
    public void Match_Hash_Test_LongTopic_MultiWildcard()
    {
        var sb = new StringBuilder();
        const int numLevels = 8;
        var levelNames = new string[numLevels];
        for (var i = 0; i < numLevels; ++i)
        {
            if (i > 0)
            {
                sb.Append('/');
            }

            string levelName;
            if (i == numLevels - 1)
            {
                // last one is #
                levelName = "#";
            }
            else
            {
                levelName = "level" + i;
            }

            levelNames[i] = levelName;
            sb.Append(levelName);
        }

        var topic = sb.ToString();

        // UInt64 is limited to 8 levels

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");

        var hashBytes = GetBytes(topicHash);
        // all bytes should contain checksum
        var count = 0;
        foreach (var h in hashBytes)
        {
            if (count < 7)
            {
                Assert.AreNotEqual(0, h, "checksum mismatch");
            }
            else
            {
                Assert.AreEqual(0, h, "checksum mismatch");
            }

            ++count;
        }

        // The mask should have ff except for last level
        var hashMaskBytes = GetBytes(topicHashMask);
        count = 0;
        foreach (var h in hashMaskBytes)
        {
            if (count < 7)
            {
                Assert.AreEqual(0xff, h, "mask mismatch");
            }
            else
            {
                Assert.AreEqual(0, h, "last mask mismatch");
            }

            ++count;
        }
    }

    /// <summary>
    ///     Test long topic name exceeding 8 levels
    /// </summary>
    [TestMethod]
    public void Match_Hash_Test_LongTopic_NoWildCard()
    {
        var sb = new StringBuilder();
        const int numLevels = 9;
        var levelNames = new string[numLevels];
        for (var i = 0; i < numLevels; ++i)
        {
            if (i > 0)
            {
                sb.Append('/');
            }

            var levelName = "level" + i;
            levelNames[i] = levelName;
            sb.Append(levelName);
        }

        var topic = sb.ToString();

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsFalse(topicHasWildcard, "Wildcard detected when not present");


        var hashBytes = GetBytes(topicHash);
        // all bytes should contain checksum
        var count = 0;
        foreach (var h in hashBytes)
        {
            Assert.AreNotEqual(0, h, "checksum mismatch");
            ++count;
        }

        // The mask should have ff
        var hashMaskBytes = GetBytes(topicHashMask);
        count = 0;
        foreach (var h in hashMaskBytes)
        {
            Assert.AreEqual(0xff, h, "mask mismatch");
            ++count;
        }
    }

    /// <summary>
    ///     Test long topic name with last level being +
    /// </summary>
    [TestMethod]
    public void Match_Hash_Test_LongTopic_SingleWildCard()
    {
        var sb = new StringBuilder();
        const int numLevels = 8;
        var levelNames = new string[numLevels];
        for (var i = 0; i < numLevels; ++i)
        {
            if (i > 0)
            {
                sb.Append('/');
            }

            string levelName;
            if (i == numLevels - 1)
            {
                // last one is +
                levelName = "+";
            }
            else
            {
                levelName = "level" + i;
            }

            levelNames[i] = levelName;
            sb.Append(levelName);
        }

        var topic = sb.ToString();

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");


        var hashBytes = GetBytes(topicHash);
        // all bytes should contain checksum
        var count = 0;
        foreach (var h in hashBytes)
        {
            if (count < 7)
            {
                Assert.AreNotEqual(0, h, "checksum mismatch");
            }
            else
            {
                // wildcard position
                Assert.AreEqual(0, h, "checksum mismatch");
            }

            ++count;
        }

        // The mask should have ff
        var hashMaskBytes = GetBytes(topicHashMask);
        count = 0;
        foreach (var h in hashMaskBytes)
        {
            if (count < 7)
            {
                Assert.AreEqual(0xff, h, "mask mismatch");
            }
            else
            {
                Assert.AreEqual(0, h, "last mask mismatch");
            }

            ++count;
        }
    }

    [TestMethod]
    public void Match_Hash_Test_MultiWildCard()
    {
        var l0 = "pub0";
        var l1 = "topic1";
        var l2 = "sub1";
        var l3 = "#";
        var topic = $"{l0}/{l1}/{l2}/{l3}";

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");

        var hashBytes = GetBytes(topicHash);
        Assert.AreNotEqual(0, hashBytes[0], "checksum 0 mismatch");
        Assert.AreNotEqual(0, hashBytes[1], "checksum 1 mismatch");
        Assert.AreNotEqual(0, hashBytes[2], "checksum 2 mismatch");
        Assert.AreEqual(0, hashBytes[3], "checksum 3 mismatch");
        Assert.AreEqual(0, hashBytes[4], "checksum 4 mismatch");
        Assert.AreEqual(0, hashBytes[5], "checksum 5 mismatch");
        Assert.AreEqual(0, hashBytes[6], "checksum 6 mismatch");
        Assert.AreEqual(0, hashBytes[7], "checksum 7 mismatch");

        // The mask should have zeroes where the wildcard is and zero onward
        var hashMaskBytes = GetBytes(topicHashMask);
        Assert.AreEqual(0xff, hashMaskBytes[0], "mask 0 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[1], "mask 1 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[2], "mask 2 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[3], "mask 3 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[4], "mask 4 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[5], "mask 5 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[6], "mask 6 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[7], "mask 7 mismatch");
    }

    [TestMethod]
    public void Match_Hash_Test_NoWildCard()
    {
        var l0 = "pub0";
        var l1 = "topic1";
        var l2 = "sub1";
        var l3 = "prop1";
        var topic = $"{l0}/{l1}/{l2}/{l3}";

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsFalse(topicHasWildcard, "Wildcard detected when not wildcard present");

        var hashBytes = GetBytes(topicHash);
        Assert.AreNotEqual(0, hashBytes[0], "checksum 0 mismatch");
        Assert.AreNotEqual(0, hashBytes[1], "checksum 1 mismatch");
        Assert.AreNotEqual(0, hashBytes[2], "checksum 2 mismatch");
        Assert.AreNotEqual(0, hashBytes[3], "checksum 3 mismatch");
        Assert.AreEqual(0, hashBytes[4], "checksum 4 mismatch");
        Assert.AreEqual(0, hashBytes[5], "checksum 5 mismatch");
        Assert.AreEqual(0, hashBytes[6], "checksum 6 mismatch");
        Assert.AreEqual(0, hashBytes[7], "checksum 7 mismatch");

        // The mask should have ff
        var hashMaskBytes = GetBytes(topicHashMask);
        Assert.AreEqual(0xff, hashMaskBytes[0], "mask 0 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[1], "mask 1 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[2], "mask 2 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[3], "mask 3 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[4], "mask 4 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[5], "mask 5 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[6], "mask 6 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[7], "mask 7 mismatch");
    }

    [TestMethod]
    public async Task Match_Hash_Test_Search_MultiWildcard()
    {
        var topics = await PrepareTopicHashSubscriptions(TopicHashSelector.MultiWildcard);
        var matchCount = CheckTopicSubscriptions(_clientSession, topics, "Multi Wildcard");
        // Should match all topics
        Assert.AreEqual(topics.Count, matchCount, "Topics not matched");
    }

    [TestMethod]
    public async Task Match_Hash_Test_Search_NoWildcard()
    {
        var topics = await PrepareTopicHashSubscriptions(TopicHashSelector.NoWildcard);

        // all match lookup
        var matchCount = CheckTopicSubscriptions(_clientSession, topics, "No Wildcard All Match");

        Assert.AreEqual(topics.Count, matchCount, "Not all topics matched");

        // no match lookup
        {
            {
                var topicsToFind = new List<string>();
                foreach (var t in topics)
                {
                    topicsToFind.Add(t + "x");
                }

                matchCount = CheckTopicSubscriptions(_clientSession, topicsToFind, "No Wildcard No Match (Append X)");

                Assert.AreEqual(0, matchCount, "Topic match count not zero");
            }

            {
                var topicsToFind = new List<string>();
                foreach (var t in topics)
                {
                    // replace last letter with x
                    topicsToFind.Add(string.Concat(t.AsSpan(0, t.Length - 1), "x"));
                }

                matchCount = CheckTopicSubscriptions(_clientSession, topicsToFind, "No Wildcard No Match (Replace X)");

                Assert.AreEqual(0, matchCount, "Topic match count not zero");
            }
        }
    }

    [TestMethod]
    public async Task Match_Hash_Test_Search_SingleWildcard()
    {
        var topics = await PrepareTopicHashSubscriptions(TopicHashSelector.SingleWildcard);
        var matchCount = CheckTopicSubscriptions(_clientSession, topics, "Single Wildcard");
        // Should match all topics
        Assert.AreEqual(topics.Count, matchCount, "Topics not matched");
    }

    [TestMethod]
    public void Match_Hash_Test_SingleWildCard()
    {
        var l0 = "pub0";
        var l1 = "topic1";
        var l2 = "+";
        var l3 = "prop1";
        var topic = $"{l0}/{l1}/{l2}/{l3}";

        MqttTopicHash.Calculate(topic, out var topicHash, out var topicHashMask, out var topicHasWildcard);

        Assert.IsTrue(topicHasWildcard, "Wildcard not detected");

        var hashBytes = GetBytes(topicHash);
        Assert.AreNotEqual(0, hashBytes[0], "checksum 0 mismatch");
        Assert.AreNotEqual(0, hashBytes[1], "checksum 1 mismatch");
        Assert.AreEqual(0, hashBytes[2], "checksum 2 mismatch");
        Assert.AreNotEqual(0, hashBytes[3], "checksum 3 mismatch");
        Assert.AreEqual(0, hashBytes[4], "checksum 4 mismatch");
        Assert.AreEqual(0, hashBytes[5], "checksum 5 mismatch");
        Assert.AreEqual(0, hashBytes[6], "checksum 6 mismatch");
        Assert.AreEqual(0, hashBytes[7], "checksum 7 mismatch");

        // The mask should have zeroes where the wildcard and ff at the end
        var hashMaskBytes = GetBytes(topicHashMask);
        Assert.AreEqual(0xff, hashMaskBytes[0], "mask 0 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[1], "mask 1 mismatch");
        Assert.AreEqual(0x0, hashMaskBytes[2], "mask 2 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[3], "mask 3 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[4], "mask 4 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[5], "mask 5 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[6], "mask 6 mismatch");
        Assert.AreEqual(0xff, hashMaskBytes[7], "mask 7 mismatch");
    }

    static void CheckTopicHash(string topic, ulong expectedHash, ulong expectedHashMask)
    {
        MqttTopicHash.Calculate(topic, out var topicHash, out var hashMask, out _);

        Console.WriteLine();
        Console.WriteLine("Topic: " + topic);
        Console.WriteLine($"Hash: {topicHash:X8}");
        Console.WriteLine($"Hash Mask: {hashMask:X8}");

        Assert.AreEqual(expectedHash, topicHash, "Topic hash not as expected. Has the hash function changed?");
        Assert.AreEqual(expectedHashMask, hashMask, "Topic hash mask not as expected");
    }

    static int CheckTopicSubscriptions(MqttSession clientSession, List<string> topicsToFind, string subject)
    {
        var matchCount = 0;

        {
            var resultCount = 0;

            var stopWatch = new Stopwatch();
            stopWatch.Start();
            var countUp = 0;
            var countDown = topicsToFind.Count - 1;
            for (; countUp < topicsToFind.Count; ++countUp, --countDown)
            {
                var topicToFind = topicsToFind[countDown];

                MqttTopicHash.Calculate(topicToFind, out var topicHash, out _, out _);

                clientSession.TryCheckSubscriptions(topicToFind, topicHash, MqttQualityOfServiceLevel.AtMostOnce, "OtherClient", out var result);
                if (result.IsSubscribed)
                {
                    ++matchCount;
                }

                ++resultCount;
            }

            stopWatch.Stop();

            Console.Write("Match count: " + matchCount + "; ");

            Console.WriteLine(subject + " lookup milliseconds: " + stopWatch.ElapsedMilliseconds);
        }

        return matchCount;
    }

    static byte[] GetBytes(ulong value)
    {
        var bytes = BitConverter.GetBytes(value);

        // Ensure that highest byte comes first for comparison left to right
        if (BitConverter.IsLittleEndian)
        {
            bytes = bytes.AsEnumerable().Reverse().ToArray();
        }

        return bytes;
    }


    async Task<List<string>> PrepareTopicHashSubscriptions(TopicHashSelector selector)
    {
        const int numPublishers = 1;
        const int numTopicsPerPublisher = 10000;

        TopicGenerator.Generate(numPublishers, numTopicsPerPublisher, out var topicsByPublisher, out var singleWildcardTopicsByPublisher, out var multiWildcardTopicsByPublisher);

        var topics = topicsByPublisher.FirstOrDefault().Value;
        var singleWildcardTopics = singleWildcardTopicsByPublisher.FirstOrDefault().Value;
        var multiWildcardTopics = multiWildcardTopicsByPublisher.FirstOrDefault().Value;

        const string clientId = "Client1";
        var logger = new TestLogger();
        var serverOptions = new MqttServerOptions();
        var eventContainer = new MqttServerEventContainer();
        var retainedMessagesManager = new MqttRetainedMessagesManager(eventContainer, logger);
        var sessionManager = new MqttClientSessionsManager(serverOptions, retainedMessagesManager, eventContainer, logger);
        _clientSession = new MqttSession(
            new MqttConnectPacket { ClientId = clientId },
            new Dictionary<object, object>(),
            serverOptions,
            eventContainer,
            retainedMessagesManager,
            sessionManager);

        List<string> topicsToSubscribe;

        switch (selector)
        {
            case TopicHashSelector.SingleWildcard:
                topicsToSubscribe = singleWildcardTopics;
                break;
            case TopicHashSelector.MultiWildcard:
                topicsToSubscribe = multiWildcardTopics;
                break;
            default:
                topicsToSubscribe = topics;
                break;
        }

        foreach (var t in topicsToSubscribe)
        {
            var subPacket = new MqttSubscribePacket();
            var filter = new MqttTopicFilter
            {
                Topic = t
            };
            subPacket.TopicFilters.Add(filter);
            await _clientSession.Subscribe(subPacket, default);
        }

        return topics;
    }

    enum TopicHashSelector
    {
        NoWildcard,
        SingleWildcard,
        MultiWildcard
    }
}