Thread-Safe, Lock-Free, Append-Only, Copy-On-Write Dictionary

CopyOnWriteDictionary.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;

namespace Singulink.Collections
{
    /// <summary>
    /// Thread-safe append-only dictionary with internal copy-on-write behavior that enables the fastest possible lock-free lookups.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The performance characteristics of this dictionary make it perfect for use as a permanent object cache. Updates to the dictionary are temporarily
    /// cached in a synchronized mutable collection which allows copying to be delayed so that updates are batched together to reduce copying and improve write
    /// performance. This dictionary has identical lookup performance to <see cref="Dictionary{TKey, TValue}"/> for values that have been copied into the main
    /// lookup.
    /// </para>
    /// <para>
    /// Delayed updates are stored in a temporary synchronized mutable lookup. Once there have been no updates to the dictionary for the amount of time set on
    /// <see cref="CopyDelay"/> then the internal copy operation is performed on a background thread and lookups are restored to full speed.
    /// </para>
    /// </remarks>
    public class CopyOnWriteDictionary<TKey, TValue> : IDictionary<TKey, TValue> where TKey : notnull
    {
        private Dictionary<TKey, TValue> _lookup;
        private Dictionary<TKey, TValue>? _pendingUpdates;

        private int _copyDelay = 100;
        private Timer? _copyTimer;

        private readonly object _syncRoot = new object();

        /// <summary>
        /// Initializes a new copy-on-write dictionary that is empty and uses the specified comparer.
        /// </summary>
        public CopyOnWriteDictionary(IEqualityComparer<TKey>? comparer = null)
        {
            _lookup = new Dictionary<TKey, TValue>(comparer);
        }

        /// <summary>
        /// Initialized a new copy-on-write dictionary that contains elements copied from the specified collection and uses the specified comparer.
        /// </summary>
        public CopyOnWriteDictionary(IEnumerable<KeyValuePair<TKey, TValue>> keyValuePairs, IEqualityComparer<TKey>? comparer = null)
        {
            _lookup = new Dictionary<TKey, TValue>(keyValuePairs.KnownCount() ?? 0, comparer);

            foreach (var kvp in keyValuePairs)
                _lookup.Add(kvp.Key, kvp.Value);
        }

        /// <summary>
        /// Gets or sets the value associated with the specified key.
        /// </summary>
        public TValue this[TKey key] {
            get {
                if (TryGetValue(key, out var value))
                    return value;

                throw new KeyNotFoundException();
            }
            set {
                lock(_syncRoot) {
                    if (_lookup.ContainsKey(key))
                        _lookup[key] = value;
                    else if (_pendingUpdates != null)
                        _pendingUpdates[key] = value;
                    else
                        AddInternal(key, value, true, true);
                }
            }
        }

        /// <summary>
        /// Gets the keys contained in this dictionary.
        /// </summary>
        public ICollection<TKey> Keys {
            get {
                lock (_syncRoot) {
                    if (_pendingUpdates != null)
                        return new MergedCollection<TKey>(_lookup.Keys, _pendingUpdates.Keys.ToList());

                    return _lookup.Keys;
                }
            }
        }

        /// <summary>
        /// Gets the values contained in this dictionary.
        /// </summary>
        public ICollection<TValue> Values {
            get {
                lock (_syncRoot) {
                    if (_pendingUpdates != null)
                        return new MergedCollection<TValue>(_lookup.Values, _pendingUpdates.Values.ToList());

                    return _lookup.Values;
                }
            }
        }

        /// <summary>
        /// Gets or sets the number of milliseconds to delay copying updates to the internal dictionary. Default is 100ms and a value of 0 disables the copy
        /// delay feature, causing any updates to the dictionary to immediately trigger a copy operation.
        /// </summary>
        public int CopyDelay {
            get => Volatile.Read(ref _copyDelay);
            set {
                if (value < 0)
                    throw new ArgumentOutOfRangeException(nameof(value));

                lock (_syncRoot) {
                    DebugCheckState();
                    _copyDelay = value;

                    if (_pendingUpdates != null) {
                        if (value == 0)
                            CopyUpdates();
                        else
                            _copyTimer?.Change(value, Timeout.Infinite);
                    }

                    DebugCheckState();
                }
            }
        }

        /// <summary>
        /// Gets the number of items in this dictionary.
        /// </summary>
        public int Count {
            get {
                DebugNoSync();

                lock (_syncRoot)
                    return CountInternal;
            }
        }

        /// <summary>
        /// Gets a value indicating whether a copy operation is pending due to delayed updates.
        /// </summary>
        public bool IsCopyPending {
            get {
                DebugNoSync();

                lock (_syncRoot)
                    return _pendingUpdates != null;
            }
        }

        private int CountInternal {
            get {
                DebugSyncRequired();
                int count = _lookup.Count;

                if (_pendingUpdates != null)
                    count += _pendingUpdates.Count;

                return count;
            }
        }

        /// <summary>
        /// Performs a lookup for a value with the specified key. This method has a fast lock-free path with identical performance to <see
        /// cref="Dictionary{TKey, TValue}"/> if the key has been copied to the main dictionary and is not waiting in delayed pending updates.
        /// </summary>
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool TryGetValue(TKey key, [MaybeNullWhen(false)] out TValue value)
        {
            if (_lookup.TryGetValue(key, out value))
                return true;

            return TryGetValueSlow(key, out value);

            //return _lookup.TryGetValue(key, out value);
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private bool TryGetValueSlow(TKey key, out TValue value)
        {
            lock (_syncRoot) {
                // Check pending updates first since we just checked the main lookup and it wasn't there.
                if (_pendingUpdates != null && _pendingUpdates.TryGetValue(key, out value))
                    return true;

                return _lookup.TryGetValue(key, out value);
            }
        }

        /// <summary>
        /// Returns whether this dictionary contains a particular key.
        /// </summary>
        public bool ContainsKey(TKey key) => TryGetValue(key, out _);

        /// <summary>
        /// Adds the key and value to a copy of the internal dictionary under a synchronized lock. This method is mostly intended for initialization or when
        /// there is no possibility of another thread trying to add the same value.
        /// </summary>
        /// <remarks>
        /// If adding items after initialization then <see cref="AddOrGetValue(TKey, Func{TValue}, bool)"/> is the preferred method to use since the value
        /// factory only runs if the key doesn't exist and there is no risk of competing threads trying to add the same item.
        /// </remarks>
        public void Add(TKey key, TValue value, bool delayCopy = true)
        {
            lock (_syncRoot) {
                AddInternal(key, value, delayCopy, false);
            }
        }

        /// <summary>
        /// Adds the key and value to a copy of the internal dictionary under a synchronized lock. This method is mostly intended for initialization or when
        /// there is no possibility of another thread trying to add the same value.
        /// </summary>
        /// <remarks>
        /// If adding items after initialization then <see cref="AddOrGetValue(TKey, Func{TValue}, bool)"/> is the preferred method to use since the value
        /// factory only runs if the key doesn't exist and there is no risk of competing threads trying to add the same item.
        /// </remarks>
        public bool TryAdd(TKey key, TValue value, bool delayCopy = true)
        {
            lock (_syncRoot) {
                return TryAddInternal(key, value, delayCopy, false);
            }
        }

        /// <summary>
        /// Performs a lookup for a value with the specified key under a synchronized lock. If the key is not found then the factory is invoked and the value
        /// is added to the dictionary. This method should only be called after trying <see cref="TryGetValue(TKey, out TValue)"/> first to avoiding taking
        /// a lock and creating a factory delegate in the event that the key already exists.
        /// </summary>
        public TValue GetOrAdd(TKey key, TValue value, bool delayCopy = true)
        {
            lock (_syncRoot) {
                if (TryGetValue(key, out var existingValue))
                    return existingValue;

                AddInternal(key, value, delayCopy, true);
                return value;
            }
        }

        /// <summary>
        /// Returns an enumerator that iterates though the dictionary.
        /// </summary>
        public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
        {
            IEnumerable<KeyValuePair<TKey, TValue>> items;

            lock (_syncRoot) {
                items = _lookup;

                if (_pendingUpdates != null)
                    items = items.Concat(_pendingUpdates.ToList());
            }

            foreach (var item in items)
                yield return item;
        }

        private bool TryAddInternal(TKey key, TValue value, bool delayCopy, bool lookupPrechecked)
        {
            #if DEBUG
            DebugSyncRequired();
            DebugCheckState();
            int preAddCount = CountInternal;
            #endif

            if (!lookupPrechecked && _lookup.ContainsKey(key))
                return false;

            ////////////

            if (_copyDelay == 0)
                delayCopy = false;

            if (_pendingUpdates == null) {
                if (!delayCopy) {
                    var newLookup = new Dictionary<TKey, TValue>(_lookup.Count + 1, _lookup.Comparer);

                    foreach (var item in _lookup)
                        newLookup.Add(item.Key, item.Value);

                    newLookup.Add(key, value);

                    _lookup = newLookup;
                    return true;
                }

                _pendingUpdates = new Dictionary<TKey, TValue>(31, _lookup.Comparer);
            }

            if (!_pendingUpdates.TryAdd(key, value))
                return false;

            if (delayCopy) {
                if (_copyTimer == null)
                    _copyTimer = new Timer(CopyUpdatesCallback, null, _copyDelay, Timeout.Infinite);
                else
                    _copyTimer.Change(_copyDelay, Timeout.Infinite);
            }
            else {
                CopyUpdates();
            }

            #if DEBUG
            DebugCheckState();
            Debug.Assert(CountInternal == preAddCount + 1);
            #endif

            return true;
        }

        private void AddInternal(TKey key, TValue value, bool delayCopy, bool lookupPrechecked)
        {
            if (!TryAddInternal(key, value, delayCopy, lookupPrechecked))
                throw new ArgumentException("An element with the same key already exists.", nameof(key));
        }

        private void CopyUpdates()
        {
            #if DEBUG
            DebugSyncRequired();
            DebugCheckState();
            int preCopyCount = CountInternal;
            #endif

            if (_copyTimer != null) {
                _copyTimer.Dispose();
                _copyTimer = null;
            }

            if (_pendingUpdates != null) {
                var newLookup = new Dictionary<TKey, TValue>(_lookup.Count + _pendingUpdates.Count, _lookup.Comparer);

                foreach (var item in _lookup)
                    newLookup.Add(item.Key, item.Value);

                foreach (var item in _pendingUpdates)
                    newLookup.Add(item.Key, item.Value);

                _lookup = newLookup;
                _pendingUpdates = null;
            }

            #if DEBUG
            DebugCheckState();
            Debug.Assert(CountInternal == preCopyCount);
            #endif
        }

        private void CopyUpdatesCallback(object? state)
        {
            lock (_syncRoot)
                CopyUpdates();
        }

        #region Debug

        [Conditional("DEBUG")]
        private void DebugCheckState()
        {
            Debug.Assert(Monitor.IsEntered(_syncRoot), "synchronization required to check state");
            Debug.Assert((_copyTimer == null) == (_pendingUpdates == null), "invalid state");
        }

        [Conditional("DEBUG")]
        private void DebugSyncRequired() => Debug.Assert(Monitor.IsEntered(_syncRoot), "synchronization required");

        [Conditional("DEBUG")]
        private void DebugNoSync() => Debug.Assert(!Monitor.IsEntered(_syncRoot), "nested synchronization - use unsynchronized method instead");

        #endregion

        #region Explicit Interface Members

        bool ICollection<KeyValuePair<TKey, TValue>>.IsReadOnly => false;

        TValue IDictionary<TKey, TValue>.this[TKey key] {
            get => this[key];
            set => throw new NotSupportedException();
        }

        void IDictionary<TKey, TValue>.Add(TKey key, TValue value) => Add(key, value);

        void ICollection<KeyValuePair<TKey, TValue>>.Add(KeyValuePair<TKey, TValue> item) => Add(item.Key, item.Value);

        bool ICollection<KeyValuePair<TKey, TValue>>.Contains(KeyValuePair<TKey, TValue> item)
        {
            if (TryGetValue(item.Key, out var value))
                return EqualityComparer<TValue>.Default.Equals(item.Value, value);

            return false;
        }

        void ICollection<KeyValuePair<TKey, TValue>>.CopyTo(KeyValuePair<TKey, TValue>[] array, int arrayIndex)
        {
            ICollection<KeyValuePair<TKey, TValue>> lookup;

            lock (_syncRoot) {
                if (_pendingUpdates != null)
                    lookup = new MergedCollection<KeyValuePair<TKey, TValue>>(_lookup, _pendingUpdates.ToList());
                else
                    lookup = _lookup;
            }

            lookup.CopyTo(array, arrayIndex);
        }

        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

        #endregion

        #region Not Supported

        bool IDictionary<TKey, TValue>.Remove(TKey key) => throw new NotSupportedException();

        bool ICollection<KeyValuePair<TKey, TValue>>.Remove(KeyValuePair<TKey, TValue> item) => throw new NotSupportedException();

        void ICollection<KeyValuePair<TKey, TValue>>.Clear() => throw new NotSupportedException();

        #endregion
    }
    
    /// <summary>
    /// Provides a wrapper around two collections to be presented as one combined collection.
    /// </summary>
    public sealed class MergedCollection<T> : ICollection<T>
    {
        private readonly ICollection<T> _first;
        private readonly ICollection<T> _second;

        /// <summary>
        /// Create a new instance of <see cref="MergedCollection{T}"/> using the provided collections.
        /// </summary>
        /// <param name="first">The first collection to wrap.</param>
        /// <param name="second">The second collection to wrap.</param>
        public MergedCollection(ICollection<T> first, ICollection<T> second)
        {
            _first = first ?? throw new ArgumentNullException(nameof(first));
            _second = second ?? throw new ArgumentNullException(nameof(second));
        }

        /// <summary>
        /// Gets the number of elements contained in this collection.
        /// </summary>
        public int Count => _first.Count + _second.Count;

        /// <summary>
        /// Gets a value indicating whether this collection is read-only. Always returns true.
        /// </summary>
        public bool IsReadOnly => true;

        /// <summary>
        /// Determines whether the collection contains the specified value.
        /// </summary>
        public bool Contains(T item) => _first.Contains(item) || _second.Contains(item);

        /// <inheritdoc cref="IEnumerable{T}.GetEnumerator"/>
        public IEnumerator<T> GetEnumerator()
        {
            foreach (var item in _first)
                yield return item;

            foreach (var item in _second)
                yield return item;
        }

        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

        void ICollection<T>.CopyTo(T[] array, int arrayIndex)
        {
            _first.CopyTo(array, arrayIndex);
            _second.CopyTo(array, arrayIndex + _first.Count);
        }

        void ICollection<T>.Add(T item) => throw new NotSupportedException();

        void ICollection<T>.Clear() => throw new NotSupportedException();

        bool ICollection<T>.Remove(T item) => throw new NotSupportedException();
    }

    public static class EnumerableExtensions
    {
        public static int? KnownCount<T>(this IEnumerable<T> source)
        {
            if (source is ICollection<T> genericCollection)
                return genericCollection.Count;

            return null;
        }
    }
}