Fix threadpool load induced delayed responses to plugin requests

src/NuGet.Core/NuGet.Protocol/Plugins/InboundRequestContext.cs

@@ -17,6 +17,7 @@ public sealed class InboundRequestContext : IDisposable
         private readonly IConnection _connection;
         private bool _isDisposed;
         private readonly IPluginLogger _logger;
+        private readonly InboundRequestProcessingHandler _inboundRequestProcessingHandler;
 
         /// <summary>
         /// Gets the request ID.
@@ -37,7 +38,7 @@ public InboundRequestContext(
             IConnection connection,
             string requestId,
             CancellationToken cancellationToken)
-            : this(connection, requestId, cancellationToken, PluginLogger.DefaultInstance)
+            : this(connection, requestId, cancellationToken, new InboundRequestProcessingHandler(), PluginLogger.DefaultInstance)
         {
         }
 
@@ -52,12 +53,15 @@ public InboundRequestContext(
         /// is <c>null</c>.</exception>
         /// <exception cref="ArgumentException">Thrown if <paramref name="requestId" />
         /// is either <c>null</c> or an empty string.</exception>
-        /// <exception cref="ArgumentNullException">Thrown if <paramref name="logger" />
+        /// <exception cref="ArgumentNullException">Thrown if <paramref name="inboundRequestProcessingHandler" />
+        /// is <c>null</c>.</exception>
+        /// /// <exception cref="ArgumentNullException">Thrown if <paramref name="logger" />
         /// is <c>null</c>.</exception>
         internal InboundRequestContext(
             IConnection connection,
             string requestId,
             CancellationToken cancellationToken,
+            InboundRequestProcessingHandler inboundRequestProcessingHandler,
             IPluginLogger logger)
         {
             if (connection == null)
@@ -75,6 +79,11 @@ internal InboundRequestContext(
                 throw new ArgumentNullException(nameof(logger));
             }
 
+            if (inboundRequestProcessingHandler == null)
+            {
+                throw new ArgumentNullException(nameof(inboundRequestProcessingHandler));
+            }
+
             _connection = connection;
             RequestId = requestId;
 
@@ -85,6 +94,42 @@ internal InboundRequestContext(
             _cancellationToken = _cancellationTokenSource.Token;
 
             _logger = logger;
+
+            _inboundRequestProcessingHandler = inboundRequestProcessingHandler;
+        }
+
+        private async Task ProcessResponseAsync(IRequestHandler requestHandler, Message request, IResponseHandler responseHandler)
+        {
+            try
+            {
+                if (_logger.IsEnabled)
+                {
+                    _logger.Write(new TaskLogMessage(_logger.Now, request.RequestId, request.Method, request.Type, TaskState.Executing));
+                }
+
+                await requestHandler.HandleResponseAsync(
+                    _connection,
+                    request,
+                    responseHandler,
+                    _cancellationToken);
+            }
+            catch (OperationCanceledException) when (_cancellationToken.IsCancellationRequested)
+            {
+                var response = MessageUtilities.Create(request.RequestId, MessageType.Cancel, request.Method);
+
+                await _connection.SendAsync(response, CancellationToken.None);
+            }
+            catch (Exception ex)
+            {
+                BeginFaultAsync(request, ex);
+            }
+            finally
+            {
+                if (_logger.IsEnabled)
+                {
+                    _logger.Write(new TaskLogMessage(_logger.Now, request.RequestId, request.Method, request.Type, TaskState.Completed));
+                }
+            }
         }
 
         /// <summary>
@@ -108,8 +153,7 @@ public void Dispose()
             {
             }
 
-            // Do not dispose of _connection or _logger.  This context does not own them.
-
+            // Do not dispose of the _connection, _logger or _requestProcessingContext.  This context does not own them.
             GC.SuppressFinalize(this);
 
             _isDisposed = true;
@@ -210,42 +254,9 @@ public void BeginResponseAsync(
             {
                 _logger.Write(new TaskLogMessage(_logger.Now, request.RequestId, request.Method, request.Type, TaskState.Queued));
             }
+            Func<Task> task = async () => await ProcessResponseAsync(requestHandler, request, responseHandler);
 
-            Task.Run(async () =>
-                {
-                    // Top-level exception handler for a worker pool thread.
-                    try
-                    {
-                        if (_logger.IsEnabled)
-                        {
-                            _logger.Write(new TaskLogMessage(_logger.Now, request.RequestId, request.Method, request.Type, TaskState.Executing));
-                        }
-
-                        await requestHandler.HandleResponseAsync(
-                            _connection,
-                            request,
-                            responseHandler,
-                            _cancellationToken);
-                    }
-                    catch (OperationCanceledException) when (_cancellationToken.IsCancellationRequested)
-                    {
-                        var response = MessageUtilities.Create(request.RequestId, MessageType.Cancel, request.Method);
-
-                        await _connection.SendAsync(response, CancellationToken.None);
-                    }
-                    catch (Exception ex)
-                    {
-                        BeginFaultAsync(request, ex);
-                    }
-                    finally
-                    {
-                        if (_logger.IsEnabled)
-                        {
-                            _logger.Write(new TaskLogMessage(_logger.Now, request.RequestId, request.Method, request.Type, TaskState.Completed));
-                        }
-                    }
-                },
-                _cancellationToken);
+            _inboundRequestProcessingHandler.Handle(request.Method, task, _cancellationToken);
         }
 
         /// <summary>
@@ -262,4 +273,4 @@ public void Cancel()
             }
         }
     }
-}
+}

src/NuGet.Core/NuGet.Protocol/Plugins/InboundRequestProcessingHandler.cs

@@ -0,0 +1,81 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace NuGet.Protocol.Plugins
+{
+    internal class InboundRequestProcessingHandler : IDisposable
+    {
+        private readonly ISet<MessageMethod> _fastProccessingMethods;
+        private readonly Lazy<DedicatedAsynchronousProcessingThread> _processingThread;
+        private bool _isDisposed;
+
+        public InboundRequestProcessingHandler() :
+            this(new HashSet<MessageMethod>())
+        {
+        }
+
+        /// <summary>
+        /// Requests from the processing methods provided in this set are handled on a dedicated thread.
+        /// </summary>
+        /// <param name="fastProcessingMethods"></param>
+        public InboundRequestProcessingHandler(ISet<MessageMethod> fastProcessingMethods) 
+        {
+            _fastProccessingMethods = fastProcessingMethods ?? throw new ArgumentNullException(nameof(fastProcessingMethods));
+            // Lazily initialize the processing thread. It is not needed if there are no time critical methods.
+            _processingThread = new Lazy<DedicatedAsynchronousProcessingThread>(() =>
+            {
+                var thread = new DedicatedAsynchronousProcessingThread();
+                thread.Start();
+                return thread;
+            });
+
+        }
+
+        /// <summary>
+        /// Methods that are in the fast processing method list will be handled on a separate thread. Everything else will be queued on the threadpool.
+        /// </summary>
+        /// <param name="messageMethod"></param>
+        /// <param name="task"></param>
+        /// <param name="cancellationToken"></param>
+        internal void Handle(MessageMethod messageMethod, Func<Task> task, CancellationToken cancellationToken)
+        {
+            ThrowIfDisposed();
+            if (_fastProccessingMethods.Contains(messageMethod))
+            {
+                _processingThread.Value.Enqueue(task);
+            }
+            else
+            {
+                Task.Run(async () =>
+                {
+                    await task();
+                }, cancellationToken);
+            }
+        }
+
+        public void Dispose()
+        {
+            if (_isDisposed)
+            {
+                return;
+            }
+            _processingThread.Value.Dispose();
+            GC.SuppressFinalize(this);
+
+            _isDisposed = true;
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (_isDisposed)
+            {
+                throw new ObjectDisposedException(GetType().Name);
+            }
+        }
+    }
+}

src/NuGet.Core/NuGet.Protocol/Plugins/Logging/AssemblyLogMessage.cs

@@ -12,15 +12,18 @@ internal sealed class AssemblyLogMessage : PluginLogMessage
         private readonly string _fileVersion;
         private readonly string _fullName;
         private readonly string _informationalVersion;
+        private readonly string _entryAssemblyFullName;
 
         internal AssemblyLogMessage(DateTimeOffset now)
             : base(now)
         {
             var assembly = typeof(PluginFactory).Assembly;
+            var entryAssembly = Assembly.GetEntryAssembly();
             var informationalVersionAttribute = assembly.GetCustomAttribute<AssemblyInformationalVersionAttribute>();
             var fileVersionAttribute = assembly.GetCustomAttribute<AssemblyFileVersionAttribute>();
 
             _fullName = assembly.FullName;
+            _entryAssemblyFullName = entryAssembly?.FullName;
 
             if (fileVersionAttribute != null)
             {
@@ -35,7 +38,9 @@ internal AssemblyLogMessage(DateTimeOffset now)
 
         public override string ToString()
         {
-            var message = new JObject(new JProperty("assembly full name", _fullName));
+            var message = new JObject(
+                new JProperty("assembly full name", _fullName),
+                new JProperty("entry assembly full name", _entryAssemblyFullName));
 
             if (!string.IsNullOrEmpty(_fileVersion))
             {
@@ -50,4 +55,4 @@ public override string ToString()
             return ToString("assembly", message);
         }
     }
-}
+}

src/NuGet.Core/NuGet.Protocol/Plugins/Logging/ProcessLogMessage.cs

@@ -11,6 +11,7 @@ internal sealed class ProcessLogMessage : PluginLogMessage
     {
         private readonly int _processId;
         private readonly string _processName;
+        private readonly DateTime _processStartTime;
 
         internal ProcessLogMessage(DateTimeOffset now)
             : base(now)
@@ -19,16 +20,18 @@ internal ProcessLogMessage(DateTimeOffset now)
             {
                 _processId = process.Id;
                 _processName = process.ProcessName;
+                _processStartTime = process.StartTime.ToUniversalTime();
             }
         }
 
         public override string ToString()
         {
             var message = new JObject(
                 new JProperty("process ID", _processId),
-                new JProperty("process name", _processName));
+                new JProperty("process name", _processName),
+                new JProperty("process start time", _processStartTime.ToString("O")));
 
             return ToString("process", message);
         }
     }
-}
+}

src/NuGet.Core/NuGet.Protocol/Plugins/MessageDispatcher.cs

@@ -22,6 +22,7 @@ public sealed class MessageDispatcher : IMessageDispatcher, IResponseHandler
         private readonly ConcurrentDictionary<string, InboundRequestContext> _inboundRequestContexts;
         private readonly IPluginLogger _logger;
         private readonly ConcurrentDictionary<string, OutboundRequestContext> _outboundRequestContexts;
+        private readonly InboundRequestProcessingHandler _inboundRequestProcessingContext;
 
         /// <summary>
         /// Gets the request handlers for use by the dispatcher.
@@ -38,7 +39,7 @@ public sealed class MessageDispatcher : IMessageDispatcher, IResponseHandler
         /// <exception cref="ArgumentNullException">Thrown if <paramref name="idGenerator" />
         /// is <c>null</c>.</exception>
         public MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGenerator)
-            : this(requestHandlers, idGenerator, PluginLogger.DefaultInstance)
+            : this(requestHandlers, idGenerator, new InboundRequestProcessingHandler(), PluginLogger.DefaultInstance)
         {
         }
 
@@ -51,9 +52,11 @@ public MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGenera
         /// is <c>null</c>.</exception>
         /// <exception cref="ArgumentNullException">Thrown if <paramref name="idGenerator" />
         /// is <c>null</c>.</exception>
-        /// <exception cref="ArgumentNullException">Thrown if <paramref name="logger" />
+        /// <exception cref="ArgumentNullException">Thrown if <paramref name="inboundRequestProcessingHandler" />
         /// is <c>null</c>.</exception>
-        internal MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGenerator, IPluginLogger logger)
+        /// /// <exception cref="ArgumentNullException">Thrown if <paramref name="logger" />
+        /// is <c>null</c>.</exception>
+        internal MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGenerator, InboundRequestProcessingHandler inboundRequestProcessingHandler, IPluginLogger logger)
         {
             if (requestHandlers == null)
             {
@@ -65,6 +68,10 @@ internal MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGene
                 throw new ArgumentNullException(nameof(idGenerator));
             }
 
+            if(inboundRequestProcessingHandler == null)
+            {
+                throw new ArgumentNullException(nameof(inboundRequestProcessingHandler));
+            }
             if (logger == null)
             {
                 throw new ArgumentNullException(nameof(logger));
@@ -76,6 +83,7 @@ internal MessageDispatcher(IRequestHandlers requestHandlers, IIdGenerator idGene
 
             _inboundRequestContexts = new ConcurrentDictionary<string, InboundRequestContext>();
             _outboundRequestContexts = new ConcurrentDictionary<string, OutboundRequestContext>();
+            _inboundRequestProcessingContext = inboundRequestProcessingHandler;
         }
 
         /// <summary>
@@ -89,7 +97,7 @@ public void Dispose()
             }
 
             Close();
-
+            _inboundRequestProcessingContext.Dispose();
             SetConnection(connection: null);
 
             // Do not dispose of _logger.  This message dispatcher does not own it.
@@ -645,19 +653,6 @@ private IRequestHandler GetInboundRequestHandler(MessageMethod method)
             return handler;
         }
 
-        private InboundRequestContext GetInboundRequestContext(string requestId)
-        {
-            InboundRequestContext requestContext;
-
-            if (!_inboundRequestContexts.TryGetValue(requestId, out requestContext))
-            {
-                throw new ProtocolException(
-                    string.Format(CultureInfo.CurrentCulture, Strings.Plugin_RequestContextDoesNotExist, requestId));
-            }
-
-            return requestContext;
-        }
-
         private OutboundRequestContext GetOutboundRequestContext(string requestId)
         {
             OutboundRequestContext requestContext;
@@ -699,6 +694,7 @@ private InboundRequestContext CreateInboundRequestContext(
                 _connection,
                 message.RequestId,
                 cancellationToken,
+                _inboundRequestProcessingContext,
                 _logger);
         }
 
@@ -741,4 +737,4 @@ private sealed class NullPayload
         {
         }
     }
-}
+}