aspnetcore/src/Microsoft.AspNetCore.SignalR.Core/SerializedHubMessage.cs

// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

using System;
using System.Collections.Generic;
using Microsoft.AspNetCore.SignalR.Protocol;

namespace Microsoft.AspNetCore.SignalR
{
    /// <summary>
    /// Represents a serialization cache for a single message.
    /// </summary>
    public class SerializedHubMessage
    {
        private SerializedMessage _cachedItem1;
        private SerializedMessage _cachedItem2;
        private IList<SerializedMessage> _cachedItems;
        private readonly object _lock = new object();

        /// <summary>
        /// Gets the hub message for the serialization cache.
        /// </summary>
        public HubMessage Message { get; }

        /// <summary>
        /// Initializes a new instance of the <see cref="SerializedHubMessage"/> class.
        /// </summary>
        /// <param name="messages">A collection of already serialized messages to cache.</param>
        public SerializedHubMessage(IReadOnlyList<SerializedMessage> messages)
        {
            // A lock isn't needed here because nobody has access to this type until the constructor finishes.
            for (var i = 0; i < messages.Count; i++)
            {
                var message = messages[i];
                SetCache(message.ProtocolName, message.Serialized);
            }
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="SerializedHubMessage"/> class.
        /// </summary>
        /// <param name="message">The hub message for the cache. This will be serialized with an <see cref="IHubProtocol"/> in <see cref="GetSerializedMessage"/> to get the message's serialized representation.</param>
        public SerializedHubMessage(HubMessage message)
        {
            Message = message;
        }

        /// <summary>
        /// Gets the serialized representation of the <see cref="HubMessage"/> using the specified <see cref="IHubProtocol"/>.
        /// </summary>
        /// <param name="protocol">The protocol used to create the serialized representation.</param>
        /// <returns>The serialized representation of the <see cref="HubMessage"/>.</returns>
        public ReadOnlyMemory<byte> GetSerializedMessage(IHubProtocol protocol)
        {
            lock (_lock)
            {
                if (!TryGetCached(protocol.Name, out var serialized))
                {
                    if (Message == null)
                    {
                        throw new InvalidOperationException(
                            "This message was received from another server that did not have the requested protocol available.");
                    }

                    serialized = protocol.GetMessageBytes(Message);
                    SetCache(protocol.Name, serialized);
                }

                return serialized;
            }
        }

        // Used for unit testing.
        internal IReadOnlyList<SerializedMessage> GetAllSerializations()
        {
            // Even if this is only used in tests, let's do it right.
            lock (_lock)
            {
                if (_cachedItem1.ProtocolName == null)
                {
                    return Array.Empty<SerializedMessage>();
                }

                var list = new List<SerializedMessage>(2);
                list.Add(_cachedItem1);

                if (_cachedItem2.ProtocolName != null)
                {
                    list.Add(_cachedItem2);

                    if (_cachedItems != null)
                    {
                        list.AddRange(_cachedItems);
                    }
                }

                return list;
            }
        }

        private void SetCache(string protocolName, ReadOnlyMemory<byte> serialized)
        {
            // We set the fields before moving on to the list, if we need it to hold more than 2 items.
            // We have to read/write these fields under the lock because the structs might tear and another
            // thread might observe them half-assigned

            if (_cachedItem1.ProtocolName == null)
            {
                _cachedItem1 = new SerializedMessage(protocolName, serialized);
            }
            else if (_cachedItem2.ProtocolName == null)
            {
                _cachedItem2 = new SerializedMessage(protocolName, serialized);
            }
            else
            {
                if (_cachedItems == null)
                {
                    _cachedItems = new List<SerializedMessage>();
                }

                foreach (var item in _cachedItems)
                {
                    if (string.Equals(item.ProtocolName, protocolName, StringComparison.Ordinal))
                    {
                        // No need to add
                        return;
                    }
                }

                _cachedItems.Add(new SerializedMessage(protocolName, serialized));
            }
        }

        private bool TryGetCached(string protocolName, out ReadOnlyMemory<byte> result)
        {
            if (string.Equals(_cachedItem1.ProtocolName, protocolName, StringComparison.Ordinal))
            {
                result = _cachedItem1.Serialized;
                return true;
            }

            if (string.Equals(_cachedItem2.ProtocolName, protocolName, StringComparison.Ordinal))
            {
                result = _cachedItem2.Serialized;
                return true;
            }

            if (_cachedItems != null)
            {
                foreach (var serializedMessage in _cachedItems)
                {
                    if (string.Equals(serializedMessage.ProtocolName, protocolName, StringComparison.Ordinal))
                    {
                        result = serializedMessage.Serialized;
                        return true;
                    }
                }
            }

            result = default;
            return false;
        }
    }
}