Merge pull request #2288 from Suor/pool-fixes

remote: pool fixes

dvc/main.py

@@ -8,6 +8,7 @@ from dvc.cli import parse_args
 from dvc.config import ConfigError
 from dvc.analytics import Analytics
 from dvc.exceptions import NotDvcRepoError, DvcParserError
+from dvc.remote.ssh.pool import close_ssh_pools
 
 
 logger = logging.getLogger("dvc")
@@ -53,6 +54,9 @@ def main(argv=None):
         ret = 255
     finally:
         logger.setLevel(outerLogLevel)
+        # Python 2 fails to close these clean occasionally and users see
+        # weird error messages, so we do it manually
+        close_ssh_pools()
 
     Analytics().send_cmd(cmd, args, ret)
 

dvc/remote/ssh/pool.py

@@ -13,6 +13,7 @@ def ssh_connection(*conn_args, **conn_kwargs):
         yield conn
     except BaseException:
         conn.close()
+        raise
     else:
         pool.release(conn)
 
@@ -22,6 +23,12 @@ def get_ssh_pool(*conn_args, **conn_kwargs):
     return SSHPool(conn_args, conn_kwargs)
 
 
+def close_ssh_pools():
+    for pool in get_ssh_pool.memory.values():
+        pool.close()
+    get_ssh_pool.memory.clear()
+
+
 class SSHPool(object):
     def __init__(self, conn_args, conn_kwargs):
         self._conn_args = conn_args

tests/conftest.py

@@ -83,17 +83,18 @@ key_path = os.path.join(here, "{0}.key".format(user))
 def ssh_server():
     users = {user: key_path}
     with mockssh.Server(users) as s:
+        s.test_creds = {
+            "host": s.host,
+            "port": s.port,
+            "username": user,
+            "key_filename": key_path,
+        }
         yield s
 
 
 @pytest.fixture
 def ssh(ssh_server):
-    yield SSHConnection(
-        ssh_server.host,
-        username=user,
-        port=ssh_server.port,
-        key_filename=key_path,
-    )
+    yield SSHConnection(**ssh_server.test_creds)
 
 
 @pytest.fixture

tests/func/test_data_cloud.py

@@ -31,8 +31,6 @@ from dvc.utils import file_md5
 from dvc.utils.stage import load_stage_file, dump_stage_file
 
 from tests.basic_env import TestDvc
-from tests.conftest import user
-from tests.conftest import key_path
 from tests.utils import spy
 
 
@@ -460,10 +458,11 @@ class TestRemoteSSHMocked(TestDataCloudBase):
         self.method_name = request.function.__name__
 
     def _get_url(self):
+        user = self.ssh_server.test_creds["username"]
         return get_ssh_url_mocked(user, self.ssh_server.port)
 
     def _get_keyfile(self):
-        return key_path
+        return self.ssh_server.test_creds["key_filename"]
 
     def _should_test(self):
         return True

tests/unit/remote/ssh/test_pool.py

@@ -0,0 +1,11 @@
+import pytest
+
+from dvc.remote.ssh.pool import ssh_connection
+
+
+def test_doesnt_swallow_errors(ssh_server):
+    class MyError(Exception):
+        pass
+
+    with pytest.raises(MyError), ssh_connection(**ssh_server.test_creds):
+        raise MyError