Allow more specific LSP services to override Any LSP services

src/LanguageServer/Microsoft.CodeAnalysis.LanguageServer/Telemetry/VSCodeRequestTelemetryLogger.cs

@@ -0,0 +1,62 @@
+// 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.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Composition;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Internal.Log;
+using Microsoft.CodeAnalysis.LanguageServer.Handler;
+using Microsoft.CodeAnalysis.Telemetry;
+
+//namespace Microsoft.CodeAnalysis.LanguageServer.Telemetry;
+
+//internal class VSCodeRequestTelemetryLogger : RequestTelemetryLogger
+//{
+//    /// <summary>
+//    /// Tracks whether or not the initial project load has completed so we can see
+//    /// how often we get misc file requests after we've loaded.
+//    /// </summary>
+//    private static bool _initialProjectLoadCompleted = false;
+
+//    private readonly ConcurrentDictionary<bool, ConcurrentDictionary<string, Counter>> _findDocumentCounters;
+
+//    public VSCodeRequestTelemetryLogger(string serverTypeName) : base(serverTypeName)
+//    {
+//    }
+
+//    public static void ReportProjectInitializationComplete()
+//    {
+//        _initialProjectLoadCompleted = true;
+//        Logger.Log(FunctionId.VSCode_Projects_Load_Completed, logLevel: LogLevel.Information);
+//    }
+
+//    public static void ReportProjectLoadStarted()
+//    {
+//        Logger.Log(FunctionId.VSCode_Project_Load_Started, logLevel: LogLevel.Information);
+//    }
+
+//    protected override void IncreaseFindDocumentCount(string workspaceInfo)
+//    {
+//        TelemetryLogging.LogAggregated(FunctionId.LSP_FindDocumentInWorkspace, KeyValueLogMessage.Create(m =>
+//        {
+//            m[TelemetryLogging.KeyName] = _serverTypeName;
+//            m[TelemetryLogging.KeyValue] = (int)queuedDuration.TotalMilliseconds;
+//            m[TelemetryLogging.KeyMetricName] = "Count";
+//            m["server"] = _serverTypeName;
+//            m["method"] = methodName;
+//            m["language"] = language;
+//        }));
+
+//        base.IncreaseFindDocumentCount(workspaceInfo);
+//    }
+
+//    protected override void ReportFindDocumentCounter()
+//    {
+//        base.ReportFindDocumentCounter();
+//    }
+//}

src/LanguageServer/Protocol/Handler/RequestTelemetryLogger.cs → src/LanguageServer/Protocol/Handler/Telemetry/RequestTelemetryLogger.cs

@@ -14,7 +14,7 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 /// Logs metadata on LSP requests (duration, success / failure metrics)
 /// for this particular LSP server instance.
 /// </summary>
-internal sealed class RequestTelemetryLogger : IDisposable, ILspService
+internal class RequestTelemetryLogger : IDisposable, ILspService
 {
     private readonly string _serverTypeName;
 
@@ -46,10 +46,15 @@ public void UpdateFindDocumentTelemetryData(bool success, string? workspaceKind)
 
         if (workspaceKindTelemetryProperty != null)
         {
-            _findDocumentResults.IncreaseCount(workspaceKindTelemetryProperty);
+            IncreaseFindDocumentCount(workspaceKindTelemetryProperty);
         }
     }
 
+    protected virtual void IncreaseFindDocumentCount(string workspaceInfo)
+    {
+        _findDocumentResults.IncreaseCount(workspaceInfo);
+    }
+
     public void UpdateUsedForkedSolutionCounter(bool usedForkedSolution)
     {
         _usedForkedSolutionCounter.IncreaseCount(usedForkedSolution);
@@ -103,6 +108,22 @@ public void Dispose()
         TelemetryLogging.Flushed -= OnFlushed;
     }
 
+    protected virtual void ReportFindDocumentCounter()
+    {
+        if (!_findDocumentResults.IsEmpty)
+        {
+            TelemetryLogging.Log(FunctionId.LSP_FindDocumentInWorkspace, KeyValueLogMessage.Create(LogType.Trace, m =>
+            {
+                m["server"] = _serverTypeName;
+                foreach (var kvp in _findDocumentResults)
+                {
+                    var info = kvp.Key.ToString()!;
+                    m[info] = kvp.Value.GetCount();
+                }
+            }));
+        }
+    }
+
     private void OnFlushed(object? sender, EventArgs e)
     {
         foreach (var kvp in _requestCounters)
@@ -118,18 +139,7 @@ private void OnFlushed(object? sender, EventArgs e)
             }));
         }
 
-        if (!_findDocumentResults.IsEmpty)
-        {
-            TelemetryLogging.Log(FunctionId.LSP_FindDocumentInWorkspace, KeyValueLogMessage.Create(LogType.Trace, m =>
-            {
-                m["server"] = _serverTypeName;
-                foreach (var kvp in _findDocumentResults)
-                {
-                    var info = kvp.Key.ToString()!;
-                    m[info] = kvp.Value.GetCount();
-                }
-            }));
-        }
+        ReportFindDocumentCounter();
 
         if (!_usedForkedSolutionCounter.IsEmpty)
         {
@@ -149,7 +159,7 @@ private void OnFlushed(object? sender, EventArgs e)
         _usedForkedSolutionCounter.Clear();
     }
 
-    private class Counter
+    protected class Counter
     {
         private int _succeededCount;
         private int _failedCount;

src/LanguageServer/Protocol/LspServices/LspServices.cs

@@ -7,6 +7,7 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CommonLanguageServerProtocol.Framework;
@@ -40,36 +41,48 @@ public LspServices(
     {
         var serviceMap = new Dictionary<string, Lazy<ILspService, LspServiceMetadataView>>();
 
-        // Convert MEF exported service factories to the lazy LSP services that they create.
-        foreach (var lazyServiceFactory in mefLspServiceFactories)
-        {
-            var metadata = lazyServiceFactory.Metadata;
+        // Add services from factories exported for this server kind.
+        foreach (var lazyServiceFactory in mefLspServiceFactories.Where(f => f.Metadata.ServerKind == serverKind))
+            AddSpecificService(new(() => lazyServiceFactory.Value.CreateILspService(this, serverKind), lazyServiceFactory.Metadata));
 
-            // Make sure that we only include services exported for the specified server kind (or NotSpecified).
-            if (metadata.ServerKind == serverKind ||
-                metadata.ServerKind == WellKnownLspServerKinds.Any)
-            {
-                serviceMap.Add(
-                    metadata.TypeRef.TypeName,
-                    new(() => lazyServiceFactory.Value.CreateILspService(this, serverKind), metadata));
-            }
-        }
+        // Add services exported for this server kind.
+        foreach (var lazyService in mefLspServices.Where(s => s.Metadata.ServerKind == serverKind))
+            AddSpecificService(lazyService);
 
-        foreach (var lazyService in mefLspServices)
-        {
-            var metadata = lazyService.Metadata;
+        // Add services from factories exported for any (if there is not already an existing service for the specific server kind).
+        foreach (var lazyServiceFactory in mefLspServiceFactories.Where(f => f.Metadata.ServerKind == WellKnownLspServerKinds.Any))
+            TryAddAnyService(new(() => lazyServiceFactory.Value.CreateILspService(this, serverKind), lazyServiceFactory.Metadata));
 
-            // Make sure that we only include services exported for the specified server kind (or NotSpecified).
-            if (metadata.ServerKind == serverKind ||
-                metadata.ServerKind == WellKnownLspServerKinds.Any)
-            {
-                serviceMap.Add(metadata.TypeRef.TypeName, lazyService);
-            }
-        }
+        // Add services exported for any (if there is not already an existing service for the specific server kind).
+        foreach (var lazyService in mefLspServices.Where(s => s.Metadata.ServerKind == WellKnownLspServerKinds.Any))
+            TryAddAnyService(lazyService);
 
         _lazyMefLspServices = serviceMap.ToFrozenDictionary();
 
         _baseServices = baseServices;
+
+        void AddSpecificService(Lazy<ILspService, LspServiceMetadataView> serviceGetter)
+        {
+            var metadata = serviceGetter.Metadata;
+            Contract.ThrowIfFalse(metadata.ServerKind == serverKind);
+            serviceMap.Add(metadata.TypeRef.TypeName, serviceGetter);
+        }
+
+        void TryAddAnyService(Lazy<ILspService, LspServiceMetadataView> serviceGetter)
+        {
+            var metadata = serviceGetter.Metadata;
+            Contract.ThrowIfFalse(metadata.ServerKind == WellKnownLspServerKinds.Any);
+            if (!serviceMap.TryGetValue(metadata.TypeRef.TypeName, out var existing))
+            {
+                serviceMap.Add(metadata.TypeRef.TypeName, serviceGetter);
+            }
+            else
+            {
+                // Make sure we're not trying to add a duplicate Any service, but otherwise we should skip adding
+                // this service as we already have a more specific service available.
+                Contract.ThrowIfTrue(existing.Metadata.ServerKind == WellKnownLspServerKinds.Any);
+            }
+        }
     }
 
     public T GetRequiredService<T>() where T : notnull

src/LanguageServer/ProtocolUnitTests/LspServicesTests.cs

@@ -0,0 +1,120 @@
+// 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.
+
+using System;
+using System.Composition;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Test.Utilities;
+using Roslyn.Test.Utilities;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.UnitTests;
+
+[UseExportProvider]
+public class LspServicesTests(ITestOutputHelper testOutputHelper) : AbstractLanguageServerProtocolTests(testOutputHelper)
+{
+    [Theory, CombinatorialData]
+    public async Task ReturnsSpecificLspService(bool mutatingLspWorkspace)
+    {
+        var composition = base.Composition.AddParts(typeof(CSharpLspService), typeof(CSharpLspServiceFactory));
+        await using var server = await CreateTestLspServerAsync("", mutatingLspWorkspace, initializationOptions: new() { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer }, composition);
+
+        var lspService = server.GetRequiredLspService<TestLspService>();
+        Assert.True(lspService is CSharpLspService);
+
+        var lspServiceFromFactory = server.GetRequiredLspService<TestLspServiceFromFactory>();
+        Assert.Equal(typeof(CSharpLspServiceFactory).Name, lspServiceFromFactory.FactoryName);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task SpecificLspServiceOverridesAny(bool mutatingLspWorkspace)
+    {
+        var composition = base.Composition.AddParts(typeof(CSharpLspService), typeof(AnyLspService), typeof(CSharpLspServiceFactory), typeof(AnyLspServiceFactory));
+        await using var server = await CreateTestLspServerAsync("", mutatingLspWorkspace, initializationOptions: new() { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer }, composition);
+
+        var lspService = server.GetRequiredLspService<TestLspService>();
+        Assert.True(lspService is CSharpLspService);
+
+        var lspServiceFromFactory = server.GetRequiredLspService<TestLspServiceFromFactory>();
+        Assert.Equal(typeof(CSharpLspServiceFactory).Name, lspServiceFromFactory.FactoryName);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task ReturnsAnyLspService(bool mutatingLspWorkspace)
+    {
+        var composition = base.Composition.AddParts(typeof(AnyLspService), typeof(AnyLspServiceFactory));
+        await using var server = await CreateTestLspServerAsync("", mutatingLspWorkspace, initializationOptions: new() { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer }, composition);
+
+        var lspService = server.GetRequiredLspService<TestLspService>();
+        Assert.True(lspService is AnyLspService);
+
+        var lspServiceFromFactory = server.GetRequiredLspService<TestLspServiceFromFactory>();
+        Assert.Equal(typeof(AnyLspServiceFactory).Name, lspServiceFromFactory.FactoryName);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task DuplicateSpecificServicesThrow(bool mutatingLspWorkspace)
+    {
+        var composition = base.Composition.AddParts(typeof(CSharpLspService), typeof(CSharpLspServiceFactory), typeof(DuplicateCSharpLspService), typeof(DuplicateCSharpLspServiceFactory));
+        await Assert.ThrowsAnyAsync<Exception>(async () => await CreateTestLspServerAsync("", mutatingLspWorkspace, initializationOptions: new() { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer }, composition));
+    }
+
+    [Theory, CombinatorialData]
+    public async Task DuplicateAnyServicesThrow(bool mutatingLspWorkspace)
+    {
+        var composition = base.Composition.AddParts(typeof(AnyLspService), typeof(AnyLspServiceFactory), typeof(DuplicateAnyLspService), typeof(DuplicateAnyLspServiceFactory));
+        await Assert.ThrowsAnyAsync<Exception>(async () => await CreateTestLspServerAsync("", mutatingLspWorkspace, initializationOptions: new() { ServerKind = WellKnownLspServerKinds.CSharpVisualBasicLspServer }, composition));
+    }
+
+    internal class TestLspService : ILspService { }
+
+    internal record class TestLspServiceFromFactory(string FactoryName) : ILspService { }
+
+    internal class TestLspServiceFactory : ILspServiceFactory
+    {
+        public ILspService CreateILspService(LspServices lspServices, WellKnownLspServerKinds serverKind) => new TestLspServiceFromFactory(this.GetType().Name);
+    }
+
+    [ExportStatelessLspService(typeof(TestLspService), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.CSharpVisualBasicLspServer), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class CSharpLspService() : TestLspService { }
+
+    [ExportLspServiceFactory(typeof(TestLspServiceFromFactory), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.CSharpVisualBasicLspServer), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class CSharpLspServiceFactory() : TestLspServiceFactory { }
+
+    [ExportStatelessLspService(typeof(TestLspService), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.Any), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class AnyLspService() : TestLspService { }
+
+    [ExportLspServiceFactory(typeof(TestLspServiceFromFactory), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.Any), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class AnyLspServiceFactory() : TestLspServiceFactory { }
+
+    [ExportStatelessLspService(typeof(TestLspService), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.CSharpVisualBasicLspServer), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class DuplicateCSharpLspService() : TestLspService { }
+
+    [ExportLspServiceFactory(typeof(TestLspServiceFromFactory), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.CSharpVisualBasicLspServer), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class DuplicateCSharpLspServiceFactory() : CSharpLspServiceFactory { }
+
+    [ExportStatelessLspService(typeof(TestLspService), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.Any), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class DuplicateAnyLspService() : TestLspService { }
+
+    [ExportLspServiceFactory(typeof(TestLspServiceFromFactory), ProtocolConstants.RoslynLspLanguagesContract, WellKnownLspServerKinds.Any), Shared]
+    [method: ImportingConstructor]
+    [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+    internal class DuplicateAnyLspServiceFactory() : CSharpLspServiceFactory { }
+}

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Log/FunctionId.cs

@@ -634,5 +634,8 @@ internal enum FunctionId
     Copilot_On_The_Fly_Docs_Results_Canceled = 814,
     Copilot_On_The_Fly_Docs_Get_Counts = 815,
     Copilot_On_The_Fly_Docs_Content_Excluded = 816,
-    Copilot_Rename = 851
+    Copilot_Rename = 851,
+
+    VSCode_Project_Load_Started = 860,
+    VSCode_Projects_Load_Completed = 861,
 }