Tweaked entry point per MZT8's Connect.VBA.cs; still experiencing hos…

…t app failure on shutdown.
retailcoder · Sep 20, 2016 · 8ac768d · 8ac768d
1 parent 21e8668
commit 8ac768d
Showing 1 changed file with 73 additions and 26 deletions.
diff --git a/RetailCoder.VBE/Extension.cs b/RetailCoder.VBE/Extension.cs
@@ -6,7 +6,6 @@
 using Rubberduck.UI;
 using System;
 using System.ComponentModel;
-using System.Diagnostics;
 using System.Globalization;
 using System.IO;
 using System.Reflection;
@@ -20,6 +19,9 @@
 
 namespace Rubberduck
 {
+    /// <remarks>
+    /// Special thanks to Carlos Quintero (MZ-Tools) for providing the general structure here.
+    /// </remarks>
     [ComVisible(true)]
     [Guid(ClassId)]
     [ProgId(ProgId)]
@@ -30,27 +32,40 @@ public class _Extension : IDTExtensibility2
         private const string ClassId = "8D052AD8-BBD2-4C59-8DEC-F697CA1F8A66";
         private const string ProgId = "Rubberduck.Extension";
 
-        private dynamic _ide;
+        private VBE _ide;
         private AddIn _addin;
+        private bool _isInitialized;
+        private bool _isBeginShutdownExecuted;
 
         private IKernel _kernel;
         private App _app;
         private readonly Logger _logger = LogManager.GetCurrentClassLogger();
 
-        public void OnAddInsUpdate(ref Array custom)
-        {
-        }
-
-        public void OnBeginShutdown(ref Array custom)
-        {
-
-        }
+        public void OnAddInsUpdate(ref Array custom) { }
 
         // ReSharper disable InconsistentNaming
         public void OnConnection(object Application, ext_ConnectMode ConnectMode, object AddInInst, ref Array custom)
         {
-            _ide = Application;
-            _addin = (AddIn)AddInInst;
+            try
+            {
+                _ide = (VBE)Application;
+                _addin = (AddIn)AddInInst;
+                _addin.Object = this;
+
+                switch (ConnectMode)
+                {
+                    case ext_ConnectMode.ext_cm_Startup:
+                        // normal execution path - don't initialize just yet, wait for OnStartupComplete to be called by the host.
+                        break;
+                    case ext_ConnectMode.ext_cm_AfterStartup:
+                        InitializeAddIn();
+                        break;
+                }
+            }
+            catch (Exception e)
+            {
+                Console.WriteLine(e);
+            }
         }
 
         Assembly LoadFromSameFolder(object sender, ResolveEventArgs args)
@@ -68,7 +83,48 @@ Assembly LoadFromSameFolder(object sender, ResolveEventArgs args)
 
         public void OnStartupComplete(ref Array custom)
         {
-            Debug.Assert(_kernel == null);
+            InitializeAddIn();
+        }
+
+        public void OnBeginShutdown(ref Array custom)
+        {
+            _isBeginShutdownExecuted = true;
+            ShutdownAddIn();
+        }
+
+        // ReSharper disable InconsistentNaming
+        public void OnDisconnection(ext_DisconnectMode RemoveMode, ref Array custom)
+        {
+            switch (RemoveMode)
+            {
+                case ext_DisconnectMode.ext_dm_UserClosed:
+                    ShutdownAddIn();
+                    break;
+
+                case ext_DisconnectMode.ext_dm_HostShutdown:
+                    if (_isBeginShutdownExecuted)
+                    {
+                        // this is the normal case: nothing to do here, we already ran ShutdownAddIn.
+                    }
+                    else
+                    {
+                        // some hosts do not call OnBeginShutdown: this mitigates it.
+                        ShutdownAddIn();
+                    }
+                    break;
+            }
+        }
+
+        private void InitializeAddIn()
+        {
+            if (_isInitialized)
+            {
+                // The add-in is already initialized. See:
+                // The strange case of the add-in initialized twice
+                // http://msmvps.com/blogs/carlosq/archive/2013/02/14/the-strange-case-of-the-add-in-initialized-twice.aspx
+                return;
+            }
+
             _kernel = new StandardKernel(new NinjectSettings { LoadExtensions = true }, new FuncModule(), new DynamicProxyModule());
 
             try
@@ -92,21 +148,11 @@ public void OnStartupComplete(ref Array custom)
                     catch (CultureNotFoundException) { }
                 }
 
-                var vbe = _ide as VBE;
-                if (vbe != null)
-                {
-                    _kernel.Load(new RubberduckModule(vbe, _addin));
-                }
-                else
-                {
-                    Type type = _ide.GetType();
-                    var ideName = type.GetProperty("Name");
-                    // VB6?
-                    // todo: wrap up the VBE interface with... ...generic members?    
-                }
+                _kernel.Load(new RubberduckModule(_ide, _addin));
 
                 _app = _kernel.Get<App>();
                 _app.Startup();
+                _isInitialized = true;
             }
             catch (Exception exception)
             {
@@ -115,7 +161,7 @@ public void OnStartupComplete(ref Array custom)
             }
         }
 
-        public void OnDisconnection(ext_DisconnectMode RemoveMode, ref Array custom)
+        private void ShutdownAddIn()
         {
             if (_app != null)
             {
@@ -131,6 +177,7 @@ public void OnDisconnection(ext_DisconnectMode RemoveMode, ref Array custom)
 
             Marshal.ReleaseComObject(_addin);
             Marshal.ReleaseComObject(_ide);
+            _isInitialized = false;
         }
     }
 }