changed IDiagnosticUpdateSource to be registration based rather than MEF based.

MEF based approach had race issue where if aggregator created after source, aggregator will miss events fired before aggregator created.

src/EditorFeatures/Core/Implementation/EditAndContinue/EditAndContinueDiagnosticUpdateSource.cs

@@ -13,7 +13,6 @@
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.EditAndContinue
 {
-    [Export(typeof(IDiagnosticUpdateSource))]
     [Export(typeof(EditAndContinueDiagnosticUpdateSource))]
     [Shared]
     internal sealed class EditAndContinueDiagnosticUpdateSource : IDiagnosticUpdateSource
@@ -21,8 +20,10 @@ internal sealed class EditAndContinueDiagnosticUpdateSource : IDiagnosticUpdateS
         internal static object DebuggerErrorId = new object();
         internal static object EmitErrorId = new object();
 
-        public EditAndContinueDiagnosticUpdateSource()
+        [ImportingConstructor]
+        public EditAndContinueDiagnosticUpdateSource(IDiagnosticUpdateSourceRegistrationService registrationService)
         {
+            registrationService.Register(this);
         }
 
         public bool SupportGetDiagnostics { get { return false; } }

src/EditorFeatures/Test/Diagnostics/DiagnosticServiceTests.cs

@@ -24,7 +24,9 @@ public void TestGetDiagnostics1()
                 var document = workspace.CurrentSolution.AddProject("TestProject", "TestProject", LanguageNames.CSharp).AddDocument("TestDocument", string.Empty);
 
                 var source = new TestDiagnosticUpdateSource(false, null);
-                var diagnosticService = new DiagnosticService(SpecializedCollections.SingletonCollection(source), AggregateAsynchronousOperationListener.EmptyListeners);
+                var diagnosticService = new DiagnosticService(AggregateAsynchronousOperationListener.EmptyListeners);
+                diagnosticService.Register(source);
+
                 diagnosticService.DiagnosticsUpdated += (s, o) => { set.Set(); };
 
                 var id = Tuple.Create(workspace, document);
@@ -54,7 +56,9 @@ public void TestGetDiagnostics2()
                 var document2 = document.Project.AddDocument("TestDocument2", string.Empty);
 
                 var source = new TestDiagnosticUpdateSource(false, null);
-                var diagnosticService = new DiagnosticService(SpecializedCollections.SingletonCollection(source), AggregateAsynchronousOperationListener.EmptyListeners);
+                var diagnosticService = new DiagnosticService(AggregateAsynchronousOperationListener.EmptyListeners);
+                diagnosticService.Register(source);
+
                 diagnosticService.DiagnosticsUpdated += (s, o) => { set.Set(); };
 
                 var id = Tuple.Create(workspace, document);

src/EditorFeatures/Test/Diagnostics/DiagnosticsSquiggleTaggerProviderTests.cs

@@ -69,7 +69,8 @@ private DiagnosticTaggerWrapper(TestWorkspace workspace, DiagnosticAnalyzerServi
                 ValueTuple.Create(FeatureAttribute.ErrorSquiggles, asyncListener));
 
             this.analyzerService = analyzerService;
-            var diagnosticService = new DiagnosticService(SpecializedCollections.SingletonEnumerable(updateSource), listeners);
+            var diagnosticService = new DiagnosticService(listeners);
+            diagnosticService.Register(updateSource);
 
             this.TaggerProvider = new DiagnosticsSquiggleTaggerProvider(
                 workspace.Services.GetService<IOptionService>(), diagnosticService,

src/EditorFeatures/Test/Diagnostics/MockDiagnosticUpdateSourceRegistrationService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using Microsoft.CodeAnalysis.Diagnostics;
+
+namespace Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics
+{
+    internal class MockDiagnosticUpdateSourceRegistrationService : IDiagnosticUpdateSourceRegistrationService
+    {
+        public void Register(IDiagnosticUpdateSource source)
+        {
+            // do nothing
+        }
+    }
+}
\ No newline at end of file

src/EditorFeatures/Test/Diagnostics/TestDiagnosticAnalyzerService.cs

@@ -4,6 +4,7 @@
 using System.Collections.Immutable;
 using System.Linq;
 using Microsoft.CodeAnalysis.Diagnostics.Log;
+using Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Diagnostics
@@ -33,13 +34,13 @@ internal TestDiagnosticAnalyzerService(ImmutableArray<AnalyzerReference> hostAna
         }
 
         private TestDiagnosticAnalyzerService(HostAnalyzerManager hostAnalyzerManager, AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource, Action<Exception, DiagnosticAnalyzer, Diagnostic> onAnalyzerException)
-            : base(hostAnalyzerManager, hostDiagnosticUpdateSource)
+            : base(hostAnalyzerManager, hostDiagnosticUpdateSource, new MockDiagnosticUpdateSourceRegistrationService())
         {
             _onAnalyzerException = onAnalyzerException;
         }
 
         internal TestDiagnosticAnalyzerService(AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource = null, Action<Exception, DiagnosticAnalyzer, Diagnostic> onAnalyzerException = null)
-           : base(SpecializedCollections.EmptyEnumerable<HostDiagnosticAnalyzerPackage>(), hostDiagnosticUpdateSource)
+           : base(SpecializedCollections.EmptyEnumerable<HostDiagnosticAnalyzerPackage>(), hostDiagnosticUpdateSource, new MockDiagnosticUpdateSourceRegistrationService())
         {
             _onAnalyzerException = onAnalyzerException;
         }

src/EditorFeatures/Test/EditorServicesTest.csproj

@@ -222,6 +222,7 @@
     <Compile Include="Diagnostics\AbstractSuppressionDiagnosticTest.cs" />
     <Compile Include="Diagnostics\AbstractDiagnosticProviderBasedUserDiagnosticTest.cs" />
     <Compile Include="Diagnostics\AbstractUserDiagnosticTest.cs" />
+    <Compile Include="Diagnostics\MockDiagnosticUpdateSourceRegistrationService.cs" />
     <Compile Include="Diagnostics\TestDiagnosticAnalyzerDriver.cs" />
     <Compile Include="Diagnostics\TestDiagnosticAnalyzerService.cs" />
     <Compile Include="Diagnostics\DiagnosticDataTests.cs" />

src/EditorFeatures/TestUtilities/ExceptionUtilities.cs

 // Copyright (c) Microsoft.  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.Diagnostics;
-using System.Linq;
 using System.Runtime.InteropServices;
-using System.Text;
-using System.Threading.Tasks;
 
 namespace Microsoft.CodeAnalysis.Test.Utilities
 {

src/EditorFeatures/TestUtilities/ServicesTestUtilities.csproj

@@ -88,4 +88,4 @@
     <Import Project="..\..\..\build\Targets\VSL.Imports.targets" />
     <Import Project="..\..\..\build\Targets\Roslyn.Toolsets.Xunit.targets" />
   </ImportGroup>
-</Project>
+</Project>
\ No newline at end of file

src/Features/Core/Portable/Diagnostics/DiagnosticAnalyzerService.cs

@@ -24,25 +24,34 @@ internal partial class DiagnosticAnalyzerService : IDiagnosticAnalyzerService
 
         [ImportingConstructor]
         public DiagnosticAnalyzerService(
+            IDiagnosticUpdateSourceRegistrationService registrationService,
             [ImportMany] IEnumerable<Lazy<IAsynchronousOperationListener, FeatureMetadata>> asyncListeners,
             [Import(AllowDefault = true)]IWorkspaceDiagnosticAnalyzerProviderService diagnosticAnalyzerProviderService = null,
             [Import(AllowDefault = true)]AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource = null)
             : this(diagnosticAnalyzerProviderService != null ? diagnosticAnalyzerProviderService.GetHostDiagnosticAnalyzerPackages() : SpecializedCollections.EmptyEnumerable<HostDiagnosticAnalyzerPackage>(),
-                   hostDiagnosticUpdateSource,
-                   new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.DiagnosticService))
+                hostDiagnosticUpdateSource,
+                registrationService, new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.DiagnosticService))
         {
         }
 
         public IAsynchronousOperationListener Listener => _listener;
 
         // protected for testing purposes.
-        protected DiagnosticAnalyzerService(IEnumerable<HostDiagnosticAnalyzerPackage> workspaceAnalyzerPackages, AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource, IAsynchronousOperationListener listener = null)
-            : this (new HostAnalyzerManager(workspaceAnalyzerPackages, hostDiagnosticUpdateSource), hostDiagnosticUpdateSource, listener)
+        protected DiagnosticAnalyzerService(
+            IEnumerable<HostDiagnosticAnalyzerPackage> workspaceAnalyzerPackages,
+            AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource,
+            IDiagnosticUpdateSourceRegistrationService registrationService,
+            IAsynchronousOperationListener listener = null)
+            : this(new HostAnalyzerManager(workspaceAnalyzerPackages, hostDiagnosticUpdateSource), hostDiagnosticUpdateSource, registrationService, listener)
         {
         }
 
         // protected for testing purposes.
-        protected DiagnosticAnalyzerService(HostAnalyzerManager hostAnalyzerManager, AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource = null, IAsynchronousOperationListener listener = null) : this()
+        protected DiagnosticAnalyzerService(
+            HostAnalyzerManager hostAnalyzerManager,
+            AbstractHostDiagnosticUpdateSource hostDiagnosticUpdateSource,
+            IDiagnosticUpdateSourceRegistrationService registrationService,
+            IAsynchronousOperationListener listener = null) : this(registrationService)
         {
             _hostAnalyzerManager = hostAnalyzerManager;
             _hostDiagnosticUpdateSource = hostDiagnosticUpdateSource;

src/Features/Core/Portable/Diagnostics/DiagnosticAnalyzerService_UpdateSource.cs

@@ -2,17 +2,20 @@
 
 using System;
 using System.Collections.Immutable;
-using System.Composition;
 using System.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Diagnostics
 {
-    [Export(typeof(IDiagnosticUpdateSource))]
     internal partial class DiagnosticAnalyzerService : IDiagnosticUpdateSource
     {
         public event EventHandler<DiagnosticsUpdatedArgs> DiagnosticsUpdated;
 
+        public DiagnosticAnalyzerService(IDiagnosticUpdateSourceRegistrationService registrationService) : this()
+        {
+            registrationService.Register(this);
+        }
+
         internal void RaiseDiagnosticsUpdated(object sender, DiagnosticsUpdatedArgs state)
         {
             var handler = this.DiagnosticsUpdated;

src/Features/Core/Portable/Diagnostics/DiagnosticService.cs

@@ -5,7 +5,6 @@
 using System.Collections.Immutable;
 using System.Composition;
 using System.Diagnostics;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
@@ -14,35 +13,28 @@
 namespace Microsoft.CodeAnalysis.Diagnostics
 {
     [Export(typeof(IDiagnosticService)), Shared]
-    internal class DiagnosticService : IDiagnosticService
+    internal partial class DiagnosticService : IDiagnosticService
     {
         private const string DiagnosticsUpdatedEventName = "DiagnosticsUpdated";
 
         private readonly IAsynchronousOperationListener _listener;
         private readonly EventMap _eventMap;
         private readonly SimpleTaskQueue _eventQueue;
-        private readonly ImmutableArray<IDiagnosticUpdateSource> _updateSources;
 
         private readonly object _gate;
         private readonly Dictionary<IDiagnosticUpdateSource, Dictionary<object, Data>> _map;
 
         [ImportingConstructor]
-        public DiagnosticService(
-            [ImportMany] IEnumerable<IDiagnosticUpdateSource> diagnosticUpdateSource,
-            [ImportMany] IEnumerable<Lazy<IAsynchronousOperationListener, FeatureMetadata>> asyncListeners)
+        public DiagnosticService([ImportMany] IEnumerable<Lazy<IAsynchronousOperationListener, FeatureMetadata>> asyncListeners) : this()
         {
             // queue to serialize events.
             _eventMap = new EventMap();
             _eventQueue = new SimpleTaskQueue(TaskScheduler.Default);
 
-            _updateSources = diagnosticUpdateSource.AsImmutable();
             _listener = new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.DiagnosticService);
 
             _gate = new object();
             _map = new Dictionary<IDiagnosticUpdateSource, Dictionary<object, Data>>();
-
-            // connect each diagnostic update source to events
-            ConnectDiagnosticsUpdatedEvents();
         }
 
         public event EventHandler<DiagnosticsUpdatedArgs> DiagnosticsUpdated
@@ -80,11 +72,11 @@ private void UpdateDataMap(object sender, DiagnosticsUpdatedArgs args)
                 return;
             }
 
-            Contract.Requires(_updateSources.IndexOf(updateSource) >= 0);
-
             // we expect someone who uses this ability to small.
             lock (_gate)
             {
+                Contract.Requires(_updateSources.Contains(updateSource));
+
                 var list = _map.GetOrAdd(updateSource, _ => new Dictionary<object, Data>());
                 var data = new Data(args);
 
@@ -99,14 +91,6 @@ private void UpdateDataMap(object sender, DiagnosticsUpdatedArgs args)
             }
         }
 
-        private void ConnectDiagnosticsUpdatedEvents()
-        {
-            foreach (var source in _updateSources)
-            {
-                source.DiagnosticsUpdated += OnDiagnosticsUpdated;
-            }
-        }
-
         private void OnDiagnosticsUpdated(object sender, DiagnosticsUpdatedArgs e)
         {
             AssertIfNull(e.Diagnostics);

src/Features/Core/Portable/Diagnostics/DiagnosticService_UpdateSourceRegistrationService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Collections.Immutable;
+using System.Composition;
+
+namespace Microsoft.CodeAnalysis.Diagnostics
+{
+    [Export(typeof(IDiagnosticUpdateSourceRegistrationService))]
+    internal partial class DiagnosticService : IDiagnosticUpdateSourceRegistrationService
+    {
+        private ImmutableHashSet<IDiagnosticUpdateSource> _updateSources;
+
+        public DiagnosticService()
+        {
+            // we use registry service rather than doing MEF import since MEF import method can have race issue where
+            // update source gets created before aggregator - diagnostic service - is created and we will lose events fired before
+            // the aggregator is created.
+            _updateSources = ImmutableHashSet<IDiagnosticUpdateSource>.Empty;
+        }
+
+        public void Register(IDiagnosticUpdateSource source)
+        {
+            lock (_gate)
+            {
+                if (_updateSources.Contains(source))
+                {
+                    return;
+                }
+
+                _updateSources = _updateSources.Add(source);
+                source.DiagnosticsUpdated += OnDiagnosticsUpdated;
+            }
+        }
+    }
+}

src/Features/Core/Portable/Diagnostics/IDiagnosticUpdateSourceRegistrationService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.CodeAnalysis.Diagnostics
+{
+    /// <summary>
+    /// A service that let people to register new IDiagnosticUpdateSource
+    /// </summary>
+    internal interface IDiagnosticUpdateSourceRegistrationService
+    {
+        /// <summary>
+        /// Register new IDiagnosticUpdateSource
+        /// 
+        /// Currently, it doesn't support unregister since our event is asynchronous and unregistering source that deal with asynchronous event is not straight forward.
+        /// </summary>
+        void Register(IDiagnosticUpdateSource source);
+    }
+}

src/Features/Core/Portable/Features.csproj

@@ -161,8 +161,10 @@
     <Compile Include="Diagnostics\AnalyzerHelper.cs" />
     <Compile Include="Diagnostics\AbstractHostDiagnosticUpdateSource.cs" />
     <Compile Include="Diagnostics\AnalyzerUpdateArgsId.cs" />
+    <Compile Include="Diagnostics\DiagnosticService_UpdateSourceRegistrationService.cs" />
     <Compile Include="Diagnostics\Extensions.cs" />
     <Compile Include="Diagnostics\HostDiagnosticAnalyzerPackage.cs" />
+    <Compile Include="Diagnostics\IDiagnosticUpdateSourceRegistrationService.cs" />
     <Compile Include="Diagnostics\InternalDiagnosticsOptionsProvider.cs" />
     <Compile Include="Diagnostics\PredefinedBuildTools.cs" />
     <Compile Include="Diagnostics\DiagnosticAnalyzerService_BuildSynchronization.cs" />

src/VisualStudio/Core/Def/Implementation/Diagnostics/MiscellaneousDiagnosticAnalyzerService.cs

@@ -17,16 +17,17 @@
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.Diagnostics
 {
     [ExportIncrementalAnalyzerProvider(WorkspaceKind.MiscellaneousFiles)]
-    [Export(typeof(IDiagnosticUpdateSource))]
     [Shared]
     internal partial class MiscellaneousDiagnosticAnalyzerService : IIncrementalAnalyzerProvider, IDiagnosticUpdateSource
     {
         private readonly IDiagnosticAnalyzerService _analyzerService;
 
         [ImportingConstructor]
-        public MiscellaneousDiagnosticAnalyzerService(IDiagnosticAnalyzerService analyzerService)
+        public MiscellaneousDiagnosticAnalyzerService(IDiagnosticAnalyzerService analyzerService, IDiagnosticUpdateSourceRegistrationService registrationService)
         {
             _analyzerService = analyzerService;
+
+            registrationService.Register(this);
         }
 
         public IIncrementalAnalyzer CreateIncrementalAnalyzer(Workspace workspace)

src/VisualStudio/Core/Def/Implementation/TaskList/ExternalErrorDiagnosticUpdateSource.cs

@@ -18,7 +18,6 @@
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.TaskList
 {
-    [Export(typeof(IDiagnosticUpdateSource))]
     [Export(typeof(ExternalErrorDiagnosticUpdateSource))]
     internal class ExternalErrorDiagnosticUpdateSource : IDiagnosticUpdateSource
     {
@@ -36,8 +35,9 @@ internal class ExternalErrorDiagnosticUpdateSource : IDiagnosticUpdateSource
         public ExternalErrorDiagnosticUpdateSource(
             VisualStudioWorkspaceImpl workspace,
             IDiagnosticAnalyzerService diagnosticService,
+            IDiagnosticUpdateSourceRegistrationService registrationService,
             [ImportMany] IEnumerable<Lazy<IAsynchronousOperationListener, FeatureMetadata>> asyncListeners) :
-                this(workspace, diagnosticService, new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.ErrorList))
+                this(workspace, diagnosticService, registrationService, new AggregateAsynchronousOperationListener(asyncListeners, FeatureAttribute.ErrorList))
         {
             Contract.Requires(!KnownUIContexts.SolutionBuildingContext.IsActive);
             KnownUIContexts.SolutionBuildingContext.UIContextChanged += OnSolutionBuild;
@@ -49,6 +49,7 @@ internal class ExternalErrorDiagnosticUpdateSource : IDiagnosticUpdateSource
         internal ExternalErrorDiagnosticUpdateSource(
             Workspace workspace,
             IDiagnosticAnalyzerService diagnosticService,
+            IDiagnosticUpdateSourceRegistrationService registrationService,
             IAsynchronousOperationListener listener)
         {
             // use queue to serialize work. no lock needed
@@ -61,6 +62,8 @@ internal class ExternalErrorDiagnosticUpdateSource : IDiagnosticUpdateSource
             _diagnosticService = diagnosticService;
 
             _notificationService = _workspace.Services.GetService<IGlobalOperationNotificationService>();
+
+            registrationService.Register(this);
         }
 
         public event EventHandler<bool> BuildStarted;

src/VisualStudio/Core/Def/Implementation/TaskList/HostDiagnosticUpdateSource.cs

@@ -12,7 +12,6 @@ namespace Microsoft.VisualStudio.LanguageServices.Implementation.TaskList
 {
     // exporting both Abstract and HostDiagnosticUpdateSource is just to make testing easier.
     // use HostDiagnosticUpdateSource when abstract one is not needed for testing purpose
-    [Export(typeof(IDiagnosticUpdateSource))]
     [Export(typeof(AbstractHostDiagnosticUpdateSource))]
     [Export(typeof(HostDiagnosticUpdateSource))]
     internal sealed class HostDiagnosticUpdateSource : AbstractHostDiagnosticUpdateSource
@@ -23,9 +22,11 @@ internal sealed class HostDiagnosticUpdateSource : AbstractHostDiagnosticUpdateS
         private readonly Dictionary<ProjectId, HashSet<object>> _diagnosticMap = new Dictionary<ProjectId, HashSet<object>>();
 
         [ImportingConstructor]
-        public HostDiagnosticUpdateSource(VisualStudioWorkspaceImpl workspace)
+        public HostDiagnosticUpdateSource(VisualStudioWorkspaceImpl workspace, IDiagnosticUpdateSourceRegistrationService registrationService)
         {
             _workspace = workspace;
+
+            registrationService.Register(this);
         }
 
         internal override Workspace Workspace

src/VisualStudio/Core/Test/Diagnostics/ExternalDiagnosticUpdateSourceTests.vb

@@ -6,6 +6,7 @@ Imports System.Threading.Tasks
 Imports Microsoft.CodeAnalysis
 Imports Microsoft.CodeAnalysis.Diagnostics
 Imports Microsoft.CodeAnalysis.Editor.Implementation.Diagnostics
+Imports Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Shared.TestHooks
 Imports Microsoft.CodeAnalysis.Text
@@ -20,7 +21,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Diagnostics
             Using workspace = CSharpWorkspaceFactory.CreateWorkspaceFromLines(String.Empty)
                 Dim waiter = New Waiter()
                 Dim service = New TestDiagnosticAnalyzerService()
-                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, waiter)
+                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, New MockDiagnosticUpdateSourceRegistrationService(), waiter)
 
                 Assert.False(source.SupportGetDiagnostics)
             End Using
@@ -31,7 +32,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Diagnostics
             Using workspace = CSharpWorkspaceFactory.CreateWorkspaceFromLines(String.Empty)
                 Dim waiter = New Waiter()
                 Dim service = New TestDiagnosticAnalyzerService()
-                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, waiter)
+                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, New MockDiagnosticUpdateSourceRegistrationService(), waiter)
 
                 Dim project = workspace.CurrentSolution.Projects.First()
                 Dim diagnostic = GetDiagnosticData(workspace, project.Id)
@@ -63,7 +64,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Diagnostics
                 Dim diagnostic = GetDiagnosticData(workspace, project.Id)
 
                 Dim service = New TestDiagnosticAnalyzerService(ImmutableArray.Create(Of DiagnosticData)(diagnostic))
-                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, waiter)
+                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, New MockDiagnosticUpdateSourceRegistrationService(), waiter)
 
                 Dim map = New Dictionary(Of DocumentId, HashSet(Of DiagnosticData))()
                 map.Add(project.DocumentIds.First(), New HashSet(Of DiagnosticData)(
@@ -88,7 +89,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Diagnostics
                 Dim diagnostic = GetDiagnosticData(workspace, project.Id)
 
                 Dim service = New TestDiagnosticAnalyzerService(ImmutableArray(Of DiagnosticData).Empty)
-                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, waiter)
+                Dim source = New ExternalErrorDiagnosticUpdateSource(workspace, service, New MockDiagnosticUpdateSourceRegistrationService(), waiter)
                 AddHandler source.BuildStarted, Sub(o, started)
                                                     If Not started Then
                                                         Assert.Equal(2, source.GetBuildErrors().Length)

src/VisualStudio/Core/Test/Diagnostics/MiscDiagnosticUpdateSourceTests.vb

@@ -6,14 +6,12 @@ Imports Microsoft.CodeAnalysis
 Imports Microsoft.CodeAnalysis.Diagnostics
 Imports Microsoft.CodeAnalysis.Editor
 Imports Microsoft.CodeAnalysis.Editor.Implementation.Diagnostics
-Imports Microsoft.CodeAnalysis.Editor.Shared.Tagging
-Imports Microsoft.CodeAnalysis.Editor.UnitTests
+Imports Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics
 Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
 Imports Microsoft.CodeAnalysis.Options
 Imports Microsoft.CodeAnalysis.Shared.TestHooks
 Imports Microsoft.CodeAnalysis.Text.Shared.Extensions
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.Diagnostics
-Imports Microsoft.VisualStudio.Text
 Imports Microsoft.VisualStudio.Text.Tagging
 Imports Roslyn.Test.Utilities
 Imports Roslyn.Utilities
@@ -26,7 +24,10 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.Diagnostics
 class 123 { }
                        </code>
             Using workspace = CSharpWorkspaceFactory.CreateWorkspaceFromLines(code.ToString())
-                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()))
+                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(
+                    New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()),
+                    New MockDiagnosticUpdateSourceRegistrationService())
+
                 Assert.False(miscService.SupportGetDiagnostics)
 
                 Dim listener = New AsynchronousOperationListener()
@@ -34,7 +35,8 @@ class 123 { }
                     ValueTuple.Create(FeatureAttribute.DiagnosticService, listener),
                     ValueTuple.Create(FeatureAttribute.ErrorSquiggles, listener))
 
-                Dim diagnosticService = New DiagnosticService(New IDiagnosticUpdateSource() {miscService}, listeners)
+                Dim diagnosticService = New DiagnosticService(listeners)
+                diagnosticService.Register(miscService)
 
                 Dim optionsService = workspace.Services.GetService(Of IOptionService)()
 
@@ -64,7 +66,10 @@ class 123 { }
 class 123 { }
                        </code>
             Using workspace = CSharpWorkspaceFactory.CreateWorkspaceFromLines(code.ToString())
-                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()))
+                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(
+                    New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()),
+                    New MockDiagnosticUpdateSourceRegistrationService())
+
                 Dim buildTool = String.Empty
 
                 AddHandler miscService.DiagnosticsUpdated, Sub(e, a)
@@ -86,7 +91,10 @@ Class 123
 End Class
                        </code>
             Using workspace = VisualBasicWorkspaceFactory.CreateWorkspaceFromLines(code.ToString())
-                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()))
+                Dim miscService = New MiscellaneousDiagnosticAnalyzerService(
+                    New TestDiagnosticAnalyzerService(DiagnosticExtensions.GetCompilerDiagnosticAnalyzersMap()),
+                    New MockDiagnosticUpdateSourceRegistrationService())
+
                 Dim buildTool = String.Empty
 
                 AddHandler miscService.DiagnosticsUpdated, Sub(e, a)

src/VisualStudio/Core/Test/ProjectSystemShim/VisualStudioAnalyzerTests.vb

@@ -4,7 +4,7 @@ Imports System.IO
 Imports System.Reflection
 Imports Microsoft.CodeAnalysis
 Imports Microsoft.CodeAnalysis.Diagnostics
-Imports Microsoft.CodeAnalysis.Host
+Imports Microsoft.CodeAnalysis.Editor.UnitTests.Diagnostics
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.ProjectSystem
 Imports Microsoft.VisualStudio.LanguageServices.Implementation.TaskList
 Imports Microsoft.VisualStudio.LanguageServices.UnitTests.ProjectSystemShim.Framework
@@ -24,7 +24,7 @@ Namespace Microsoft.VisualStudio.LanguageServices.UnitTests.ProjectSystemShim
 
         <Fact, Trait(Traits.Feature, Traits.Features.Diagnostics)>
         Public Sub AnalyzerErrorsAreUpdated()
-            Dim hostDiagnosticUpdateSource = New HostDiagnosticUpdateSource(Nothing)
+            Dim hostDiagnosticUpdateSource = New HostDiagnosticUpdateSource(Nothing, New MockDiagnosticUpdateSourceRegistrationService())
 
             Dim file = Path.GetTempFileName()
             Dim eventHandler = New EventHandlers(file)