feature: cache downloaded wheel information (#2276)

Co-authored-by: Bernát Gábor <gaborjbernat@gmail.com>

6 files changed, 199 insertions, 74 deletions

diff --git a/docs/changelog/2268.feature.rst b/docs/changelog/2268.feature.rst
new file mode 100644
index 0000000..73c9580
--- /dev/null
+++ b/docs/changelog/2268.feature.rst
@@ -0,0 +1,2 @@
+Add downloaded wheel information in the relevant JSON embed file to
+prevent additional downloads of the same wheel. - by :user:`mayeut`.
diff --git a/src/virtualenv/app_data/via_disk_folder.py b/src/virtualenv/app_data/via_disk_folder.py
index 265db1b..3f6afd5 100644
--- a/src/virtualenv/app_data/via_disk_folder.py
+++ b/src/virtualenv/app_data/via_disk_folder.py
@@ -14,7 +14,7 @@ virtualenv-app-data
 │       │           └── <install class> -> CopyPipInstall / SymlinkPipInstall
 │       │               └── <wheel name> -> pip-20.1.1-py2.py3-none-any
 │       └── embed
-│           └── 2 -> json format versioning
+│           └── 3 -> json format versioning
 │               └── *.json -> for every distribution contains data about newer embed versions and releases
 └─── unzip <in zip app we cannot refer to some internal files, so first extract them>
      └── <virtualenv version>
@@ -101,7 +101,7 @@ class AppDataDiskFolder(AppData):
                             filename.unlink()
 
     def embed_update_log(self, distribution, for_py_version):
-        return EmbedDistributionUpdateStoreDisk(self.lock / "wheel" / for_py_version / "embed" / "2", distribution)
+        return EmbedDistributionUpdateStoreDisk(self.lock / "wheel" / for_py_version / "embed" / "3", distribution)
 
     @property
     def house(self):
diff --git a/src/virtualenv/seed/wheels/acquire.py b/src/virtualenv/seed/wheels/acquire.py
index 963077b..3744026 100644
--- a/src/virtualenv/seed/wheels/acquire.py
+++ b/src/virtualenv/seed/wheels/acquire.py
@@ -10,6 +10,7 @@ from virtualenv.util.six import ensure_str
 from virtualenv.util.subprocess import Popen, subprocess
 
 from .bundle import from_bundle
+from .periodic_update import add_wheel_to_update_log
 from .util import Version, Wheel, discover_wheels
 
 
@@ -35,6 +36,8 @@ def get_wheel(distribution, version, for_py_version, search_dirs, download, app_
             to_folder=app_data.house,
             env=env,
         )
+        if wheel is not None and app_data.can_update:
+            add_wheel_to_update_log(wheel, for_py_version, app_data)
 
     return wheel
 
diff --git a/src/virtualenv/seed/wheels/periodic_update.py b/src/virtualenv/seed/wheels/periodic_update.py
index 1a6cbe8..4f0336b 100644
--- a/src/virtualenv/seed/wheels/periodic_update.py
+++ b/src/virtualenv/seed/wheels/periodic_update.py
@@ -82,6 +82,19 @@ def handle_auto_update(distribution, for_py_version, wheel, search_dirs, app_dat
         trigger_update(distribution, for_py_version, wheel, search_dirs, app_data, periodic=True, env=env)
 
 
+def add_wheel_to_update_log(wheel, for_py_version, app_data):
+    embed_update_log = app_data.embed_update_log(wheel.distribution, for_py_version)
+    logging.debug("adding %s information to %s", wheel.name, embed_update_log.file)
+    u_log = UpdateLog.from_dict(embed_update_log.read())
+    if any(version.filename == wheel.name for version in u_log.versions):
+        logging.warning("%s already present in %s", wheel.name, embed_update_log.file)
+        return
+    # we don't need a release date for sources other than "periodic"
+    version = NewVersion(wheel.name, datetime.now(), None, "download")
+    u_log.versions.append(version)  # always write at the end for proper updates
+    embed_update_log.write(u_log.to_dict())
+
+
 DATETIME_FMT = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
@@ -248,23 +261,27 @@ def _run_do_update(app_data, distribution, embed_filename, for_py_version, perio
     embed_update_log = app_data.embed_update_log(distribution, for_py_version)
     u_log = UpdateLog.from_dict(embed_update_log.read())
     now = datetime.now()
+
+    update_versions, other_versions = [], []
+    for version in u_log.versions:
+        if version.source in {"periodic", "manual"}:
+            update_versions.append(version)
+        else:
+            other_versions.append(version)
+
     if periodic:
         source = "periodic"
-        # mark everything not updated manually as source "periodic"
-        for version in u_log.versions:
-            if version.source != "manual":
-                version.source = source
     else:
         source = "manual"
-        # mark everything as source "manual"
-        for version in u_log.versions:
-            version.source = source
+        # mark the most recent one as source "manual"
+        if update_versions:
+            update_versions[0].source = source
 
     if wheel_filename is not None:
         dest = wheelhouse / wheel_filename.name
         if not dest.exists():
             copy2(str(wheel_filename), str(wheelhouse))
-    last, last_version, versions = None, None, []
+    last, last_version, versions, filenames = None, None, [], set()
     while last is None or not last.use(now, ignore_grace_period_ci=True):
         download_time = datetime.now()
         dest = acquire.download_wheel(
@@ -276,13 +293,14 @@ def _run_do_update(app_data, distribution, embed_filename, for_py_version, perio
             to_folder=wheelhouse,
             env=os.environ,
         )
-        if dest is None or (u_log.versions and u_log.versions[0].filename == dest.name):
+        if dest is None or (update_versions and update_versions[0].filename == dest.name):
             break
         release_date = release_date_for_wheel_path(dest.path)
         last = NewVersion(filename=dest.path.name, release_date=release_date, found_date=download_time, source=source)
         logging.info("detected %s in %s", last, datetime.now() - download_time)
         versions.append(last)
-        last_wheel = Wheel(Path(last.filename))
+        filenames.add(last.filename)
+        last_wheel = last.wheel
         last_version = last_wheel.version
         if embed_version is not None:
             if embed_version >= last_wheel.version_tuple:  # stop download if we reach the embed version
@@ -290,7 +308,9 @@ def _run_do_update(app_data, distribution, embed_filename, for_py_version, perio
     u_log.periodic = periodic
     if not u_log.periodic:
         u_log.started = now
-    u_log.versions = versions + u_log.versions
+    # update other_versions by removing version we just found
+    other_versions = [version for version in other_versions if version.filename not in filenames]
+    u_log.versions = versions + update_versions + other_versions
     u_log.completed = datetime.now()
     embed_update_log.write(u_log.to_dict())
     return versions
@@ -395,6 +415,7 @@ def _run_manual_upgrade(app_data, distribution, for_py_version, env):
 
 
 __all__ = (
+    "add_wheel_to_update_log",
     "periodic_update",
     "do_update",
     "manual_upgrade",
diff --git a/tests/unit/seed/wheels/test_acquire.py b/tests/unit/seed/wheels/test_acquire.py
index e255311..dcce0f6 100644
--- a/tests/unit/seed/wheels/test_acquire.py
+++ b/tests/unit/seed/wheels/test_acquire.py
@@ -2,16 +2,25 @@ from __future__ import absolute_import, unicode_literals
 
 import os
 import sys
+from datetime import datetime
 from subprocess import CalledProcessError
 
 import pytest
 
+from virtualenv.app_data import AppDataDiskFolder
+from virtualenv.info import IS_PYPY, PY2
 from virtualenv.seed.wheels.acquire import download_wheel, get_wheel, pip_wheel_env_run
 from virtualenv.seed.wheels.embed import BUNDLE_FOLDER, get_embed_wheel
+from virtualenv.seed.wheels.periodic_update import dump_datetime
 from virtualenv.seed.wheels.util import Wheel, discover_wheels
 from virtualenv.util.path import Path
 
 
+@pytest.fixture(autouse=True)
+def fake_release_date(mocker):
+    mocker.patch("virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path", return_value=None)
+
+
 def test_pip_wheel_env_run_could_not_find(session_app_data, mocker):
     mocker.patch("virtualenv.seed.wheels.acquire.from_bundle", return_value=None)
     with pytest.raises(RuntimeError, match="could not find the embedded pip"):
@@ -74,24 +83,64 @@ def test_download_fails(mocker, for_py_version, session_app_data):
 @pytest.fixture
 def downloaded_wheel(mocker):
     wheel = Wheel.from_path(Path("setuptools-0.0.0-py2.py3-none-any.whl"))
-    mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", return_value=wheel)
-    yield wheel
+    yield wheel, mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", return_value=wheel)
 
 
 @pytest.mark.parametrize("version", ["bundle", "0.0.0"])
-def test_get_wheel_download_called(for_py_version, session_app_data, downloaded_wheel, version):
+def test_get_wheel_download_called(mocker, for_py_version, session_app_data, downloaded_wheel, version):
     distribution = "setuptools"
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
     wheel = get_wheel(distribution, version, for_py_version, [], True, session_app_data, False, os.environ)
     assert wheel is not None
-    assert wheel.name == downloaded_wheel.name
+    assert wheel.name == downloaded_wheel[0].name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
 
 
 @pytest.mark.parametrize("version", ["embed", "pinned"])
-def test_get_wheel_download_not_called(for_py_version, session_app_data, downloaded_wheel, version):
+def test_get_wheel_download_not_called(mocker, for_py_version, session_app_data, downloaded_wheel, version):
     distribution = "setuptools"
     expected = get_embed_wheel(distribution, for_py_version)
     if version == "pinned":
         version = expected.version
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
     wheel = get_wheel(distribution, version, for_py_version, [], True, session_app_data, False, os.environ)
     assert wheel is not None
     assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 0
+    assert write.call_count == 0
+
+
+@pytest.mark.skipif(IS_PYPY and PY2, reason="mocker.spy failing on PyPy 2.x")
+def test_get_wheel_download_cached(tmp_path, freezer, mocker, for_py_version, downloaded_wheel):
+    from virtualenv.app_data.via_disk_folder import JSONStoreDisk
+
+    app_data = AppDataDiskFolder(folder=str(tmp_path))
+    expected = downloaded_wheel[0]
+    write = mocker.spy(JSONStoreDisk, "write")
+    # 1st call, not cached, download is called
+    wheel = get_wheel(expected.distribution, expected.version, for_py_version, [], True, app_data, False, os.environ)
+    assert wheel is not None
+    assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
+    # 2nd call, cached, download is not called
+    wheel = get_wheel(expected.distribution, expected.version, for_py_version, [], True, app_data, False, os.environ)
+    assert wheel is not None
+    assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
+    wrote_json = write.call_args[0][1]
+    assert wrote_json == {
+        "completed": None,
+        "periodic": None,
+        "started": None,
+        "versions": [
+            {
+                "filename": expected.name,
+                "release_date": None,
+                "found_date": dump_datetime(datetime.now()),
+                "source": "download",
+            },
+        ],
+    }
diff --git a/tests/unit/seed/wheels/test_periodic_update.py b/tests/unit/seed/wheels/test_periodic_update.py
index e35dfed..7adb359 100644
--- a/tests/unit/seed/wheels/test_periodic_update.py
+++ b/tests/unit/seed/wheels/test_periodic_update.py
@@ -170,9 +170,9 @@ def test_manual_update_honored(mocker, session_app_data, for_py_version):
     assert str(result.path) == expected_path
 
 
-def wheel_path(wheel, of):
+def wheel_path(wheel, of, pre_release=""):
     new_version = ".".join(str(i) for i in (tuple(sum(x) for x in zip_longest(wheel.version_tuple, of, fillvalue=0))))
-    new_name = wheel.name.replace(wheel.version, new_version)
+    new_name = wheel.name.replace(wheel.version, new_version + pre_release)
     return str(wheel.path.parent / new_name)
 
 
@@ -488,21 +488,25 @@ def test_get_release_fails(mocker, caplog):
     assert repr(exc) in caplog.text
 
 
+def mock_download(mocker, pip_version_remote):
+    def download():
+        index = 0
+        while True:
+            path = pip_version_remote[index]
+            index += 1
+            yield Wheel(Path(path))
+
+    do = download()
+    return mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(do))
+
+
 def test_download_stop_with_embed(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
     pip_version_remote = [wheel_path(wheel, (0, 0, 2)), wheel_path(wheel, (0, 0, 1)), wheel_path(wheel, (-1, 0, 0))]
-    at = {"index": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["index"]]
-            at["index"] += 1
-            yield Wheel(Path(path))
 
-    do = download()
-    download_wheel = mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(do))
+    download_wheel = mock_download(mocker, pip_version_remote)
     url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
 
     last_update = _UP_NOW - timedelta(days=14)
@@ -519,32 +523,71 @@ def test_download_stop_with_embed(tmp_path, mocker, freezer):
     assert write.call_count == 1
 
 
+def test_download_manual_stop_after_one_download(tmp_path, mocker, freezer):
+    freezer.move_to(_UP_NOW)
+    wheel = get_embed_wheel("pip", "3.9")
+    app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
+    pip_version_remote = [wheel_path(wheel, (0, 1, 1))]
+
+    download_wheel = mock_download(mocker, pip_version_remote)
+    url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
+
+    last_update = _UP_NOW - timedelta(days=14)
+    u_log = UpdateLog(started=last_update, completed=last_update, versions=[], periodic=True)
+    read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
+
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
+
+    assert download_wheel.call_count == 1
+    assert url_o.call_count == 2
+    assert read_dict.call_count == 1
+    assert write.call_count == 1
+
+
+def test_download_manual_ignores_pre_release(tmp_path, mocker, freezer):
+    freezer.move_to(_UP_NOW)
+    wheel = get_embed_wheel("pip", "3.9")
+    app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
+    pip_version_remote = [wheel_path(wheel, (0, 0, 1))]
+    pip_version_pre = NewVersion(Path(wheel_path(wheel, (0, 1, 0), "b1")).name, _UP_NOW, None, "downloaded")
+
+    download_wheel = mock_download(mocker, pip_version_remote)
+    url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
+
+    last_update = _UP_NOW - timedelta(days=14)
+    u_log = UpdateLog(started=last_update, completed=last_update, versions=[pip_version_pre], periodic=True)
+    read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
+
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
+
+    assert download_wheel.call_count == 1
+    assert url_o.call_count == 2
+    assert read_dict.call_count == 1
+    assert write.call_count == 1
+    wrote_json = write.call_args[0][0]
+    assert wrote_json["versions"] == [
+        {
+            "filename": Path(pip_version_remote[0]).name,
+            "release_date": None,
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "manual",
+        },
+        pip_version_pre.to_dict(),
+    ]
+
+
 def test_download_periodic_stop_at_first_usable(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
     pip_version_remote = [wheel_path(wheel, (0, 1, 1)), wheel_path(wheel, (0, 1, 0))]
     rel_date_remote = [_UP_NOW - timedelta(days=1), _UP_NOW - timedelta(days=30)]
-    at = {"download": 0, "release_date": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["download"]]
-            at["download"] += 1
-            yield Wheel(Path(path))
 
-    download_gen = download()
-    download_wheel = mocker.patch(
-        "virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(download_gen)
-    )
-
-    def rel_date():
-        while True:
-            value = rel_date_remote[at["release_date"]]
-            at["release_date"] += 1
-            yield value
+    download_wheel = mock_download(mocker, pip_version_remote)
 
-    rel_date_gen = rel_date()
+    rel_date_gen = iter(rel_date_remote)
     release_date = mocker.patch(
         "virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path",
         side_effect=lambda *a, **k: next(rel_date_gen),
@@ -559,51 +602,58 @@ def test_download_periodic_stop_at_first_usable(tmp_path, mocker, freezer):
 
     assert download_wheel.call_count == 2
     assert release_date.call_count == 2
-
     assert read_dict.call_count == 1
     assert write.call_count == 1
 
 
-def test_download_manual_stop_at_first_usable(tmp_path, mocker, freezer):
+def test_download_periodic_stop_at_first_usable_with_previous_minor(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
-    pip_version_remote = [wheel_path(wheel, (0, 1, 1))]
-    rel_date_remote = [_UP_NOW + timedelta(hours=1)]
-    at = {"download": 0, "release_date": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["download"]]
-            at["download"] += 1
-            yield Wheel(Path(path))
-
-    download_gen = download()
-    download_wheel = mocker.patch(
-        "virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(download_gen)
-    )
+    pip_version_remote = [wheel_path(wheel, (0, 1, 1)), wheel_path(wheel, (0, 1, 0)), wheel_path(wheel, (0, -1, 0))]
+    rel_date_remote = [_UP_NOW - timedelta(days=1), _UP_NOW - timedelta(days=30), _UP_NOW - timedelta(days=40)]
+    downloaded_versions = [
+        NewVersion(Path(pip_version_remote[2]).name, rel_date_remote[2], None, "download"),
+        NewVersion(Path(pip_version_remote[0]).name, rel_date_remote[0], None, "download"),
+    ]
 
-    def rel_date():
-        while True:
-            value = rel_date_remote[at["release_date"]]
-            at["release_date"] += 1
-            yield value
+    download_wheel = mock_download(mocker, pip_version_remote)
 
-    rel_date_gen = rel_date()
+    rel_date_gen = iter(rel_date_remote)
     release_date = mocker.patch(
         "virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path",
         side_effect=lambda *a, **k: next(rel_date_gen),
     )
 
     last_update = _UP_NOW - timedelta(days=14)
-    u_log = UpdateLog(started=last_update, completed=last_update, versions=[], periodic=True)
+    u_log = UpdateLog(
+        started=last_update,
+        completed=last_update,
+        versions=downloaded_versions,
+        periodic=True,
+    )
     read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
     write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
 
-    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
-
-    assert download_wheel.call_count == 1
-    assert release_date.call_count == 1
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], True)
 
+    assert download_wheel.call_count == 2
+    assert release_date.call_count == 2
     assert read_dict.call_count == 1
     assert write.call_count == 1
+    wrote_json = write.call_args[0][0]
+    assert wrote_json["versions"] == [
+        {
+            "filename": Path(pip_version_remote[0]).name,
+            "release_date": dump_datetime(rel_date_remote[0]),
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "periodic",
+        },
+        {
+            "filename": Path(pip_version_remote[1]).name,
+            "release_date": dump_datetime(rel_date_remote[1]),
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "periodic",
+        },
+        downloaded_versions[0].to_dict(),
+    ]