Simplifies IdKeyMap interface, moves the weird add mapping logic in PublishedSnapshotService to a composer.

src/Umbraco.Abstractions/Services/IIdKeyMap.cs

@@ -5,7 +5,6 @@ namespace Umbraco.Core.Services
 {
     public interface IIdKeyMap
     {
-        void SetMapper(UmbracoObjectTypes umbracoObjectType, Func<int, Guid> id2key, Func<Guid, int> key2id);
         Attempt<int> GetIdForKey(Guid key, UmbracoObjectTypes umbracoObjectType);
         Attempt<int> GetIdForUdi(Udi udi);
         Attempt<Udi> GetUdiForId(int id, UmbracoObjectTypes umbracoObjectType);

src/Umbraco.Tests/PublishedContent/NuCacheChildrenTests.cs

@@ -150,7 +150,6 @@ namespace Umbraco.Tests.PublishedContent
                 runtime,
                 serviceContext,
                 contentTypeFactory,
-                null,
                 _snapshotAccessor,
                 _variationAccesor,
                 Mock.Of<IProfilingLogger>(),

src/Umbraco.Tests/PublishedContent/NuCacheTests.cs

@@ -192,7 +192,6 @@ namespace Umbraco.Tests.PublishedContent
                 runtime,
                 serviceContext,
                 contentTypeFactory,
-                null,
                 new TestPublishedSnapshotAccessor(),
                 _variationAccesor,
                 Mock.Of<IProfilingLogger>(),

src/Umbraco.Tests/Scoping/ScopedNuCacheTests.cs

@@ -92,7 +92,6 @@ namespace Umbraco.Tests.Scoping
                 runtimeStateMock.Object,
                 ServiceContext,
                 contentTypeFactory,
-                null,
                 publishedSnapshotAccessor,
                 Mock.Of<IVariationContextAccessor>(),
                 ProfilingLogger,

src/Umbraco.Tests/Services/ContentTypeServiceVariantsTests.cs

@@ -65,7 +65,6 @@ namespace Umbraco.Tests.Services
                 runtimeStateMock.Object,
                 ServiceContext,
                 contentTypeFactory,
-                null,
                 publishedSnapshotAccessor,
                 Mock.Of<IVariationContextAccessor>(),
                 ProfilingLogger,

src/Umbraco.Web/PublishedCache/NuCache/NuCacheComposer.cs

@@ -1,5 +1,8 @@
 using Umbraco.Core;
 using Umbraco.Core.Composing;
+using Umbraco.Core.Models;
+using Umbraco.Core.Scoping;
+using Umbraco.Core.Services;
 using Umbraco.Web.PublishedCache.NuCache.DataSource;
 
 namespace Umbraco.Web.PublishedCache.NuCache
@@ -18,6 +21,20 @@ namespace Umbraco.Web.PublishedCache.NuCache
             composition.Register(factory => new PublishedSnapshotServiceOptions());
             composition.SetPublishedSnapshotService<PublishedSnapshotService>();
 
+            // replace this service since we want to improve the content/media
+            // mapping lookups if we are using nucache.
+            composition.RegisterUnique<IIdKeyMap>(factory =>
+            {
+                var idkSvc = new IdKeyMap(factory.GetInstance<IScopeProvider>());
+                var publishedSnapshotService = factory.GetInstance<IPublishedSnapshotService>() as PublishedSnapshotService;
+                if (publishedSnapshotService != null)
+                {
+                    idkSvc.SetMapper(UmbracoObjectTypes.Document, id => publishedSnapshotService.GetDocumentUid(id), uid => publishedSnapshotService.GetDocumentId(uid));
+                    idkSvc.SetMapper(UmbracoObjectTypes.Media, id => publishedSnapshotService.GetMediaUid(id), uid => publishedSnapshotService.GetMediaId(uid));
+                }                
+                return idkSvc;
+            });
+
             // add the NuCache health check (hidden from type finder)
             // TODO: no NuCache health check yet
             //composition.HealthChecks().Add<NuCacheIntegrityHealthCheck>();

src/Umbraco.Web/PublishedCache/NuCache/PublishedSnapshotService.cs

@@ -76,7 +76,7 @@ namespace Umbraco.Web.PublishedCache.NuCache
         //private static int _singletonCheck;
 
         public PublishedSnapshotService(PublishedSnapshotServiceOptions options, IMainDom mainDom, IRuntimeState runtime,
-            ServiceContext serviceContext, IPublishedContentTypeFactory publishedContentTypeFactory, IIdKeyMap idKeyMap,
+            ServiceContext serviceContext, IPublishedContentTypeFactory publishedContentTypeFactory,
             IPublishedSnapshotAccessor publishedSnapshotAccessor, IVariationContextAccessor variationContextAccessor, IProfilingLogger logger, IScopeProvider scopeProvider,
             IDocumentRepository documentRepository, IMediaRepository mediaRepository, IMemberRepository memberRepository,
             IDefaultCultureAccessor defaultCultureAccessor,
@@ -156,17 +156,22 @@ namespace Umbraco.Web.PublishedCache.NuCache
 
                 LoadCachesOnStartup();
             }
-
-            Guid GetUid(ContentStore store, int id) => store.LiveSnapshot.Get(id)?.Uid ?? default;
-            int GetId(ContentStore store, Guid uid) => store.LiveSnapshot.Get(uid)?.Id ?? default;
-
-            if (idKeyMap != null)
-            {
-                idKeyMap.SetMapper(UmbracoObjectTypes.Document, id => GetUid(_contentStore, id), uid => GetId(_contentStore, uid));
-                idKeyMap.SetMapper(UmbracoObjectTypes.Media, id => GetUid(_mediaStore, id), uid => GetId(_mediaStore, uid));
-            }
         }
 
+        #region Id <-> Key methods
+
+        // NOTE: These aren't used within this object but are made available internally to improve the IdKey lookup performance
+        // when nucache is enabled.
+
+        internal int GetDocumentId(Guid udi) => GetId(_contentStore, udi);
+        internal int GetMediaId(Guid udi) => GetId(_mediaStore, udi);
+        internal Guid GetDocumentUid(int id) => GetUid(_contentStore, id);
+        internal Guid GetMediaUid(int id) => GetUid(_mediaStore, id);
+        private int GetId(ContentStore store, Guid uid) => store.LiveSnapshot.Get(uid)?.Id ?? default;
+        private Guid GetUid(ContentStore store, int id) => store.LiveSnapshot.Get(id)?.Uid ?? default;
+
+        #endregion
+
         /// <summary>
         /// Install phase of <see cref="IMainDom"/>
         /// </summary>

src/Umbraco.Web/Runtime/WebInitialComposer.cs

@@ -143,7 +143,7 @@ namespace Umbraco.Web.Runtime
             composition.RegisterUnique<IEventMessagesFactory, DefaultEventMessagesFactory>();
             composition.RegisterUnique<IEventMessagesAccessor, HybridEventMessagesAccessor>();
             composition.RegisterUnique<ITreeService, TreeService>();
-            composition.RegisterUnique<ISectionService, SectionService>();
+            composition.RegisterUnique<ISectionService, SectionService>();            
 
             composition.RegisterUnique<IDashboardService, DashboardService>();