git: use asyncssh as dulwich SSH vendor (#6797)

* git: support client kwargs for auth in remote methods

* tests: use docker for git/ssh server

* move dulwich backend into subdir

* add base sync<->async support

* dulwich: add initial asyncssh SSH vendor

* dulwich: use asyncssh as the default SSH vendor

* add test for git over ssh

* test read (fetch) over ssh

* ci: re-enable ssh tests

temporary workaround for #6788, #6789

* move async stubs into dvc.scm.asyn

* tests: use sshfs and not dvc.fs.ssh

* allow explicit 0-length read in asyncssh vendor

.github/workflows/tests.yaml

@@ -109,7 +109,7 @@ jobs:
       run: >-
         python -m tests -n=4
         --cov-report=xml --cov-report=term
-        ${{ env.extra_test_args }}
+        --enable-ssh ${{ env.extra_test_args }}
     - name: upload coverage report
       uses: codecov/[email protected]
       with:

dvc/scm/asyn.py

@@ -0,0 +1,48 @@
+"""DVC re-implementation of fsspec's dedicated async event loop."""
+import asyncio
+import os
+import threading
+from typing import List, Optional
+
+from fsspec.asyn import (  # noqa: F401, pylint:disable=unused-import
+    _selector_policy,
+    sync,
+    sync_wrapper,
+)
+
+# dedicated async IO thread
+iothread: List[Optional[threading.Thread]] = [None]
+# global DVC event loop
+default_loop: List[Optional[asyncio.AbstractEventLoop]] = [None]
+lock = threading.Lock()
+
+
+def get_loop() -> asyncio.AbstractEventLoop:
+    """Create or return the global DVC event loop."""
+    if default_loop[0] is None:
+        with lock:
+            if default_loop[0] is None:
+                with _selector_policy():
+                    default_loop[0] = asyncio.new_event_loop()
+                loop = default_loop[0]
+                th = threading.Thread(
+                    target=loop.run_forever,  # type: ignore[attr-defined]
+                    name="dvcIO",
+                )
+                th.daemon = True
+                th.start()
+                iothread[0] = th
+    assert default_loop[0] is not None
+    return default_loop[0]
+
+
+class BaseAsyncObject:
+    def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None):
+        self._loop: asyncio.AbstractEventLoop = loop or get_loop()
+        self._pid = os.getpid()
+
+    @property
+    def loop(self):
+        # AsyncMixin is not fork-safe
+        assert self._pid == os.getpid()
+        return self._loop

dvc/scm/git/backend/base.py

@@ -178,7 +178,7 @@ class BaseGitBackend(ABC):
         """
 
     @abstractmethod
-    def iter_remote_refs(self, url: str, base: Optional[str] = None):
+    def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs):
         """Iterate over all refs in the specified remote Git repo.
 
         If base is specified, only refs which begin with base will be yielded.
@@ -197,6 +197,7 @@ class BaseGitBackend(ABC):
         dest: str,
         force: bool = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         """Push refspec to a remote Git repo.
 
@@ -222,6 +223,7 @@ class BaseGitBackend(ABC):
         refspecs: Iterable[str],
         force: Optional[bool] = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         """Fetch refspecs from a remote Git repo.
 
@@ -349,5 +351,5 @@ class BaseGitBackend(ABC):
         """
 
     @abstractmethod
-    def validate_git_remote(self, url: str):
+    def validate_git_remote(self, url: str, **kwargs):
         """Verify that url is a valid git URL or remote name."""

dvc/scm/git/backend/dulwich/__init__.py

@@ -23,13 +23,13 @@ from dvc.progress import Tqdm
 from dvc.scm.base import GitAuthError, InvalidRemoteSCMRepo, SCMError
 from dvc.utils import relpath
 
-from ..objects import GitObject
-from .base import BaseGitBackend
+from ...objects import GitObject
+from ..base import BaseGitBackend
 
 if TYPE_CHECKING:
     from dvc.types import StrPath
 
-    from ..objects import GitCommit
+    from ...objects import GitCommit
 
 logger = logging.getLogger(__name__)
 
@@ -80,6 +80,13 @@ class DulwichObject(GitObject):
 class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     """Dulwich Git backend."""
 
+    from dulwich import client
+
+    from .asyncssh_vendor import AsyncSSHVendor
+
+    # monkeypatch dulwich client's default SSH vendor to use asyncssh
+    client.get_ssh_vendor = AsyncSSHVendor
+
     # Dulwich progress will return messages equivalent to git CLI,
     # our pbars should just display the messages as formatted by dulwich
     BAR_FMT_NOTOTAL = "{desc}{bar:b}|{postfix[info]} [{elapsed}]"
@@ -354,14 +361,14 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
             else:
                 yield os.fsdecode(key)
 
-    def iter_remote_refs(self, url: str, base: Optional[str] = None):
+    def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs):
         from dulwich.client import HTTPUnauthorized, get_transport_and_path
         from dulwich.errors import NotGitRepository
         from dulwich.porcelain import get_remote_repo
 
         try:
             _remote, location = get_remote_repo(self.repo, url)
-            client, path = get_transport_and_path(location)
+            client, path = get_transport_and_path(location, **kwargs)
         except Exception as exc:
             raise InvalidRemoteSCMRepo(url) from exc
 
@@ -389,6 +396,7 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
         dest: str,
         force: bool = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         from dulwich.client import HTTPUnauthorized, get_transport_and_path
         from dulwich.errors import NotGitRepository, SendPackError
@@ -402,7 +410,7 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
 
         try:
             _remote, location = get_remote_repo(self.repo, url)
-            client, path = get_transport_and_path(location)
+            client, path = get_transport_and_path(location, **kwargs)
         except Exception as exc:
             raise SCMError(
                 f"'{url}' is not a valid Git remote or URL"
@@ -476,6 +484,7 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
         refspecs: Iterable[str],
         force: Optional[bool] = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         from dulwich.client import get_transport_and_path
         from dulwich.objectspec import parse_reftuples
@@ -504,7 +513,7 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
 
         try:
             _remote, location = get_remote_repo(self.repo, url)
-            client, path = get_transport_and_path(location)
+            client, path = get_transport_and_path(location, **kwargs)
         except Exception as exc:
             raise SCMError(
                 f"'{url}' is not a valid Git remote or URL"
@@ -659,13 +668,13 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
     ) -> Optional[str]:
         raise NotImplementedError
 
-    def validate_git_remote(self, url: str):
+    def validate_git_remote(self, url: str, **kwargs):
         from dulwich.client import LocalGitClient, get_transport_and_path
         from dulwich.porcelain import get_remote_repo
 
         try:
             _, location = get_remote_repo(self.repo, url)
-            client, path = get_transport_and_path(location)
+            client, path = get_transport_and_path(location, **kwargs)
         except Exception as exc:
             raise InvalidRemoteSCMRepo(url) from exc
         if isinstance(client, LocalGitClient) and not os.path.exists(

dvc/scm/git/backend/dulwich/asyncssh_vendor.py

@@ -0,0 +1,92 @@
+"""asyncssh SSH vendor for Dulwich."""
+from typing import List, Optional
+
+from dulwich.client import SSHVendor
+
+from dvc.scm.asyn import BaseAsyncObject, sync_wrapper
+
+
+class _StderrWrapper:
+    def __init__(self, stderr):
+        self.stderr = stderr
+
+    async def _readlines(self):
+        lines = []
+        while True:
+            line = await self.stderr.readline()
+            if not line:
+                break
+            lines.append(line)
+        return lines
+
+    readlines = sync_wrapper(_readlines)
+
+
+class AsyncSSHWrapper(BaseAsyncObject):
+    def __init__(self, conn, proc, **kwargs):
+        super().__init__(**kwargs)
+        self.conn = conn
+        self.proc = proc
+        self.stderr = _StderrWrapper(proc.stderr)
+
+    def can_read(self) -> bool:
+        # pylint:disable=protected-access
+        return self.proc.stdout._session._recv_buf_len > 0
+
+    async def _read(self, n: Optional[int] = None) -> bytes:
+        if self.proc.stdout.at_eof():
+            return b""
+
+        return await self.proc.stdout.read(n=n if n is not None else -1)
+
+    read = sync_wrapper(_read)
+
+    def write(self, data: bytes):
+        self.proc.stdin.write(data)
+
+    def close(self):
+        self.conn.close()
+
+
+class AsyncSSHVendor(BaseAsyncObject, SSHVendor):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    async def _run_command(
+        self,
+        host: str,
+        command: List[str],
+        username: Optional[str] = None,
+        port: Optional[int] = None,
+        password: Optional[str] = None,
+        key_filename: Optional[str] = None,
+        **kwargs,
+    ):
+        """Connect to an SSH server.
+
+        Run a command remotely and return a file-like object for interaction
+        with the remote command.
+
+        Args:
+          host: Host name
+          command: Command to run (as argv array)
+          username: Optional ame of user to log in as
+          port: Optional SSH port to use
+          password: Optional ssh password for login or private key
+          key_filename: Optional path to private keyfile
+        """
+        import asyncssh
+
+        conn = await asyncssh.connect(
+            host,
+            port=port,
+            username=username,
+            password=password,
+            client_keys=[key_filename] if key_filename else [],
+            known_hosts=None,
+            encoding=None,
+        )
+        proc = await conn.create_process(command, encoding=None)
+        return AsyncSSHWrapper(conn, proc)
+
+    run_command = sync_wrapper(_run_command)

dvc/scm/git/backend/gitpython.py

@@ -438,7 +438,7 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
         for ref in Reference.iter_items(self.repo, common_path=base):
             yield ref.path
 
-    def iter_remote_refs(self, url: str, base: Optional[str] = None):
+    def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs):
         raise NotImplementedError
 
     def get_refs_containing(self, rev: str, pattern: Optional[str] = None):
@@ -465,6 +465,7 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
         dest: str,
         force: bool = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         raise NotImplementedError
 
@@ -474,6 +475,7 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
         refspecs: Iterable[str],
         force: Optional[bool] = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         raise NotImplementedError
 
@@ -627,5 +629,5 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
             raise SCMError("Merge failed") from exc
         return None
 
-    def validate_git_remote(self, url: str):
+    def validate_git_remote(self, url: str, **kwargs):
         raise NotImplementedError

dvc/scm/git/backend/pygit2.py

@@ -374,6 +374,7 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         dest: str,
         force: bool = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         raise NotImplementedError
 
@@ -383,6 +384,7 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
         refspecs: Iterable[str],
         force: Optional[bool] = False,
         on_diverged: Optional[Callable[[str, str], bool]] = None,
+        **kwargs,
     ):
         raise NotImplementedError
 
@@ -537,7 +539,7 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
                     index.add(entry.path)
                 index.write()
 
-    def iter_remote_refs(self, url: str, base: Optional[str] = None):
+    def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs):
         raise NotImplementedError
 
     def status(
@@ -584,5 +586,5 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
                 self.repo.index.write()
         return None
 
-    def validate_git_remote(self, url: str):
+    def validate_git_remote(self, url: str, **kwargs):
         raise NotImplementedError

tests/docker-compose.yml

@@ -37,3 +37,14 @@ services:
       - 2222
     volumes:
       - ./remotes/user.key.pub:/tmp/key
+
+  git-server:
+    image: ghcr.io/linuxserver/openssh-server
+    environment:
+      - USER_NAME=user
+      - PUBLIC_KEY_FILE=/tmp/key
+    ports:
+      - 2222
+    volumes:
+      - ./remotes/user.key.pub:/tmp/key
+      - ./remotes/git-init:/config/custom-cont-init.d

tests/func/test_scm.py

@@ -331,3 +331,52 @@ def test_git_stash_clear(tmp_dir, scm, ref):
         not os.path.exists(log_path)
         or not open(log_path, encoding="utf-8").read()
     )
+
+
+@pytest.mark.needs_internet
+@pytest.mark.parametrize("server", [pytest.lazy_fixture("git_ssh")])
+def test_git_ssh(tmp_dir, scm, server):
+    from dulwich.repo import Repo as DulwichRepo
+    from sshfs import SSHFileSystem
+
+    from dvc.utils.fs import remove
+    from tests.remotes.ssh import TEST_SSH_KEY_PATH, TEST_SSH_USER
+
+    fs = SSHFileSystem(
+        host=server.host,
+        port=server.port,
+        username=TEST_SSH_USER,
+        client_keys=[TEST_SSH_KEY_PATH],
+    )
+    server._ssh.execute("git init --bare test-repo.git")
+    url = f"ssh://{TEST_SSH_USER}@{server.host}:{server.port}/~/test-repo.git"
+
+    tmp_dir.scm_gen("foo", "foo", commit="init")
+    rev = scm.get_rev()
+
+    scm.push_refspec(
+        url,
+        "refs/heads/master",
+        "refs/heads/master",
+        force=True,
+        key_filename=TEST_SSH_KEY_PATH,
+    )
+
+    assert (
+        rev.encode("ascii")
+        == fs.open("test-repo.git/refs/heads/master").read().strip()
+    )
+
+    remove(tmp_dir / ".git")
+    remove(tmp_dir / "foo")
+    DulwichRepo.init(str(tmp_dir))
+
+    scm.fetch_refspecs(
+        url,
+        ["refs/heads/master"],
+        force=True,
+        key_filename=TEST_SSH_KEY_PATH,
+    )
+    assert rev == scm.get_ref("refs/heads/master")
+    scm.checkout("master")
+    assert "foo" == (tmp_dir / "foo").read_text()

tests/remotes/__init__.py

@@ -8,6 +8,7 @@ from .gdrive import (  # noqa: F401; noqa: F401
     GDrive,
     gdrive,
 )
+from .git_server import git_server, git_ssh  # noqa: F401
 from .gs import (  # noqa: F401; noqa: F401
     GCP,
     TEST_GCP_CREDS_FILE,

tests/remotes/git-init/git.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+apk add --no-cache git

tests/remotes/git_server.py

@@ -0,0 +1,60 @@
+import pytest
+
+from .ssh import SSH, TEST_SSH_KEY_PATH, TEST_SSH_USER
+
+
+class GitSSH(SSH):
+    @staticmethod
+    def get_url(host, port):  # pylint: disable=arguments-differ
+        return f"ssh://{host}:{port}/tmp/data/git"
+
+
+@pytest.fixture
+def git_server(test_config, docker_compose, docker_services):
+    import asyncssh
+    from sshfs import SSHFileSystem
+
+    test_config.requires("ssh")
+    conn_info = {
+        "host": "127.0.0.1",
+        "port": docker_services.port_for("git-server", 2222),
+    }
+
+    def get_fs():
+        return SSHFileSystem(
+            **conn_info,
+            username=TEST_SSH_USER,
+            client_keys=[TEST_SSH_KEY_PATH],
+        )
+
+    def _check():
+        try:
+            fs = get_fs()
+            fs.exists("/")
+            fs.execute("git --version")
+        except asyncssh.Error:
+            return False
+        else:
+            return True
+
+    docker_services.wait_until_responsive(timeout=30.0, pause=1, check=_check)
+    return conn_info
+
+
+@pytest.fixture
+def git_ssh_connection(git_server):
+    from sshfs import SSHFileSystem
+
+    yield SSHFileSystem(
+        host=git_server["host"],
+        port=git_server["port"],
+        username=TEST_SSH_USER,
+        client_keys=[TEST_SSH_KEY_PATH],
+    )
+
+
+@pytest.fixture
+def git_ssh(git_server, monkeypatch):
+    url = GitSSH(GitSSH.get_url(**git_server))
+    url.mkdir(exist_ok=True, parents=True)
+    return url