MQTTnet/Source/MQTTnet/PacketDispatcher/MqttPacketDispatcher.cs

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System;
using System.Collections.Generic;
using MQTTnet.Packets;

namespace MQTTnet.PacketDispatcher
{
    public sealed class MqttPacketDispatcher : IDisposable
    {
        public const int PingRespPacketFakeIdentifier = 2147483647;
        public const int AuthPacketFakeIdentifier = 2147483646;
        public const int SuccessAuthPacketFakeIdentifier = 2147483647;
        
        readonly List<IMqttPacketAwaitable> _waiters = new List<IMqttPacketAwaitable>();

        bool _isDisposed;

        public MqttPacketAwaitable<TResponsePacket> AddAwaitable<TResponsePacket>(int packetIdentifier) where TResponsePacket : MqttPacket
        {
            var awaitable = new MqttPacketAwaitable<TResponsePacket>(packetIdentifier, this);

            lock (_waiters)
            {
                _waiters.Add(awaitable);
            }

            return awaitable;
        }

        public void CancelAll()
        {
            lock (_waiters)
            {
                foreach (var awaitable in _waiters)
                {
                    awaitable.Cancel();
                }

                _waiters.Clear();
            }
        }

        public void Dispose()
        {
            Dispose(new ObjectDisposedException(nameof(MqttPacketDispatcher)));
        }

        public void Dispose(Exception exception)
        {
            if (exception == null)
            {
                throw new ArgumentNullException(nameof(exception));
            }

            lock (_waiters)
            {
                FailAll(exception);
                
                // Make sure that no task can start waiting after this instance is already disposed.
                // This will prevent unexpected freezes.
                _isDisposed = true;
            }
        }

        public void FailAll(Exception exception)
        {
            if (exception == null)
            {
                throw new ArgumentNullException(nameof(exception));
            }

            lock (_waiters)
            {
                foreach (var awaitable in _waiters)
                {
                    awaitable.Fail(exception);
                }

                _waiters.Clear();
            }
        }

        public void RemoveAwaitable(IMqttPacketAwaitable awaitable)
        {
            if (awaitable == null)
            {
                throw new ArgumentNullException(nameof(awaitable));
            }

            lock (_waiters)
            {
                _waiters.Remove(awaitable);
            }
        }

        public bool TryDispatch(MqttPacket packet)
        {
            if (packet == null)
            {
                throw new ArgumentNullException(nameof(packet));
            }

            var identifier = 0;
            if (packet is MqttPacketWithIdentifier packetWithIdentifier)
            {
                identifier = packetWithIdentifier.PacketIdentifier;
            }
            else if (packet is MqttPingRespPacket)
            {
                identifier = PingRespPacketFakeIdentifier;
            }
            else
            {
                throw new InvalidOperationException($"Cannot dispatch {packet.GetRfcName()} packet.");
            }
            
            var waiters = new List<IMqttPacketAwaitable>();
            lock (_waiters)
            {
                ThrowIfDisposed();
                
                for (var i = _waiters.Count - 1; i >= 0; i--)
                {
                    var entry = _waiters[i];
                    
                    if (entry.PacketIdentifier != identifier)
                    {
                        continue;
                    }

                    waiters.Add(entry);
                    _waiters.RemoveAt(i);
                }
            }

            foreach (var matchingEntry in waiters)
            {
                matchingEntry.Complete(packet);
            }

            return waiters.Count > 0;
        }

        void ThrowIfDisposed()
        {
            if (_isDisposed)
            {
                throw new ObjectDisposedException(nameof(MqttPacketDispatcher));
            }
        }
    }
}