Add a shared code cache test.

Fix a long standing bug in calculating the shared code hash along the
way!

pex/pex_builder.py

@@ -11,7 +11,7 @@
 from textwrap import dedent
 from zipimport import ZipImportError
 
-from pex import pex_warnings
+from pex import layout, pex_warnings
 from pex.atomic_directory import atomic_directory
 from pex.cache.dirs import CacheDir
 from pex.common import (
@@ -461,7 +461,9 @@ def _precompile_source(self):
             self._chroot.touch(compiled, label="bytecode")
 
     def _prepare_code(self):
-        self._pex_info.code_hash = CacheHelper.pex_code_hash(self._chroot.path())
+        self._pex_info.code_hash = CacheHelper.pex_code_hash(
+            self._chroot.path(), exclude_dirs=(layout.BOOTSTRAP_DIR, layout.DEPS_DIR)
+        )
         self._pex_info.pex_hash = hashlib.sha1(self._pex_info.dump().encode("utf-8")).hexdigest()
         self._chroot.write(self._pex_info.dump().encode("utf-8"), PexInfo.PATH, label="manifest")
 

pex/util.py

@@ -22,7 +22,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import IO, Any, Callable, Iterator, Optional, Text
+    from typing import IO, Any, Callable, Container, Iterator, Optional, Text
 
     from pex.hashing import Hasher
 
@@ -80,14 +80,18 @@ def hash(cls, path, digest=None, hasher=sha1):
         return digest.hexdigest()
 
     @classmethod
-    def pex_code_hash(cls, directory):
-        # type: (str) -> str
+    def pex_code_hash(
+        cls,
+        directory,
+        exclude_dirs=(),  # type: Container[str]
+    ):
+        # type: (...) -> str
         """Return a reproducible hash of the contents of a loose PEX; excluding all `.pyc` files."""
         digest = hashlib.sha1()
         hashing.dir_hash(
             directory=directory,
             digest=digest,
-            dir_filter=lambda d: not is_pyc_dir(d),
+            dir_filter=lambda d: not is_pyc_dir(d) and d not in exclude_dirs,
             file_filter=lambda file_path: not is_pyc_file(file_path)
             and not file_path.startswith("."),
         )

tests/integration/cli/commands/test_cache_prune.py

@@ -15,7 +15,7 @@
 import pytest
 
 from pex.cache import access
-from pex.cache.dirs import BootstrapDir, CacheDir, InstalledWheelDir, UnzipDir
+from pex.cache.dirs import BootstrapDir, CacheDir, InstalledWheelDir, UnzipDir, UserCodeDir
 from pex.cli.commands.cache.du import DiskUsage
 from pex.common import safe_open
 from pex.pep_503 import ProjectName
@@ -123,24 +123,33 @@ class AnsicolorsPex(object):
     path = attr.ib()  # type: str
 
 
-def create_ansicolors_pex(
-    tmpdir,  # type: Tempdir
-    *extra_args  # type: str
-):
-    # type: (...) -> AnsicolorsPex
-    pex = tmpdir.join("ansicolors.pex")
-    with safe_open(tmpdir.join("src", "app.py"), "w") as fp:
+def write_app_py(path):
+    # type: (str) -> None
+    with safe_open(path, "w") as fp:
         fp.write(
             dedent(
                 """\
-                import colors
+                try:
+                    from colors import green
+                except ImportError:
+                    def green(text):
+                        return text
 
 
                 if __name__ == "__main__":
-                    print(colors.green("Hello Cache!"))
+                    print(green("Hello Cache!"))
                 """
             )
         )
+
+
+def create_ansicolors_pex(
+    tmpdir,  # type: Tempdir
+    *extra_args  # type: str
+):
+    # type: (...) -> AnsicolorsPex
+    pex = tmpdir.join("ansicolors.pex")
+    write_app_py(tmpdir.join("src", "app.py"))
     run_pex_command(
         args=["ansicolors==1.1.8", "-D", "src", "-m" "app", "-o", pex] + list(extra_args),
         cwd=tmpdir.path,
@@ -269,6 +278,43 @@ def test_zipapp_prune_shared_bootstrap(
     assert_installed_wheels(expected_pip_wheels())
 
 
+def test_zipapp_prune_shared_code(
+    ansicolors_zipapp_pex,  # type: AnsicolorsPex
+    tmpdir,  # type: Tempdir
+):
+    # type: (...) -> None
+
+    execute_ansicolors_pex(ansicolors_zipapp_pex)
+    code_hash = PexInfo.from_pex(ansicolors_zipapp_pex.path).code_hash
+    assert code_hash is not None
+
+    all_user_code = list(UserCodeDir.iter_all())
+    assert len(all_user_code) == 1
+    assert code_hash == all_user_code[0].code_hash
+
+    write_app_py(tmpdir.join("app.py"))
+    no_colors_pex = tmpdir.join("no-colors.pex")
+    run_pex_command(
+        args=["-M" "app", "-m", "app", "-o", no_colors_pex], cwd=tmpdir.path
+    ).assert_success()
+    assert b"Hello Cache!\n" == subprocess.check_output(args=[no_colors_pex])
+    assert all_user_code == list(
+        UserCodeDir.iter_all()
+    ), "Expected the shared code cache to be re-used since the code is the same for both PEXes."
+
+    set_last_access_one_day_ago(ansicolors_zipapp_pex.path)
+    run_pex3("cache", "prune", "--older-than", "1 hour").assert_success()
+    assert all_user_code == list(
+        UserCodeDir.iter_all()
+    ), "Expected the shared code cache to be un-pruned since no_colors_pex still needs it."
+
+    run_pex3("cache", "prune", "--older-than", "0 seconds").assert_success()
+    assert len(list(UserCodeDir.iter_all())) == 0, (
+        "Expected the shared code cache to be pruned since the last remaining user, no_colors_pex,"
+        "is now pruned."
+    )
+
+
 @attr.s(frozen=True)
 class CowsayPex(object):
     path = attr.ib()  # type: str