Fix leaking ssh connections with VS Attach to Process

src/MIDebugEngine/AD7.Impl/AD7Engine.cs

@@ -229,6 +229,7 @@ public int Attach(IDebugProgram2[] portProgramArray, IDebugProgramNode2[] progra
                     if (port is IDebugUnixShellPort)
                     {
                         _unixPort = (IDebugUnixShellPort)port;
+                        (_unixPort as IDebugPortCleanup)?.AddSessionRef();
                     }
                     StartDebugging(launchOptions);
                 }

src/Microsoft.VisualStudio.Debugger.Interop.UnixPortSupplier/Microsoft.VisualStudio.Debugger.Interop.UnixPortSupplier.cs

@@ -73,7 +73,7 @@ public interface IDebugUnixShellPort
     }
 
     /// <summary>
-    /// Interface implemented by a port that supports explicit cleanup
+    /// Interface implemented by a port that supports explicit cleanup of shared connections.
     /// </summary>
     [ComImport()]
     [ComVisible(true)]
@@ -82,9 +82,14 @@ public interface IDebugUnixShellPort
     public interface IDebugPortCleanup
     {
         /// <summary>
-        /// Clean up debugging resources
+        /// Decrement the session reference count and close the connection when it reaches zero.
         /// </summary>
         void Clean();
+
+        /// <summary>
+        /// Increment the session reference count on this port.
+        /// </summary>
+        void AddSessionRef();
     }
 
     /// <summary>

src/SSHDebugPS/AD7/AD7Port.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Globalization;
 using System.Linq;
 using System.Threading;
@@ -21,6 +22,8 @@ internal abstract class AD7Port : IDebugPort2, IDebugUnixShellPort, IDebugPortCl
         private Connection _connection;
         private readonly Dictionary<uint, IDebugPortEvents2> _eventCallbacks = new Dictionary<uint, IDebugPortEvents2>();
         private uint _lastCallbackCookie;
+        private int _sessionRefCount;
+        private int _activeAsyncCommands;
 
         protected string Name { get; private set; }
 
@@ -37,16 +40,19 @@ public AD7Port(AD7PortSupplier portSupplier, string name, bool isInAddPort)
 
         protected Connection GetConnection()
         {
-            if (_connection == null)
+            lock (_lock)
             {
-                _connection = GetConnectionInternal();
-                if (_connection != null)
+                if (_connection == null)
                 {
-                    Name = _connection.Name;
+                    _connection = GetConnectionInternal();
+                    if (_connection != null)
+                    {
+                        Name = _connection.Name;
+                    }
                 }
-            }
 
-            return _connection;
+                return _connection;
+            }
         }
 
         protected abstract Connection GetConnectionInternal();
@@ -60,7 +66,10 @@ public bool IsConnected
         {
             get
             {
-                return _connection != null;
+                lock (_lock)
+                {
+                    return _connection != null;
+                }
             }
         }
 
@@ -88,6 +97,8 @@ private AD7Process[] EnumProcessesInternal()
                 result = processList.Select((proc) => new AD7Process(this, proc)).ToArray();
             });
 
+            CloseConnectionIfIdle();
+
             return result;
         }
 
@@ -155,7 +166,92 @@ void IDebugUnixShellPort.ExecuteSyncCommand(string commandDescription, string co
 
         void IDebugUnixShellPort.BeginExecuteAsyncCommand(string commandText, bool runInShell, IDebugUnixShellCommandCallback callback, out IDebugUnixShellAsyncCommand asyncCommand)
         {
-            GetConnection().BeginExecuteAsyncCommand(commandText, runInShell, callback, out asyncCommand);
+            var wrappedCallback = new AsyncCommandCallback(this, callback);
+            var connection = GetConnection();
+            lock (_lock)
+            {
+                _activeAsyncCommands++;
+            }
+            connection.BeginExecuteAsyncCommand(commandText, runInShell, wrappedCallback, out asyncCommand);
+            asyncCommand = new AsyncCommandWrapper(this, asyncCommand);
+        }
+
+        private class AsyncCommandWrapper : IDebugUnixShellAsyncCommand
+        {
+            private readonly AD7Port _port;
+            private readonly IDebugUnixShellAsyncCommand _inner;
+            private int _notified;
+
+            public AsyncCommandWrapper(AD7Port port, IDebugUnixShellAsyncCommand inner)
+            {
+                _port = port;
+                _inner = inner;
+            }
+
+            public void Write(string text) => _inner.Write(text);
+            public void WriteLine(string text) => _inner.WriteLine(text);
+
+            public void Abort()
+            {
+                _inner.Abort();
+                // If OnExit didn't fire (abort path), notify the port now
+                if (Interlocked.CompareExchange(ref _notified, 1, 0) == 0)
+                {
+                    _port.OnAsyncCommandExited();
+                }
+            }
+        }
+
+        private void OnAsyncCommandExited()
+        {
+            lock (_lock)
+            {
+                _activeAsyncCommands--;
+            }
+            CloseConnectionIfIdle();
+        }
+
+        private void CloseConnectionIfIdle()
+        {
+            lock (_lock)
+            {
+                if (_sessionRefCount > 0 || _activeAsyncCommands > 0 || _connection == null)
+                {
+                    return;
+                }
+
+                var conn = _connection;
+                _connection = null;
+                try { conn.Close(); } catch (Exception) { }
+            }
+        }
+
+        /// <summary>
+        /// Wraps an IDebugUnixShellCommandCallback to detect when an async command exits,
+        /// allowing the port to close idle connections for engines that do not call
+        /// IDebugPortCleanup.AddSessionRef/Clean (e.g., vsdbg).
+        /// </summary>
+        private class AsyncCommandCallback : IDebugUnixShellCommandCallback
+        {
+            private readonly AD7Port _port;
+            private readonly IDebugUnixShellCommandCallback _inner;
+
+            public AsyncCommandCallback(AD7Port port, IDebugUnixShellCommandCallback inner)
+            {
+                _port = port;
+                _inner = inner;
+            }
+
+            public void OnOutputLine(string line)
+            {
+                _inner.OnOutputLine(line);
+            }
+
+            public void OnExit(string exitCode)
+            {
+                _inner.OnExit(exitCode);
+                _port.OnAsyncCommandExited();
+            }
         }
 
         void IConnectionPointContainer.EnumConnectionPoints(out IEnumConnectionPoints ppEnum)
@@ -241,14 +337,37 @@ public bool IsLinux()
             return GetConnection().IsLinux();
         }
 
+        public void AddSessionRef()
+        {
+            lock (_lock)
+            {
+                _sessionRefCount++;
+            }
+            EnsureConnected();
+        }
+
         public void Clean()
         {
-            try
+            lock (_lock)
             {
-                _connection?.Close();
+                Debug.Assert(_sessionRefCount > 0, "Unbalanced call to Clean -- no matching AddSessionRef");
+                _sessionRefCount--;
+
+                if (_sessionRefCount > 0)
+                {
+                    return;
+                }
+
+                var conn = _connection;
+                _connection = null;
+
+                try
+                {
+                    conn?.Close();
+                }
+                // Dev15 632648: Liblinux sometimes throws exceptions on shutdown - we are shutting down anyways, so ignore to not crash
+                catch (Exception) { }
             }
-            // Dev15 632648: Liblinux sometimes throws exceptions on shutdown - we are shutting down anyways, so ignore to not crash
-            catch (Exception) { }
         }
     }
 }

src/SSHDebugPS/SSH/SSHHelper.cs

@@ -23,53 +23,61 @@ internal static SSHConnection CreateSSHConnectionFromConnectionInfo(ConnectionIn
             if (connectionInfo != null)
             {
                 UnixSystem remoteSystem = new UnixSystem();
-                string name = SSHPortSupplier.GetFormattedSSHConnectionName(connectionInfo);
-
-                while (true)
+                bool success = false;
+                try
                 {
-                    try
-                    {
-                        VSOperationWaiter.Wait(
-                            StringResources.WaitingOp_Connecting.FormatCurrentCultureWithArgs(name), 
-                            throwOnCancel: false, 
-                            action: (cancellationToken) =>
-                                remoteSystem.Connect(connectionInfo));
-                        break;
-                    }
-                    catch (RemoteAuthenticationException)
+                    string name = SSHPortSupplier.GetFormattedSSHConnectionName(connectionInfo);
+
+                    while (true)
                     {
-                        IVsConnectionManager connectionManager = (IVsConnectionManager)ServiceProvider.GlobalProvider.GetService(typeof(IVsConnectionManager));
-                        if (connectionManager != null)
+                        try
                         {
-                            IConnectionManagerResult result = connectionManager.ShowDialog(StringResources.AuthenticationFailureHeader, StringResources.AuthenticationFailureDescription, connectionInfo);
-
-                            if (result != null && (result.DialogResult & ConnectionManagerDialogResult.Succeeded) == ConnectionManagerDialogResult.Succeeded)
+                            VSOperationWaiter.Wait(
+                                StringResources.WaitingOp_Connecting.FormatCurrentCultureWithArgs(name),
+                                throwOnCancel: false,
+                                action: (cancellationToken) =>
+                                    remoteSystem.Connect(connectionInfo));
+                            break;
+                        }
+                        catch (RemoteAuthenticationException)
+                        {
+                            IVsConnectionManager connectionManager = (IVsConnectionManager)ServiceProvider.GlobalProvider.GetService(typeof(IVsConnectionManager));
+                            if (connectionManager != null)
                             {
-                                connectionInfo = result.ConnectionInfo;
+                                IConnectionManagerResult result = connectionManager.ShowDialog(StringResources.AuthenticationFailureHeader, StringResources.AuthenticationFailureDescription, connectionInfo);
+
+                                if (result != null && (result.DialogResult & ConnectionManagerDialogResult.Succeeded) == ConnectionManagerDialogResult.Succeeded)
+                                {
+                                    connectionInfo = result.ConnectionInfo;
+                                }
+                                else
+                                {
+                                    return null;
+                                }
                             }
                             else
                             {
-                                return null;
+                                throw new InvalidOperationException("Why is IVsConnectionManager null?");
                             }
                         }
-                        else
+                        catch (Exception ex)
                         {
-                            throw new InvalidOperationException("Why is IVsConnectionManager null?");
+                            VsShellUtilities.ShowMessageBox(ServiceProvider.GlobalProvider, ex.Message, null,
+                                OLEMSGICON.OLEMSGICON_CRITICAL, OLEMSGBUTTON.OLEMSGBUTTON_OK, OLEMSGDEFBUTTON.OLEMSGDEFBUTTON_FIRST);
+                            return null;
                         }
                     }
-                    catch (Exception ex)
-                    {
-                        VsShellUtilities.ShowMessageBox(ServiceProvider.GlobalProvider, ex.Message, null,
-                            OLEMSGICON.OLEMSGICON_CRITICAL, OLEMSGBUTTON.OLEMSGBUTTON_OK, OLEMSGDEFBUTTON.OLEMSGDEFBUTTON_FIRST);
-                        return null;
-                    }
-                }
 
-                // NOTE: This will be null if connect is canceled
-                if (remoteSystem != null)
-                {
+                    success = true;
                     return new SSHConnection(remoteSystem);
                 }
+                finally
+                {
+                    if (!success)
+                    {
+                        remoteSystem.Dispose();
+                    }
+                }
             }
 
             return null;

src/SSHDebugPS/SSH/SSHPortSupplier.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Runtime.InteropServices;
 using liblinux;
@@ -20,6 +21,8 @@ internal class SSHPortSupplier : AD7PortSupplier
     {
         private const string _Name = "SSH";
         private readonly Guid _Id = new Guid("3FDDF14E-E758-4695-BE0C-7509920432C9");
+        private readonly object _portCacheLock = new object();
+        private readonly Dictionary<string, SSHPort> _portCache = new Dictionary<string, SSHPort>(StringComparer.OrdinalIgnoreCase);
 
         protected override Guid Id { get { return _Id; } }
         protected override string Name { get { return _Name; } }
@@ -33,11 +36,26 @@ public override int AddPort(IDebugPortRequest2 request, out IDebugPort2 port)
             string name;
             HR.Check(request.GetPortName(out name));
 
-            AD7Port newPort = new SSHPort(this, name, isInAddPort: true);
+            SSHPort sshPort;
+            lock (_portCacheLock)
+            {
+                if (!_portCache.TryGetValue(name, out sshPort))
+                {
+                    sshPort = new SSHPort(this, name, isInAddPort: true);
+                    if (sshPort.IsConnected)
+                    {
+                        _portCache[name] = sshPort;
+                    }
+                }
+                else
+                {
+                    sshPort.EnsureConnected();
+                }
+            }
 
-            if (newPort.IsConnected)
+            if (sshPort.IsConnected)
             {
-                port = newPort;
+                port = sshPort;
                 return HR.S_OK;
             }
 
@@ -53,7 +71,18 @@ public override int EnumPorts(out IEnumDebugPorts2 ppEnum)
             for (int i = 0; i < store.Connections.Count; i++)
             {
                 ConnectionInfo connectionInfo = (ConnectionInfo)store.Connections[i];
-                ports[i] = new SSHPort(this, GetFormattedSSHConnectionName(connectionInfo), isInAddPort: false);
+                string name = GetFormattedSSHConnectionName(connectionInfo);
+
+                lock (_portCacheLock)
+                {
+                    if (!_portCache.TryGetValue(name, out SSHPort port))
+                    {
+                        port = new SSHPort(this, name, isInAddPort: false);
+                        _portCache[name] = port;
+                    }
+
+                    ports[i] = port;
+                }
             }
 
             ppEnum = new AD7PortEnum(ports);