Propagate correlation id from gallery to search service

src/NuGet.Services.Search.Client/Client/SearchClient.cs

@@ -7,17 +7,32 @@
 using System.Net;
 using System.Net.Http;
 using System.Threading.Tasks;
+using Correlator.Extensions;
 using Newtonsoft.Json.Linq;
 using NuGet.Services.Client;
+using NuGet.Services.Search.Client.Correlation;
 using NuGet.Services.Search.Models;
 
 namespace NuGet.Services.Search.Client
 {
-    public class SearchClient
+    public class SearchClient : ICorrelated
     {
         private readonly RetryingHttpClientWrapper _retryingHttpClientWrapper;
         private readonly ServiceDiscoveryClient _discoveryClient;
         private readonly string _resourceType;
+        private readonly HttpClient _httpClient;
+
+        private CorrelationIdProvider _correlationIdProvider;
+
+        public CorrelationIdProvider CorrelationIdProvider
+        {
+            get { return _correlationIdProvider; }
+            set
+            {
+                _correlationIdProvider = value;
+                _httpClient.EnsureCorrelationId(_correlationIdProvider.CorrelationId);
+            }
+        }
 
         /// <summary>
         /// Create a search service client from the specified base uri and credentials.
@@ -55,10 +70,11 @@ public SearchClient(Uri baseUri, string resourceType, ICredentials credentials,
                 handler = providedHandler;
             }
 
-            var httpClient = new HttpClient(handler, disposeHandler: true);
+            _httpClient = new HttpClient(handler, disposeHandler: true);
 
-            _retryingHttpClientWrapper = new RetryingHttpClientWrapper(httpClient, healthIndicatorStore);
-            _discoveryClient = new ServiceDiscoveryClient(httpClient, baseUri);
+            _retryingHttpClientWrapper = new RetryingHttpClientWrapper(_httpClient, healthIndicatorStore);
+            _discoveryClient = new ServiceDiscoveryClient(_httpClient, baseUri);
+            CorrelationIdProvider = new CorrelationIdProvider();
         }
 
         private static readonly Dictionary<SortOrder, string> SortNames = new Dictionary<SortOrder, string>
@@ -182,4 +198,4 @@ public async Task<ServiceResponse<JObject>> GetDiagnostics()
                 await _retryingHttpClientWrapper.GetAsync(requestEndpoints));
         }
     }
-}
+}

src/NuGet.Services.Search.Client/Correlation/CorrelationIdProvider.cs

@@ -0,0 +1,29 @@
+// 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.Net.Http;
+using Correlator.Extensions;
+
+namespace NuGet.Services.Search.Client.Correlation
+{
+    public class CorrelationIdProvider
+    {
+        public Guid CorrelationId { get; private set; }
+
+        public CorrelationIdProvider()
+        {
+            CorrelationId = Guid.NewGuid();
+        }
+
+        public CorrelationIdProvider(HttpRequestMessage request)
+        {
+            CorrelationId = request.GetClientCorrelationId();
+
+            if (CorrelationId == Guid.Empty)
+            {
+                CorrelationId = Guid.NewGuid();
+            }
+        }
+    }
+}

src/NuGet.Services.Search.Client/Correlation/ICorrelated.cs

@@ -0,0 +1,13 @@
+// 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.
+
+namespace NuGet.Services.Search.Client.Correlation
+{
+    /// <summary>
+    /// Indicated that the implementing class supports correlation id propogation, and thus excepts CorrelationIdProvider.
+    /// </summary>
+    public interface ICorrelated
+    {
+        CorrelationIdProvider CorrelationIdProvider { get; set; }
+    }
+}

src/NuGet.Services.Search.Client/NuGet.Services.Search.Client.csproj

@@ -32,6 +32,11 @@
     <WarningLevel>4</WarningLevel>
   </PropertyGroup>
   <ItemGroup>
+    <Reference Include="Correlator, Version=1.1.0.0, Culture=neutral, PublicKeyToken=ad71b95c30d52163, processorArchitecture=MSIL">
+      <SpecificVersion>False</SpecificVersion>
+      <HintPath>..\..\packages\WebApiCorrelator.1.1.0.0\lib\net452\Correlator.dll</HintPath>
+      <Private>True</Private>
+    </Reference>
     <Reference Include="Newtonsoft.Json, Version=6.0.0.0, Culture=neutral, PublicKeyToken=30ad4fe6b2a6aeed, processorArchitecture=MSIL">
       <HintPath>..\..\packages\Newtonsoft.Json.6.0.8\lib\net45\Newtonsoft.Json.dll</HintPath>
       <Private>True</Private>
@@ -59,6 +64,11 @@
       <Private>True</Private>
     </Reference>
     <Reference Include="System.Net.Http.WebRequest" />
+    <Reference Include="System.Web.Http, Version=5.1.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL">
+      <SpecificVersion>False</SpecificVersion>
+      <HintPath>..\..\packages\Microsoft.AspNet.WebApi.Core.5.1.2\lib\net45\System.Web.Http.dll</HintPath>
+      <Private>True</Private>
+    </Reference>
     <Reference Include="System.Xml.Linq" />
     <Reference Include="System.Data.DataSetExtensions" />
     <Reference Include="Microsoft.CSharp" />
@@ -75,6 +85,8 @@
     <Compile Include="Client\SearchClient.cs" />
     <Compile Include="Client\ServiceDiscovery.cs" />
     <Compile Include="Client\ServiceDiscoveryClient.cs" />
+    <Compile Include="Correlation\CorrelationIdProvider.cs" />
+    <Compile Include="Correlation\ICorrelated.cs" />
     <Compile Include="Models\SearchResults.cs" />
     <Compile Include="Models\SortOrder.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
@@ -83,6 +95,9 @@
     <None Include="app.config" />
     <None Include="packages.config" />
   </ItemGroup>
+  <ItemGroup>
+    <Content Include="Scripts\jquery.ajax.correlator.min.js" />
+  </ItemGroup>
   <Import Project="$(MSBuildToolsPath)\Microsoft.CSharp.targets" />
   <Import Project="..\..\packages\Microsoft.Bcl.Build.1.0.21\build\Microsoft.Bcl.Build.targets" Condition="Exists('..\..\packages\Microsoft.Bcl.Build.1.0.21\build\Microsoft.Bcl.Build.targets')" />
   <Target Name="EnsureNuGetPackageBuildImports" BeforeTargets="PrepareForBuild">

src/NuGet.Services.Search.Client/packages.config

@@ -1,9 +1,11 @@
 <?xml version="1.0" encoding="utf-8"?>
 <packages>
   <package id="Microsoft.AspNet.WebApi.Client" version="5.2.3" targetFramework="net452" />
+  <package id="Microsoft.AspNet.WebApi.Core" version="5.1.2" targetFramework="net452" />
   <package id="Microsoft.Bcl" version="1.1.10" targetFramework="net452" />
   <package id="Microsoft.Bcl.Build" version="1.0.21" targetFramework="net452" />
   <package id="Microsoft.Net.Http" version="2.2.29" targetFramework="net452" />
   <package id="Newtonsoft.Json" version="6.0.8" targetFramework="net452" />
   <package id="NuGet.Services.Platform.Client" version="3.0.29-r-master" targetFramework="net452" />
+  <package id="WebApiCorrelator" version="1.1.0.0" targetFramework="net452" />
 </packages>

src/NuGetGallery/App_Start/AppActivator.cs

@@ -172,6 +172,15 @@ private static void AppPostStart(IAppConfiguration configuration)
             WebApiConfig.Register(GlobalConfiguration.Configuration);
             NuGetODataConfig.Register(GlobalConfiguration.Configuration);
 
+            // Attach correlator
+            var correlationHandler = new Correlator.Handlers.ClientCorrelationHandler
+            {
+                InitializeIfEmpty = true,
+                TraceCorrelation = true
+            };
+
+            GlobalConfiguration.Configuration.MessageHandlers.Add(correlationHandler);
+
             Routes.RegisterRoutes(RouteTable.Routes, configuration.FeedOnlyMode);
             AreaRegistration.RegisterAllAreas();
 

src/NuGetGallery/Controllers/ODataV2FeedController.cs

@@ -9,21 +9,22 @@
 using System.Web.Http;
 using System.Web.Http.OData;
 using System.Web.Http.OData.Query;
+using NuGet.Frameworks;
+using NuGet.Services.Search.Client.Correlation;
+using NuGet.Versioning;
 using NuGetGallery.Configuration;
 using NuGetGallery.Infrastructure;
+using NuGetGallery.Infrastructure.Lucene;
 using NuGetGallery.OData;
 using NuGetGallery.OData.QueryInterceptors;
 using NuGetGallery.WebApi;
 using QueryInterceptor;
 using WebApi.OutputCache.V2;
-using NuGet.Versioning;
-using NuGet.Frameworks;
-using NuGetGallery.Infrastructure.Lucene;
 
 // ReSharper disable once CheckNamespace
 namespace NuGetGallery.Controllers
 {
-    public class ODataV2FeedController 
+    public class ODataV2FeedController
         : NuGetODataController
     {
         private const int MaxPageSize = SearchAdaptor.MaxPageSize;
@@ -33,7 +34,7 @@ public class ODataV2FeedController
         private readonly ISearchService _searchService;
 
         public ODataV2FeedController(
-            IEntityRepository<Package> packagesRepository, 
+            IEntityRepository<Package> packagesRepository,
             ConfigurationService configurationService,
             ISearchService searchService)
             : base(configurationService)
@@ -62,6 +63,8 @@ public async Task<IHttpActionResult> Get(ODataQueryOptions<V2FeedPackage> option
                 HijackableQueryParameters hijackableQueryParameters = null;
                 if (SearchHijacker.IsHijackable(options, out hijackableQueryParameters) && _searchService is ExternalSearchService)
                 {
+                    SetCorrelation();
+
                     packages = await SearchAdaptor.FindByIdAndVersionCore(
                         _searchService, GetTraditionalHttpContext().Request, packages,
                         hijackableQueryParameters.Id, hijackableQueryParameters.Version, curatedFeed: null);
@@ -131,6 +134,8 @@ private async Task<IHttpActionResult> GetCore(ODataQueryOptions<V2FeedPackage> o
             // try the search service
             try
             {
+                SetCorrelation();
+
                 packages = await SearchAdaptor.FindByIdAndVersionCore(
                     _searchService, GetTraditionalHttpContext().Request, packages, id, version, curatedFeed: null);
 
@@ -169,8 +174,8 @@ public IHttpActionResult GetPropertyFromPackages(string propertyName, string id,
         {
             switch (propertyName.ToLowerInvariant())
             {
-                case "id": return Ok(id);
-                case "version": return Ok(version);
+            case "id": return Ok(id);
+            case "version": return Ok(version);
             }
 
             return BadRequest("Querying property " + propertyName + " is not supported.");
@@ -181,8 +186,8 @@ public IHttpActionResult GetPropertyFromPackages(string propertyName, string id,
         [HttpPost]
         [CacheOutput(ServerTimeSpan = NuGetODataConfig.SearchCacheTimeInSeconds, ClientTimeSpan = NuGetODataConfig.SearchCacheTimeInSeconds)]
         public async Task<IHttpActionResult> Search(
-            ODataQueryOptions<V2FeedPackage> options, 
-            [FromODataUri]string searchTerm = "", 
+            ODataQueryOptions<V2FeedPackage> options,
+            [FromODataUri]string searchTerm = "",
             [FromODataUri]string targetFramework = "",
             [FromODataUri]bool includePrerelease = false)
         {
@@ -202,7 +207,7 @@ public async Task<IHttpActionResult> Search(
                     targetFramework = targetFrameworkList[0];
                 }
             }
-            
+
             // Peform actual search
             var packages = _packagesRepository.GetAll()
                 .Include(p => p.PackageRegistration)
@@ -212,6 +217,8 @@ public async Task<IHttpActionResult> Search(
                 .AsNoTracking();
 
             // todo: search hijack should take options instead of manually parsing query options
+            SetCorrelation();
+
             var query = await SearchAdaptor.SearchCore(
                 _searchService, GetTraditionalHttpContext().Request, packages, searchTerm, targetFramework, includePrerelease, curatedFeed: null);
 
@@ -223,14 +230,14 @@ public async Task<IHttpActionResult> Search(
                 var pagedQueryable = query
                     .Take(options.Top != null ? Math.Min(options.Top.Value, MaxPageSize) : MaxPageSize)
                     .ToV2FeedPackageQuery(GetSiteRoot(), _configurationService.Features.FriendlyLicenses);
-                
+
                 return QueryResult(options, pagedQueryable, MaxPageSize, totalHits, (o, s, resultCount) =>
                 {
                     // The nuget.exe 2.x list command does not like the next link at the bottom when a $top is passed.
                     // Strip it of for backward compatibility.
                     if (o.Top == null || (resultCount.HasValue && o.Top.Value >= resultCount.Value))
                     {
-                        return SearchAdaptor.GetNextLink(Request.RequestUri, resultCount, new {searchTerm, targetFramework, includePrerelease}, o, s);
+                        return SearchAdaptor.GetNextLink(Request.RequestUri, resultCount, new { searchTerm, targetFramework, includePrerelease }, o, s);
                     }
                     return null;
                 });
@@ -245,8 +252,8 @@ public async Task<IHttpActionResult> Search(
         [HttpGet]
         [CacheOutput(ServerTimeSpan = NuGetODataConfig.SearchCacheTimeInSeconds, ClientTimeSpan = NuGetODataConfig.SearchCacheTimeInSeconds)]
         public async Task<IHttpActionResult> SearchCount(
-            ODataQueryOptions<V2FeedPackage> options, 
-            [FromODataUri]string searchTerm = "", 
+            ODataQueryOptions<V2FeedPackage> options,
+            [FromODataUri]string searchTerm = "",
             [FromODataUri]string targetFramework = "",
             [FromODataUri]bool includePrerelease = false)
         {
@@ -261,9 +268,9 @@ public IHttpActionResult GetUpdates(
             ODataQueryOptions<V2FeedPackage> options,
             [FromODataUri]string packageIds,
             [FromODataUri]string versions,
-            [FromODataUri]bool includePrerelease, 
-            [FromODataUri]bool includeAllVersions, 
-            [FromODataUri]string targetFrameworks = "", 
+            [FromODataUri]bool includePrerelease,
+            [FromODataUri]bool includeAllVersions,
+            [FromODataUri]string targetFrameworks = "",
             [FromODataUri]string versionConstraints = "")
         {
             if (string.IsNullOrEmpty(packageIds) || string.IsNullOrEmpty(versions))
@@ -377,5 +384,13 @@ where versionLookup[p.PackageRegistration.Id]
             }
             return updates;
         }
+
+        private void SetCorrelation()
+        {
+            if (_searchService is ICorrelated)
+            {
+                ((ICorrelated)_searchService).CorrelationIdProvider = new CorrelationIdProvider(Request);
+            }
+        }
     }
 }

src/NuGetGallery/Infrastructure/Lucene/ExternalSearchService.cs

@@ -11,12 +11,13 @@
 using System.Web;
 using Newtonsoft.Json.Linq;
 using NuGet.Services.Search.Client;
+using NuGet.Services.Search.Client.Correlation;
 using NuGetGallery.Configuration;
 using NuGetGallery.Diagnostics;
 
 namespace NuGetGallery.Infrastructure.Lucene
 {
-    public class ExternalSearchService : ISearchService, IIndexingService, IRawSearchService
+    public class ExternalSearchService : ISearchService, IIndexingService, IRawSearchService, ICorrelated
     {
         public static readonly string SearchRoundtripTimePerfCounter = "SearchRoundtripTime";
 
@@ -314,5 +315,11 @@ public void RegisterBackgroundJobs(IList<WebBackgrounder.IJob> jobs, IAppConfigu
         {
             // No background jobs to register!
         }
+
+        public CorrelationIdProvider CorrelationIdProvider
+        {
+            get { return _client.CorrelationIdProvider; }
+            set { _client.CorrelationIdProvider = value; }
+        }
     }
 }

src/NuGetGallery/NuGetGallery.csproj

@@ -131,6 +131,11 @@
       <HintPath>..\..\packages\Autofac.WebApi2.3.4.0\lib\net45\Autofac.Integration.WebApi.dll</HintPath>
       <Private>True</Private>
     </Reference>
+    <Reference Include="Correlator, Version=1.1.0.0, Culture=neutral, PublicKeyToken=ad71b95c30d52163, processorArchitecture=MSIL">
+      <SpecificVersion>False</SpecificVersion>
+      <HintPath>..\..\packages\WebApiCorrelator.1.1.0.0\lib\net452\Correlator.dll</HintPath>
+      <Private>True</Private>
+    </Reference>
     <Reference Include="DynamicData.EFCodeFirstProvider, Version=0.3.0.0, Culture=neutral, processorArchitecture=MSIL">
       <SpecificVersion>False</SpecificVersion>
       <HintPath>..\..\packages\DynamicData.EFCodeFirstProvider.0.3.0.0\lib\net40\DynamicData.EFCodeFirstProvider.dll</HintPath>
@@ -1531,6 +1536,7 @@
     <Content Include="Scripts\jquery-1.11.0.min.js" />
     <Content Include="Scripts\jquery-ui-1.10.3.js" />
     <Content Include="Scripts\jquery-ui-1.10.3.min.js" />
+    <Content Include="Scripts\jquery.ajax.correlator.min.js" />
     <Content Include="Scripts\jquery.timeago.js" />
     <None Include="Scripts\jquery.validate-vsdoc.js" />
     <Content Include="Scripts\jquery.validate.js" />

src/NuGetGallery/packages.config

@@ -92,6 +92,7 @@
   <package id="System.Spatial" version="5.6.5-beta" targetFramework="net452" />
   <package id="WebActivator" version="1.4.4" targetFramework="net452" />
   <package id="WebActivatorEx" version="2.0.6" targetFramework="net452" />
+  <package id="WebApiCorrelator" version="1.1.0.0" targetFramework="net452" />
   <package id="WebBackgrounder" version="0.2.0" targetFramework="net452" />
   <package id="WebBackgrounder.EntityFramework" version="0.1" targetFramework="net452" />
   <package id="WebGrease" version="1.1.0" targetFramework="net452" />