Fix crasher

Author: shana

src/GitHub.App/Controllers/UIController.cs

@@ -152,7 +152,7 @@ public IObservable<bool> ListenToCompletionState()
 
         void End(bool success)
         {
-            uiProvider.RemoveService(typeof(IConnection));
+            uiProvider.RemoveService(typeof(IConnection), this);
             completion?.OnNext(success);
             completion?.OnCompleted();
             transition.OnCompleted();
@@ -228,7 +228,7 @@ public void Start([AllowNull] IConnection connection)
             if (connection != null)
             {
                 if (currentFlow != UIControllerFlow.Authentication)
-                    uiProvider.AddService(connection);
+                    uiProvider.AddService(this, connection);
                 else // sanity check: it makes zero sense to pass a connection in when calling the auth flow
                     Debug.Assert(false, "Calling the auth flow with a connection makes no sense!");
 
@@ -263,15 +263,15 @@ public void Start([AllowNull] IConnection connection)
                         if (currentFlow != UIControllerFlow.Authentication)
                         {
                             if (loggedin) // register the first available connection so the viewmodel can use it
-                                uiProvider.AddService(c);
+                                uiProvider.AddService(this, c);
                             else
                             {
                                 // a connection will be added to the list when auth is done, register it so the next
                                 // viewmodel can use it
                                 connectionAdded = (s, e) =>
                                 {
                                     if (e.Action == NotifyCollectionChangedAction.Add)
-                                        uiProvider.AddService(typeof(IConnection), e.NewItems[0]);
+                                        uiProvider.AddService(typeof(IConnection), this, e.NewItems[0]);
                                 };
                                 connectionManager.Connections.CollectionChanged += connectionAdded;
                             }

src/GitHub.Exports/Services/IUIProvider.cs

@@ -18,9 +18,15 @@ public interface IUIProvider : IServiceProvider
         object TryGetService(string typename);
         T TryGetService<T>() where T : class;
 
-        void AddService(Type t, object instance);
-        void AddService<T>(T instance);
-        void RemoveService(Type t);
+        void AddService(Type t, object owner, object instance);
+        void AddService<T>(object owner, T instance);
+        /// <summary>
+        /// Removes a service from the catalog
+        /// </summary>
+        /// <param name="t">The type we want to remove</param>
+        /// <param name="owner">The owner, which either has to match what was passed to AddService,
+        /// or if it's null, the service will be removed without checking for ownership</param>
+        void RemoveService(Type t, object owner);
 
         IObservable<UserControl> SetupUI(UIControllerFlow controllerFlow, IConnection connection);
         void RunUI();

src/GitHub.VisualStudio/Services/UIProvider.cs

@@ -26,11 +26,17 @@ namespace GitHub.VisualStudio
     [PartCreationPolicy(CreationPolicy.Shared)]
     public class UIProvider : IUIProvider, IDisposable
     {
+        class OwnedComposablePart
+        {
+            public object Owner { get; set; }
+            public ComposablePart Part { get; set; }
+        }
+
         static readonly Logger log = LogManager.GetCurrentClassLogger();
         CompositeDisposable disposables = new CompositeDisposable();
         readonly IServiceProvider serviceProvider;
         CompositionContainer tempContainer;
-        readonly Dictionary<string, ComposablePart> tempParts;
+        readonly Dictionary<string, OwnedComposablePart> tempParts;
         ExportLifetimeContext<IUIController> currentUIFlow;
         readonly Version currentVersion;
         bool initializingLogging = false;
@@ -72,7 +78,7 @@ public UIProvider([Import(typeof(SVsServiceProvider))] IServiceProvider serviceP
             {
                 SourceProvider = ExportProvider
             }));
-            tempParts = new Dictionary<string, ComposablePart>();
+            tempParts = new Dictionary<string, OwnedComposablePart>();
         }
 
         [return: AllowNull]
@@ -152,12 +158,12 @@ public Ret GetService<T, Ret>() where Ret : class
             return GetService<T>() as Ret;
         }
 
-        public void AddService<T>(T instance)
+        public void AddService<T>(object owner, T instance)
         {
-            AddService(typeof(T), instance);
+            AddService(typeof(T), owner, instance);
         }
 
-        public void AddService(Type t, object instance)
+        public void AddService(Type t, object owner, object instance)
         {
             if (!Initialized)
             {
@@ -168,13 +174,23 @@ public void AddService(Type t, object instance)
             var batch = new CompositionBatch();
             string contract = AttributedModelServices.GetContractName(t);
             Debug.Assert(!string.IsNullOrEmpty(contract), "Every type must have a contract name");
+
+            // we want to remove stale instances of a service, if they exist, regardless of who put them there
+            RemoveService(t, null);
+
             var part = batch.AddExportedValue(contract, instance);
             Debug.Assert(part != null, "Adding an exported value must return a non-null part");
-            tempParts.Add(contract, part);
+            tempParts.Add(contract, new OwnedComposablePart { Owner = owner, Part = part });
             tempContainer.Compose(batch);
         }
 
-        public void RemoveService(Type t)
+        /// <summary>
+        /// Removes a service from the catalog
+        /// </summary>
+        /// <param name="t">The type we want to remove</param>
+        /// <param name="owner">The owner, which either has to match what was passed to AddService,
+        /// or if it's null, the service will be removed without checking for ownership</param>
+        public void RemoveService(Type t, [AllowNull] object owner)
         {
             if (!Initialized)
             {
@@ -185,13 +201,14 @@ public void RemoveService(Type t)
             string contract = AttributedModelServices.GetContractName(t);
             Debug.Assert(!string.IsNullOrEmpty(contract), "Every type must have a contract name");
 
-            ComposablePart part; 
-
+            OwnedComposablePart part; 
             if (tempParts.TryGetValue(contract, out part))
             {
+                if (owner != null && part.Owner != owner)
+                    return;
                 tempParts.Remove(contract);
                 var batch = new CompositionBatch();
-                batch.RemovePart(part);
+                batch.RemovePart(part.Part);
                 tempContainer.Compose(batch);
             }
         }

src/UnitTests/GitHub.App/Controllers/UIControllerTests.cs

@@ -124,7 +124,7 @@ public void CloneDialogLoggedInWithoutConnection()
                                 {
                                     Assert.Equal(1, list.Count);
                                     Assert.IsAssignableFrom<IViewFor<IRepositoryCloneViewModel>>(list[0]);
-                                    ((IUIProvider)provider).Received().AddService(connection);
+                                    ((IUIProvider)provider).Received().AddService(uiController, connection);
                                 });
                 uiController.Start(null);
             }