Remove artifact extracts when artifact expires in cachevalentindavid/extract-expiry

Fixes #561

2 files changed, 56 insertions, 0 deletions

diff --git a/buildstream/_artifactcache/cascache.py b/buildstream/_artifactcache/cascache.py
index 074899d81..a1634fcf6 100644
--- a/buildstream/_artifactcache/cascache.py
+++ b/buildstream/_artifactcache/cascache.py
@@ -30,6 +30,8 @@ from urllib.parse import urlparse
 
 import grpc
 
+from .. import _yaml
+
 from .._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
 from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
 from .._protos.buildstream.v2 import buildstream_pb2, buildstream_pb2_grpc
@@ -524,6 +526,25 @@ class CASCache(ArtifactCache):
     #
     def remove(self, ref, *, defer_prune=False):
 
+        # Remove extract if not used by other ref
+        tree = self.resolve_ref(ref)
+        ref_name, ref_hash = os.path.split(ref)
+        extract = os.path.join(self.extractdir, ref_name, tree.hash)
+        keys_file = os.path.join(extract, 'meta', 'keys.yaml')
+        if os.path.exists(keys_file):
+            keys_meta = _yaml.load(keys_file)
+            keys = [keys_meta['strong'], keys_meta['weak']]
+            remove_extract = True
+            for other_hash in keys:
+                if other_hash == ref_hash:
+                    continue
+                remove_extract = False
+                break
+
+            if remove_extract:
+                utils._force_rmtree(extract)
+
+        # Remove cache ref
         refpath = self._refpath(ref)
         if not os.path.exists(refpath):
             raise ArtifactError("Could not find artifact for ref '{}'".format(ref))
diff --git a/tests/artifactcache/expiry.py b/tests/artifactcache/expiry.py
index 9b77fbe3e..6190ae814 100644
--- a/tests/artifactcache/expiry.py
+++ b/tests/artifactcache/expiry.py
@@ -268,3 +268,38 @@ def test_invalid_cache_quota(cli, datafiles, tmpdir, quota, success):
         res.assert_success()
     else:
         res.assert_main_error(ErrorDomain.LOAD, LoadErrorReason.INVALID_DATA)
+
+
+@pytest.mark.datafiles(DATA_DIR)
+def test_extract_expiry(cli, datafiles, tmpdir):
+    project = os.path.join(datafiles.dirname, datafiles.basename)
+    element_path = 'elements'
+
+    cli.configure({
+        'cache': {
+            'quota': 10000000,
+        }
+    })
+
+    create_element_size('target.bst', project, element_path, [], 6000000)
+    res = cli.run(project=project, args=['build', 'target.bst'])
+    res.assert_success()
+    assert cli.get_element_state(project, 'target.bst') == 'cached'
+
+    # Force creating extract
+    res = cli.run(project=project, args=['checkout', 'target.bst', os.path.join(str(tmpdir), 'checkout')])
+    res.assert_success()
+
+    extractdir = os.path.join(project, 'cache', 'artifacts', 'extract', 'test', 'target')
+    extracts = os.listdir(extractdir)
+    assert(len(extracts) == 1)
+    extract = os.path.join(extractdir, extracts[0])
+
+    # Remove target.bst from artifact cache
+    create_element_size('target2.bst', project, element_path, [], 6000000)
+    res = cli.run(project=project, args=['build', 'target2.bst'])
+    res.assert_success()
+    assert cli.get_element_state(project, 'target.bst') != 'cached'
+
+    # Now the extract should be removed.
+    assert not os.path.exists(extract)