Merge pull request #1582 from dmgonch/dmgonch/fix/RequestHandlersExtensibility

Fixes bugs when there are multiple handlers are defined for a language for a given request

Author: david-driscoll

src/OmniSharp.Abstractions/Models/ICanBeEmptyResponse.cs

@@ -0,0 +1,7 @@
+namespace OmniSharp.Models
+{
+    internal interface ICanBeEmptyResponse
+    {
+        bool IsEmpty { get; }
+    }
+}

src/OmniSharp.Abstractions/Models/v1/GotoDefinition/GotoDefinitionResponse.cs

@@ -3,13 +3,14 @@
 
 namespace OmniSharp.Models.GotoDefinition
 {
-    public class GotoDefinitionResponse
+    public class GotoDefinitionResponse : ICanBeEmptyResponse
     {
         public string FileName { get; set; }
         [JsonConverter(typeof(ZeroBasedIndexConverter))]
         public int Line { get; set; }
         [JsonConverter(typeof(ZeroBasedIndexConverter))]
         public int Column { get; set; }
         public MetadataSource MetadataSource { get; set; }
+        public bool IsEmpty => FileName == null || FileName == string.Empty;
     }
 }

src/OmniSharp.Host/AssemblyInfo.cs

@@ -3,4 +3,5 @@
 [assembly: InternalsVisibleTo("OmniSharp.Http.Tests")]
 [assembly: InternalsVisibleTo("OmniSharp.MSBuild.Tests")]
 [assembly: InternalsVisibleTo("OmniSharp.Roslyn.CSharp.Tests")]
+[assembly: InternalsVisibleTo("OmniSharp.Stdio.Tests")]
 [assembly: InternalsVisibleTo("TestUtility")]

src/OmniSharp.Host/Endpoint/EndpointHandler.cs

@@ -48,7 +48,7 @@ public class EndpointHandler<TRequest, TResponse> : EndpointHandler
     {
         private readonly CompositionHost _host;
         private readonly IPredicateHandler _languagePredicateHandler;
-        private readonly Lazy<Task<Dictionary<string, ExportHandler<TRequest, TResponse>>>> _exports;
+        private readonly Lazy<Task<Dictionary<string, ExportHandler<TRequest, TResponse>[]>>> _exports;
         private readonly OmniSharpWorkspace _workspace;
         private readonly bool _hasLanguageProperty;
         private readonly bool _hasFileNameProperty;
@@ -71,10 +71,10 @@ public EndpointHandler(IPredicateHandler languagePredicateHandler, CompositionHo
             _canBeAggregated = typeof(IAggregateResponse).IsAssignableFrom(metadata.ResponseType);
             _updateBufferHandler = updateBufferHandler;
 
-            _exports = new Lazy<Task<Dictionary<string, ExportHandler<TRequest, TResponse>>>>(() => LoadExportHandlers(handlers));
+            _exports = new Lazy<Task<Dictionary<string, ExportHandler<TRequest, TResponse>[]>>>(() => LoadExportHandlers(handlers));
         }
 
-        private Task<Dictionary<string, ExportHandler<TRequest, TResponse>>> LoadExportHandlers(IEnumerable<Lazy<IRequestHandler, OmniSharpRequestHandlerMetadata>> handlers)
+        private Task<Dictionary<string, ExportHandler<TRequest, TResponse>[]>> LoadExportHandlers(IEnumerable<Lazy<IRequestHandler, OmniSharpRequestHandlerMetadata>> handlers)
         {
             var interfaceHandlers = handlers
                 .Select(export => new RequestHandlerExportHandler<TRequest, TResponse>(export.Metadata.Language, (IRequestHandler<TRequest, TResponse>)export.Value))
@@ -84,9 +84,13 @@ private Task<Dictionary<string, ExportHandler<TRequest, TResponse>>> LoadExportH
                 .Select(plugin => new PluginExportHandler<TRequest, TResponse>(EndpointName, plugin))
                 .Cast<ExportHandler<TRequest, TResponse>>();
 
+            // Group handlers by language and sort each group for consistency
             return Task.FromResult(interfaceHandlers
-               .Concat(plugins)
-               .ToDictionary(export => export.Language));
+                .Concat(plugins)
+                .GroupBy(export => export.Language, StringComparer.OrdinalIgnoreCase)
+                .ToDictionary(
+                    group => group.Key,
+                    group => group.OrderBy(g => g).ToArray()));
         }
 
         public string EndpointName { get; }
@@ -142,18 +146,88 @@ private Task<object> HandleLanguageRequest(string language, TRequest request, Re
         {
             if (!string.IsNullOrEmpty(language))
             {
-                return HandleSingleRequest(language, request, packet);
+                return HandleRequestForLanguage(language, request, packet);
             }
 
             return HandleAllRequest(request, packet);
         }
 
-        private async Task<object> HandleSingleRequest(string language, TRequest request, RequestPacket packet)
+        private async Task<IAggregateResponse> AggregateResponsesFromLanguageHandlers(ExportHandler<TRequest, TResponse>[] handlers, TRequest request)
+        {
+            if (!_canBeAggregated)
+            {
+                throw new NotSupportedException($"Must be able to aggregate responses from all handlers for {EndpointName}");
+            }
+
+            IAggregateResponse aggregateResponse = null;
+
+            if (handlers.Length == 1)
+            {
+                var response = handlers[0].Handle(request);
+                return (IAggregateResponse)await response;
+            }
+            else
+            {
+                var responses = new List<Task<TResponse>>();
+                foreach (var handler in handlers)
+                {
+                    responses.Add(handler.Handle(request));
+                }
+
+                foreach (IAggregateResponse response in await Task.WhenAll(responses))
+                {
+                    if (aggregateResponse != null)
+                    {
+                        aggregateResponse = aggregateResponse.Merge(response);
+                    }
+                    else
+                    {
+                        aggregateResponse = response;
+                    }
+                }
+            }
+
+            return aggregateResponse;
+        }
+
+        private async Task<object> GetFirstNotEmptyResponseFromHandlers(ExportHandler<TRequest, TResponse>[] handlers, TRequest request)
+        {
+            var responses = new List<Task<TResponse>>();
+            foreach (var handler in handlers)
+            {
+                responses.Add(handler.Handle(request));
+            }
+
+            foreach (object response in await Task.WhenAll(responses))
+            {
+                var canBeEmptyResponse = response as ICanBeEmptyResponse;
+                if (canBeEmptyResponse != null)
+                {
+                    if (!canBeEmptyResponse.IsEmpty)
+                    {
+                        return response;
+                    }
+                }
+                else if (response != null)
+                {
+                    return response;
+                }
+            }
+
+            return null;
+        }
+
+        private async Task<object> HandleRequestForLanguage(string language, TRequest request, RequestPacket packet)
         {
             var exports = await _exports.Value;
-            if (exports.TryGetValue(language, out var handler))
+            if (exports.TryGetValue(language, out var handlers))
             {
-                return await handler.Handle(request);
+                if (_canBeAggregated)
+                {
+                    return await AggregateResponsesFromLanguageHandlers(handlers, request);
+                }
+
+                return await GetFirstNotEmptyResponseFromHandlers(handlers, request);
             }
 
             throw new NotSupportedException($"{language} does not support {EndpointName}");
@@ -169,11 +243,10 @@ private async Task<object> HandleAllRequest(TRequest request, RequestPacket pack
             var exports = await _exports.Value;
 
             IAggregateResponse aggregateResponse = null;
-
-            var responses = new List<Task<TResponse>>();
-            foreach (var handler in exports.Values)
+            var responses = new List<Task<IAggregateResponse>>();
+            foreach (var export in exports)
             {
-                responses.Add(handler.Handle(request));
+                responses.Add(AggregateResponsesFromLanguageHandlers(export.Value, request));
             }
 
             foreach (IAggregateResponse exportResponse in await Task.WhenAll(responses))

src/OmniSharp.Host/Endpoint/Exports/ExportHandler.cs

@@ -1,15 +1,18 @@
+using System;
 using System.Threading.Tasks;
 
 namespace OmniSharp.Endpoint.Exports
 {
-    abstract class ExportHandler<TRequest, TResponse>
+    abstract class ExportHandler<TRequest, TResponse> : IComparable<ExportHandler<TRequest, TResponse>>
     {
         protected ExportHandler(string language)
         {
             Language = language;
         }
 
         public string Language { get; }
+
+        public abstract int CompareTo(ExportHandler<TRequest, TResponse> other);
         public abstract Task<TResponse> Handle(TRequest request);
     }
 }

src/OmniSharp.Host/Endpoint/Exports/PluginExportHandler.cs

@@ -14,6 +14,17 @@ public PluginExportHandler(string endpoint, Plugin plugin) : base(plugin.Config.
             _plugin = plugin;
         }
 
+        public override int CompareTo(ExportHandler<TRequest, TResponse> other)
+        {
+            var otherPlugin = other as PluginExportHandler<TRequest, TResponse>;
+            if (otherPlugin == null)
+            {
+                return 1;
+            }
+
+            return _plugin.Key.CompareTo(otherPlugin._plugin.Key);
+        }
+
         public override Task<TResponse> Handle(TRequest request)
         {
             return _plugin.Handle<TRequest, TResponse>(_endpoint, request);

src/OmniSharp.Host/Endpoint/Exports/RequestHandlerExportHandler.cs

@@ -13,6 +13,17 @@ public RequestHandlerExportHandler(string language, IRequestHandler<TRequest, TR
             _handler = handler;
         }
 
+        public override int CompareTo(ExportHandler<TRequest, TResponse> other)
+        {
+            var otherHandler = other as RequestHandlerExportHandler<TRequest, TResponse>;
+            if (otherHandler == null)
+            {
+                return 1;
+            }
+
+            return _handler.GetType().ToString().CompareTo(otherHandler._handler.GetType().ToString());
+        }
+
         public override Task<TResponse> Handle(TRequest request)
         {
             return _handler.Handle(request);

src/OmniSharp.Host/Mef/MefValueProvider.cs

@@ -7,10 +7,12 @@ namespace OmniSharp.Mef
     internal class MefValueProvider<T> : ExportDescriptorProvider
     {
         private readonly T _item;
+        private readonly IDictionary<string, object> _metadata;
 
-        public MefValueProvider(T item)
+        public MefValueProvider(T item, IDictionary<string, object> metadata)
         {
             _item = item;
+            _metadata = metadata;
         }
 
         public override IEnumerable<ExportDescriptorPromise> GetExportDescriptors(CompositionContract contract, DependencyAccessor descriptorAccessor)
@@ -19,16 +21,16 @@ public override IEnumerable<ExportDescriptorPromise> GetExportDescriptors(Compos
             {
                 yield return new ExportDescriptorPromise(contract, string.Empty, true,
                     () => Enumerable.Empty<CompositionDependency>(),
-                    deps => ExportDescriptor.Create((context, operation) => _item, new Dictionary<string, object>()));
+                    deps => ExportDescriptor.Create((context, operation) => _item, _metadata ?? new Dictionary<string, object>()));
             }
         }
     }
 
     internal static class MefValueProvider
     {
-        public static MefValueProvider<T> From<T>(T value)
+        public static MefValueProvider<T> From<T>(T value, IDictionary<string, object> metadata = null)
         {
-            return new MefValueProvider<T>(value);
+            return new MefValueProvider<T>(value, metadata);
         }
     }
 }