WeakMapExtensions

WeakMapExtensions.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

namespace WeakMap
{
    /// <summary>
    /// Allows to link additional properties to any objects by weakly linking a dictionary for any object.
    /// Weakly means the dictionary exists in memory only as long as object it linked to is still stored in memory
    /// Works via javascript's WeakMap analog for C# (System.Runtime.CompilerServices.ConditionalWeakTable)
    /// </summary>
    public static class WeakMapExtensions
    {
        static ConditionalWeakTable<object, ConcurrentDictionary<object, Lazy<object>>> globalMap
            = new ConditionalWeakTable<object, ConcurrentDictionary<object, Lazy<object>>>();

        static object defaultKey = new object();

        /// <summary>
        /// Gets or creates weakly referenced object by a key, that will be removed from memory when source is removed
        /// Can be used for extending object with dynamic properties attached to that object and removed with the object
        /// </summary>
        /// <typeparam name="TObject"></typeparam>
        /// <typeparam name="TData"></typeparam>
        /// <param name="source"></param>
        /// <param name="key"></param>
        /// <param name="createWeakData"></param>
        /// <returns></returns>
        public static TData GetOrAddWeakData<TObject, TData>(this TObject source, object key, Func<object, TData> createWeakData)
            where TObject : class
        {
            var objDictionary = GetOrCreateWeakSourceDictionary(source);
            var res = objDictionary.GetOrAdd(key, valueFactory:k => new Lazy<object>(()=>createWeakData(k))).Value;
            return (TData)res;
        }

        public static TData GetOrAddWeakData<TObject, TData>(this TObject source, Func<object, TData> createWeakData)
            where TObject : class
        {
            return GetOrAddWeakData(source, defaultKey, createWeakData);
        }

        public static bool TryGetWeakData<TData>(this object source, object key, out TData result)
        {
            result = default(TData);
            var wdict = GetOrCreateWeakSourceDictionary(source);
            if (wdict.TryGetValue(key, out var lz)) { result = (TData)lz.Value; return true; }
            return false;
        }

        public static TData GetWeakData<TData>(this object source, object key, Func<object, TData> getDefaultByKey)
        {
            if (source.TryGetWeakData<TData>(key, out var res)) return res;
            return getDefaultByKey(key);
        }
        public static TData GetWeakData<TData>(this object source, object key, TData defaultValue = default(TData))
        {
            if (source.TryGetWeakData<TData>(key, out var res)) return res;
            return defaultValue;
        }

        static ConcurrentDictionary<object, Lazy<object>> GetOrCreateWeakSourceDictionary<TObject>(TObject source)
        {
            bool corrupted = true;
            int try_attempts = 2;
            Exception lastExc = null;
            while (try_attempts-- > 0)
                try
                {
                    var objDictionary = globalMap.GetOrCreateValue(source);
                    corrupted = false;
                    return objDictionary;
                }
                catch (Exception exception)
                {
                    Serilog.Log.Error(exception, "Error happened during access to WeakTable. Attempts:{TryAttemptsLeft}", try_attempts);
                    //TODO: log it here somehow
                    lastExc = exception;
                    if (!corrupted) throw;
                    GC.Collect();
                    GC.WaitForPendingFinalizers();
                }
            if (corrupted)
            /*
            OutOfMemory exception corrupts the state of globalMap and make all future calls unusable
              to fix it we have to recreate globalMap instance.

            Otherwise it will keep throwing this exception everytime you try to access it:

             System.InvalidOperationException: A prior operation on this collection was interrupted by an exception. Collection's state is no longer trusted.
at System.Runtime.CompilerServices.ConditionalWeakTable`2.Container.VerifyIntegrity()
at System.Runtime.CompilerServices.ConditionalWeakTable`2.Container.CreateEntryNoResize(TKey key, TValue value)
at System.Runtime.CompilerServices.ConditionalWeakTable`2.GetValueLocked(TKey key, CreateValueCallback createValueCallback)
at System.Runtime.CompilerServices.ConditionalWeakTable`2.GetValue(TKey key, CreateValueCallback createValueCallback)

             */
            {
                globalMap = new ConditionalWeakTable<object, ConcurrentDictionary<object, Lazy<object>>>();
                Serilog.Log.Error(lastExc, "Weak table has to be recreated to avoid corruption");
            }

            throw lastExc;
        }

        public static TData SetOrUpdateWeakData<TObject, TData>(this TObject source, object key, Func<object, TData, TData> updateWeakData, TData defaultData = default(TData))
            where TObject : class
        {
            var objDictionary = GetOrCreateWeakSourceDictionary(source);
            return (TData)objDictionary.AddOrUpdate(
                key,
                addValueFactory: k=> new Lazy<object>(()=>updateWeakData(k, defaultData)),
                updateValueFactory: (k,old) => new Lazy<object>(()=>updateWeakData(k, (TData)old.Value))
            ).Value;
        }
        public static TData SetOrUpdateWeakData<TObject, TData>(this TObject source, Func<object, TData, TData> updateWeakData, TData defaultData = default(TData))
            where TObject : class
        {
            return SetOrUpdateWeakData<TObject, TData>(source, defaultKey, updateWeakData, defaultData);
        }

        public static void SetWeakData<TObject, TData>(this TObject source, object key, TData data)
            where TObject : class
        {
            var objDictionary = GetOrCreateWeakSourceDictionary(source);
            objDictionary[key] = new Lazy<object>(()=>data, true);
        }


        public static bool IsAlive<T>(this WeakReference<T> weak)
            where T : class
        {
            return weak.TryGetTarget(out var target);
        }

    }
}