CleanableWeakComHandleTable.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.

#nullable disable

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
using Microsoft.CodeAnalysis.Shared.TestHooks;
using Microsoft.VisualStudio.LanguageServices.Implementation.Utilities;

namespace Microsoft.VisualStudio.LanguageServices.Implementation.Interop
{
    /// <summary>
    /// Special collection for storing a table of COM objects weakly that provides
    /// logic for cleaning up dead references in a time-sliced way. Public members of this
    /// collection are affinitized to the foreground thread.
    /// </summary>
    internal class CleanableWeakComHandleTable<TKey, TValue> : ForegroundThreadAffinitizedObject
        where TValue : class
    {
        private const int DefaultCleanUpThreshold = 25;
        private static readonly TimeSpan s_defaultCleanUpTimeSlice = TimeSpan.FromMilliseconds(15);

        private readonly Dictionary<TKey, WeakComHandle<TValue, TValue>> _table;
        private readonly HashSet<TKey> _deadKeySet;

        /// <summary>
        /// The upper limit of items that the collection will store before clean up is recommended.
        /// </summary>
        public int CleanUpThreshold { get; }

        /// <summary>
        /// The amount of time that can pass during clean up it returns.
        /// </summary>
        public TimeSpan CleanUpTimeSlice { get; }

        private int _itemsAddedSinceLastCleanUp;
        private bool _needsCleanUp;

        public bool NeedsCleanUp => _needsCleanUp;

        public CleanableWeakComHandleTable(IThreadingContext threadingContext, int? cleanUpThreshold = null, TimeSpan? cleanUpTimeSlice = null)
            : base(threadingContext)
        {
            _table = new Dictionary<TKey, WeakComHandle<TValue, TValue>>();
            _deadKeySet = new HashSet<TKey>();

            CleanUpThreshold = cleanUpThreshold ?? DefaultCleanUpThreshold;
            CleanUpTimeSlice = cleanUpTimeSlice ?? s_defaultCleanUpTimeSlice;
        }

        /// <summary>
        /// Cleans up references to dead objects in the table. This operation will yield to other foreground operations
        /// any time execution exceeds <see cref="CleanUpTimeSlice"/>.
        /// </summary>
        public async Task CleanUpDeadObjectsAsync(IAsynchronousOperationListener listener)
        {
            using var _ = listener.BeginAsyncOperation(nameof(CleanUpDeadObjectsAsync));

            Debug.Assert(ThreadingContext.JoinableTaskContext.IsOnMainThread, "This method is optimized for cases where calls do not yield before checking _needsCleanUp.");

            await ThreadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(ThreadingContext.DisposalToken);

            if (!_needsCleanUp)
            {
                return;
            }

            // Immediately mark as not needing cleanup; this operation will clean up the table by the time it returns.
            _needsCleanUp = false;

            var timeSlice = new TimeSlice(CleanUpTimeSlice);

            await CollectDeadKeysAsync().ConfigureAwait(true);
            await RemoveDeadKeysAsync().ConfigureAwait(true);
            return;

            // Local functions
            async Task CollectDeadKeysAsync()
            {
                // This method returns after making a complete pass enumerating the elements of _table without finding
                // any entries that are not alive. If a pass exceeds the allowed time slice after finding one or more
                // dead entries, the pass yields before processing the elements found so far and restarting the
                // enumeration.
                //
                // ⚠ This method may interleave with other asynchronous calls to CleanUpDeadObjectsAsync.
                var cleanUpEnumerator = _table.GetEnumerator();
                while (cleanUpEnumerator.MoveNext())
                {
                    var pair = cleanUpEnumerator.Current;
                    if (!pair.Value.IsAlive())
                    {
                        _deadKeySet.Add(pair.Key);

                        if (timeSlice.IsOver)
                        {
                            // Yield before processing items found so far.
                            await ResetTimeSliceAsync().ConfigureAwait(true);

                            // Process items found prior to exceeding the time slice. Due to interleaving, it is
                            // possible for this call to process items found by another asynchronous call to
                            // CollectDeadKeysAsync, or for another asynchronous call to RemoveDeadKeysAsync to process
                            // all items prior to this call.
                            await RemoveDeadKeysAsync().ConfigureAwait(true);

                            // Obtain a new enumerator since the previous one may be invalidated.
                            cleanUpEnumerator = _table.GetEnumerator();
                        }
                    }
                }
            }

            async Task RemoveDeadKeysAsync()
            {
                while (_deadKeySet.Count > 0)
                {
                    // Fully process one item from _deadKeySet before the possibility of yielding
                    var key = _deadKeySet.First();

                    _deadKeySet.Remove(key);

                    Debug.Assert(_table.ContainsKey(key), "Key not found in table.");
                    _table.Remove(key);

                    if (timeSlice.IsOver)
                    {
                        await ResetTimeSliceAsync().ConfigureAwait(true);
                    }
                }
            }

            async Task ResetTimeSliceAsync()
            {
                await listener.Delay(DelayTimeSpan.NearImmediate, ThreadingContext.DisposalToken).ConfigureAwait(true);
                timeSlice = new TimeSlice(CleanUpTimeSlice);
            }
        }

        public void Add(TKey key, TValue value)
        {
            this.AssertIsForeground();

            if (value == null)
            {
                throw new ArgumentNullException(nameof(value));
            }

            if (_table.ContainsKey(key))
            {
                throw new InvalidOperationException($"Key already exists in table: {(key != null ? key.ToString() : "<null>")}.");
            }

            _itemsAddedSinceLastCleanUp++;
            if (_itemsAddedSinceLastCleanUp >= CleanUpThreshold)
            {
                _needsCleanUp = true;
                _itemsAddedSinceLastCleanUp = 0;
            }

            _table.Add(key, new WeakComHandle<TValue, TValue>(value));
        }

        public TValue Remove(TKey key)
        {
            this.AssertIsForeground();

            if (_deadKeySet.Contains(key))
            {
                _deadKeySet.Remove(key);
            }

            if (_table.TryGetValue(key, out var handle))
            {
                _table.Remove(key);
                return handle.ComAggregateObject;
            }

            return null;
        }

        public bool ContainsKey(TKey key)
        {
            this.AssertIsForeground();

            return _table.ContainsKey(key);
        }

        public bool TryGetValue(TKey key, out TValue value)
        {
            this.AssertIsForeground();
            if (_table.TryGetValue(key, out var handle))
            {
                value = handle.ComAggregateObject;
                return value != null;
            }

            value = null;
            return false;
        }

        public IEnumerable<TValue> Values
        {
            get
            {
                foreach (var keyValuePair in _table)
                {
                    yield return keyValuePair.Value.ComAggregateObject;
                }
            }
        }
    }
}