diff --git a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs index 845125f617..e181e9b407 100644 --- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs +++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs @@ -3,7 +3,6 @@ // See the LICENSE file in the project root for more information. using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Runtime.CompilerServices; using System.Text.Json; @@ -46,11 +45,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte _registry = Registries.GetValue(Connection, static _ => new PageRegistry()); _registry.EnsureSubscribed(Connection); - var pages = _registry.Pages.GetOrAdd(_pageId, static _ => new List>()); - lock (pages) - { - pages.Add(new WeakReference(this)); - } + _registry.Pages.Register(_pageId, this); } public event TypedEventHandler? ItemsChanged; @@ -99,17 +94,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte { _details.Dispose(); _commands.Dispose(); - if (_registry.Pages.TryGetValue(_pageId, out var pages)) - { - lock (pages) - { - pages.RemoveAll(weak => !weak.TryGetTarget(out var target) || ReferenceEquals(target, this)); - if (pages.Count == 0) - { - _registry.Pages.TryRemove(_pageId, out _); - } - } - } + _registry.Pages.Unregister(_pageId, this); base.Dispose(); } @@ -163,30 +148,12 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte } var pageId = pageProperty.GetString(); - if (pageId is null || !registry.Pages.TryGetValue(pageId, out var pageReferences)) + if (pageId is null) { return; } - List targets = []; - lock (pageReferences) - { - pageReferences.RemoveAll(weak => !weak.TryGetTarget(out _)); - foreach (var weak in pageReferences) - { - if (weak.TryGetTarget(out var target)) - { - targets.Add(target); - } - } - - if (pageReferences.Count == 0) - { - registry.Pages.TryRemove(pageId, out _); - } - } - - foreach (var target in targets) + foreach (var target in registry.Pages.GetLiveTargets(pageId)) { var handler = target.ItemsChanged; if (handler is not null) @@ -206,7 +173,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte private readonly object _subscribeLock = new(); private bool _subscribed; - public ConcurrentDictionary>> Pages { get; } = new(); + public JSWeakReferenceRegistry Pages { get; } = new(); public void EnsureSubscribed(JsonRpcConnection connection) { diff --git a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs index b5c6d983fe..79407d44f5 100644 --- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs +++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs @@ -63,7 +63,7 @@ internal sealed partial class JSLazyCache : IDisposable return; } - DisposeCreatedValue(); + // The host may still own the previous proxy, so invalidation only releases the cache reference. _value = default!; _hasValue = false; } diff --git a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs index 2a42d132a3..e635069c87 100644 --- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs +++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs @@ -3,7 +3,6 @@ // See the LICENSE file in the project root for more information. using System; -using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Runtime.CompilerServices; @@ -56,11 +55,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage _registry = Registries.GetValue(Connection, static _ => new PageRegistry()); _registry.EnsureSubscribed(Connection); - var list = _registry.Pages.GetOrAdd(_pageId, static _ => new List>()); - lock (list) - { - list.Add(new WeakReference(this)); - } + _registry.Pages.Register(_pageId, this); } public event TypedEventHandler? ItemsChanged; @@ -295,17 +290,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage _emptyContent.Dispose(); base.Dispose(); - if (_registry.Pages.TryGetValue(_pageId, out var list)) - { - lock (list) - { - list.RemoveAll(weak => !weak.TryGetTarget(out var target) || ReferenceEquals(target, this)); - if (list.Count == 0) - { - _registry.Pages.TryRemove(_pageId, out _); - } - } - } + _registry.Pages.Unregister(_pageId, this); } private static void DispatchItemsChanged(PageRegistry registry, JsonElement paramsElement) @@ -319,7 +304,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage } var pageId = pageProp.GetString(); - if (pageId == null || !registry.Pages.TryGetValue(pageId, out var proxyRefs)) + if (pageId == null) { return; } @@ -331,27 +316,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage totalItems = totalItemsProp.GetInt32(); } - // Snapshot live proxies and prune collected ones so the registry does - // not grow as pages come and go. - List targets = new(); - lock (proxyRefs) - { - proxyRefs.RemoveAll(weak => !weak.TryGetTarget(out _)); - foreach (var weak in proxyRefs) - { - if (weak.TryGetTarget(out var proxy)) - { - targets.Add(proxy); - } - } - - if (proxyRefs.Count == 0) - { - registry.Pages.TryRemove(pageId, out _); - } - } - - foreach (var proxy in targets) + foreach (var proxy in registry.Pages.GetLiveTargets(pageId)) { proxy.UpdatePageState(paramsElement); @@ -458,7 +423,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage private readonly object _subscribeLock = new(); private bool _subscribed; - public ConcurrentDictionary>> Pages { get; } = new(); + public JSWeakReferenceRegistry Pages { get; } = new(); // Binds the itemsChanged handler to the retained registry once. Binding here, // instead of inside the ConditionalWeakTable factory, keeps the handler from diff --git a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs index ab78985aca..755316df40 100644 --- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs +++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs @@ -3,8 +3,6 @@ // See the LICENSE file in the project root for more information. using System; -using System.Collections.Concurrent; -using System.Collections.Generic; using System.Runtime.CompilerServices; using System.Text.Json; using Microsoft.CmdPal.JsonRpc; @@ -18,46 +16,17 @@ internal static class JSPropertyChangeRegistry internal static void Register(JsonRpcConnection connection, string commandId, IJSPropertyChangeTarget target) { var registry = Registries.GetValue(connection, static _ => new Registry()); - while (true) - { - var targets = registry.Targets.GetOrAdd(commandId, static _ => []); - lock (targets) - { - if (!registry.Targets.TryGetValue(commandId, out var currentTargets) || - !ReferenceEquals(targets, currentTargets)) - { - continue; - } - - targets.RemoveAll(reference => !reference.TryGetTarget(out _)); - if (!targets.Exists(reference => - reference.TryGetTarget(out var current) && ReferenceEquals(current, target))) - { - targets.Add(new WeakReference(target)); - } - - return; - } - } + registry.Targets.Register(commandId, target); } internal static void Unregister(JsonRpcConnection connection, string commandId, IJSPropertyChangeTarget target) { - if (!Registries.TryGetValue(connection, out var registry) || - !registry.Targets.TryGetValue(commandId, out var targets)) + if (!Registries.TryGetValue(connection, out var registry)) { return; } - lock (targets) - { - targets.RemoveAll(reference => - !reference.TryGetTarget(out var current) || ReferenceEquals(current, target)); - if (targets.Count == 0) - { - RemoveTargets(registry, commandId, targets); - } - } + registry.Targets.Unregister(commandId, target); } internal static void Dispatch(JsonRpcConnection connection, JsonElement paramsElement) @@ -73,30 +42,12 @@ internal static class JSPropertyChangeRegistry } var commandId = commandIdProperty.GetString(); - if (commandId is null || !registry.Targets.TryGetValue(commandId, out var targets)) + if (commandId is null) { return; } - List liveTargets = []; - lock (targets) - { - targets.RemoveAll(reference => !reference.TryGetTarget(out _)); - foreach (var reference in targets) - { - if (reference.TryGetTarget(out var target)) - { - liveTargets.Add(target); - } - } - - if (targets.Count == 0) - { - RemoveTargets(registry, commandId, targets); - } - } - - foreach (var target in liveTargets) + foreach (var target in registry.Targets.GetLiveTargets(commandId)) { target.ApplyPropertyChanges(properties); } @@ -104,29 +55,13 @@ internal static class JSPropertyChangeRegistry internal static int GetRegistrationCount(JsonRpcConnection connection, string commandId) { - if (!Registries.TryGetValue(connection, out var registry) || - !registry.Targets.TryGetValue(commandId, out var targets)) - { - return 0; - } - - lock (targets) - { - return targets.Count; - } - } - - private static void RemoveTargets( - Registry registry, - string commandId, - List> targets) - { - ((ICollection>>>)registry.Targets) - .Remove(new KeyValuePair>>(commandId, targets)); + return Registries.TryGetValue(connection, out var registry) + ? registry.Targets.GetRegistrationCount(commandId) + : 0; } private sealed class Registry { - internal ConcurrentDictionary>> Targets { get; } = new(); + internal JSWeakReferenceRegistry Targets { get; } = new(); } } diff --git a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSWeakReferenceRegistry.cs b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSWeakReferenceRegistry.cs new file mode 100644 index 0000000000..573fd26475 --- /dev/null +++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSWeakReferenceRegistry.cs @@ -0,0 +1,122 @@ +// Copyright (c) Microsoft Corporation +// The Microsoft Corporation 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.Concurrent; +using System.Collections.Generic; + +namespace Microsoft.CmdPal.JsonRpc.Models; + +internal sealed partial class JSWeakReferenceRegistry + where TKey : notnull + where TTarget : class +{ + private readonly ConcurrentDictionary>> _targets = new(); + + internal void Register(TKey key, TTarget target) + { + while (true) + { + var targets = _targets.GetOrAdd(key, static _ => []); + lock (targets) + { + if (!IsCurrent(key, targets)) + { + continue; + } + + targets.RemoveAll(reference => !reference.TryGetTarget(out _)); + if (!targets.Exists(reference => + reference.TryGetTarget(out var current) && ReferenceEquals(current, target))) + { + targets.Add(new WeakReference(target)); + } + + return; + } + } + } + + internal void Unregister(TKey key, TTarget target) + { + if (!_targets.TryGetValue(key, out var targets)) + { + return; + } + + lock (targets) + { + if (!IsCurrent(key, targets)) + { + return; + } + + targets.RemoveAll(reference => + !reference.TryGetTarget(out var current) || ReferenceEquals(current, target)); + RemoveIfEmpty(key, targets); + } + } + + internal List GetLiveTargets(TKey key) + { + while (_targets.TryGetValue(key, out var targets)) + { + lock (targets) + { + if (!IsCurrent(key, targets)) + { + continue; + } + + var liveTargets = new List(targets.Count); + targets.RemoveAll(reference => !reference.TryGetTarget(out _)); + foreach (var reference in targets) + { + if (reference.TryGetTarget(out var target)) + { + liveTargets.Add(target); + } + } + + RemoveIfEmpty(key, targets); + return liveTargets; + } + } + + return []; + } + + internal int GetRegistrationCount(TKey key) + { + while (_targets.TryGetValue(key, out var targets)) + { + lock (targets) + { + if (!IsCurrent(key, targets)) + { + continue; + } + + return targets.Count; + } + } + + return 0; + } + + private bool IsCurrent(TKey key, List> targets) + { + return _targets.TryGetValue(key, out var currentTargets) && + ReferenceEquals(targets, currentTargets); + } + + private void RemoveIfEmpty(TKey key, List> targets) + { + if (targets.Count == 0) + { + ((ICollection>>>)_targets) + .Remove(new KeyValuePair>>(key, targets)); + } + } +} diff --git a/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs b/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs index 87c3e3b4b9..ee03bb5863 100644 --- a/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs +++ b/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs @@ -51,6 +51,45 @@ public partial class JSAdapterTests JSPropertyChangeRegistry.Unregister(fake.Connection, commandId, liveTarget); } + [TestMethod] + public void WeakReferenceRegistry_PrunesDeadTargetsDuringRegistration() + { + var registry = new JSWeakReferenceRegistry(); + var deadTarget = RegisterTemporaryTarget(registry, "page"); + GC.Collect(); + GC.WaitForPendingFinalizers(); + GC.Collect(); + + Assert.IsFalse(deadTarget.TryGetTarget(out _)); + Assert.AreEqual(1, registry.GetRegistrationCount("page")); + + var liveTarget = new object(); + registry.Register("page", liveTarget); + + Assert.AreEqual(1, registry.GetRegistrationCount("page")); + Assert.AreSame(liveTarget, registry.GetLiveTargets("page")[0]); + } + + [TestMethod] + public void WeakReferenceRegistry_ConcurrentReplacementKeepsNewestTarget() + { + for (var i = 0; i < 128; i++) + { + var registry = new JSWeakReferenceRegistry(); + var oldTarget = new object(); + var newTarget = new object(); + registry.Register("page", oldTarget); + + Parallel.Invoke( + () => registry.Unregister("page", oldTarget), + () => registry.Register("page", newTarget)); + + var liveTargets = registry.GetLiveTargets("page"); + Assert.AreEqual(1, liveTargets.Count); + Assert.AreSame(newTarget, liveTargets[0]); + } + } + [TestMethod] public void NestedProxyGetters_CacheIdentityAndRegistration() { @@ -139,6 +178,7 @@ public partial class JSAdapterTests fake.Connection); var originalDetails = item.Details; var originalMoreCommands = item.MoreCommands; + var originalCommand = ((ICommandContextItem)originalMoreCommands[0]).Command; JSPropertyChangeRegistry.Dispatch( fake.Connection, @@ -167,8 +207,45 @@ public partial class JSAdapterTests Assert.AreEqual("New details", item.Details?.Title); Assert.AreNotSame(originalDetails, item.Details); Assert.AreNotSame(originalMoreCommands, item.MoreCommands); - Assert.AreEqual(0, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "old-more")); + Assert.AreSame(item.MoreCommands, item.MoreCommands); + Assert.AreEqual(1, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "old-more")); Assert.AreEqual(1, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "new-more")); + + JSPropertyChangeRegistry.Dispatch( + fake.Connection, + ParseElement(new JsonObject + { + ["commandId"] = "old-more", + ["properties"] = new JsonObject { ["name"] = "Still alive" }, + })); + + Assert.AreEqual("Still alive", originalCommand.Name); + (originalCommand as IDisposable)?.Dispose(); + } + + [TestMethod] + public async Task PageRegistries_RouteToNewestProxyAfterOlderProxyIsRemoved() + { + using var fake = new JSFakeExtension(); + using var oldListPage = new JSListPageProxy("list-page", fake.Connection); + using var newListPage = new JSListPageProxy("list-page", fake.Connection); + using var oldContentPage = new JSContentPageProxy("content-page", fake.Connection); + using var newContentPage = new JSContentPageProxy("content-page", fake.Connection); + var listChanged = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var contentChanged = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + newListPage.ItemsChanged += (_, _) => listChanged.TrySetResult(); + newContentPage.ItemsChanged += (_, _) => contentChanged.TrySetResult(); + + oldListPage.Dispose(); + oldContentPage.Dispose(); + await fake.PushNotificationAsync( + "listPage/itemsChanged", + new JsonObject { ["pageId"] = "list-page" }); + await fake.PushNotificationAsync( + "contentPage/itemsChanged", + new JsonObject { ["pageId"] = "content-page" }); + + await Task.WhenAll(listChanged.Task, contentChanged.Task).WaitAsync(Timeout); } [TestMethod] @@ -428,6 +505,16 @@ public partial class JSAdapterTests return new WeakReference(target); } + [MethodImpl(MethodImplOptions.NoInlining)] + private static WeakReference RegisterTemporaryTarget( + JSWeakReferenceRegistry registry, + string key) + { + var target = new object(); + registry.Register(key, target); + return new WeakReference(target); + } + private static JsonObject Command(string id) => new() { ["id"] = id,