casserver.py: Use quota instead of headroom

3 files changed, 21 insertions, 67 deletions

diff --git a/src/buildstream/_cas/casserver.py b/src/buildstream/_cas/casserver.py
index 71c4efde2..5a4c2b7ac 100644
--- a/src/buildstream/_cas/casserver.py
+++ b/src/buildstream/_cas/casserver.py
@@ -56,10 +56,8 @@ _MAX_PAYLOAD_BYTES = 1024 * 1024
 #     enable_push (bool): Whether to allow blob uploads and artifact updates
 #
 @contextmanager
-def create_server(repo, *, enable_push,
-                  max_head_size=int(10e9),
-                  min_head_size=int(2e9)):
-    cas = CASCache(os.path.abspath(repo))
+def create_server(repo, *, enable_push, quota):
+    cas = CASCache(os.path.abspath(repo), cache_quota=quota, protect_session_blobs=False)
 
     try:
         artifactdir = os.path.join(os.path.abspath(repo), 'artifacts', 'refs')
@@ -109,23 +107,19 @@ def create_server(repo, *, enable_push,
 @click.option('--client-certs', help="Public client certificates for TLS (PEM-encoded)")
 @click.option('--enable-push', default=False, is_flag=True,
               help="Allow clients to upload blobs and update artifact cache")
-@click.option('--head-room-min', type=click.INT,
-              help="Disk head room minimum in bytes",
-              default=2e9)
-@click.option('--head-room-max', type=click.INT,
-              help="Disk head room maximum in bytes",
+@click.option('--quota', type=click.INT,
+              help="Maximum disk usage in bytes",
               default=10e9)
 @click.argument('repo')
 def server_main(repo, port, server_key, server_cert, client_certs, enable_push,
-                head_room_min, head_room_max):
+                quota):
     # Handle SIGTERM by calling sys.exit(0), which will raise a SystemExit exception,
     # properly executing cleanup code in `finally` clauses and context managers.
     # This is required to terminate buildbox-casd on SIGTERM.
     signal.signal(signal.SIGTERM, lambda signalnum, frame: sys.exit(0))
 
     with create_server(repo,
-                       max_head_size=head_room_max,
-                       min_head_size=head_room_min,
+                       quota=quota,
                        enable_push=enable_push) as server:
 
         use_tls = bool(server_key)
diff --git a/tests/frontend/push.py b/tests/frontend/push.py
index e0a6c1e99..9c3947c2a 100644
--- a/tests/frontend/push.py
+++ b/tests/frontend/push.py
@@ -287,25 +287,22 @@ def test_push_after_pull(cli, tmpdir, datafiles):
 # the least recently pushed artifact is deleted in order to make room for
 # the incoming artifact.
 @pytest.mark.datafiles(DATA_DIR)
-@pytest.mark.xfail()
 def test_artifact_expires(cli, datafiles, tmpdir):
     project = str(datafiles)
     element_path = 'elements'
 
     # Create an artifact share (remote artifact cache) in the tmpdir/artifactshare
-    # Mock a file system with 12 MB free disk space
+    # Set a 22 MB quota
     with create_artifact_share(os.path.join(str(tmpdir), 'artifactshare'),
-                               min_head_size=int(2e9),
-                               max_head_size=int(2e9),
-                               total_space=int(10e9), free_space=(int(12e6) + int(2e9))) as share:
+                               quota=int(22e6)) as share:
 
         # Configure bst to push to the cache
         cli.configure({
             'artifacts': {'url': share.repo, 'push': True},
         })
 
-        # Create and build an element of 5 MB
-        create_element_size('element1.bst', project, element_path, [], int(5e6))
+        # Create and build an element of 15 MB
+        create_element_size('element1.bst', project, element_path, [], int(15e6))
         result = cli.run(project=project, args=['build', 'element1.bst'])
         result.assert_success()
 
@@ -343,7 +340,6 @@ def test_artifact_expires(cli, datafiles, tmpdir):
 # Test that a large artifact, whose size exceeds the quota, is not pushed
 # to the remote share
 @pytest.mark.datafiles(DATA_DIR)
-@pytest.mark.xfail()
 def test_artifact_too_large(cli, datafiles, tmpdir):
     project = str(datafiles)
     element_path = 'elements'
@@ -351,7 +347,7 @@ def test_artifact_too_large(cli, datafiles, tmpdir):
     # Create an artifact share (remote cache) in tmpdir/artifactshare
     # Mock a file system with 5 MB total space
     with create_artifact_share(os.path.join(str(tmpdir), 'artifactshare'),
-                               total_space=int(5e6) + int(2e9)) as share:
+                               quota=int(5e6)) as share:
 
         # Configure bst to push to the remote cache
         cli.configure({
@@ -380,29 +376,26 @@ def test_artifact_too_large(cli, datafiles, tmpdir):
 
 # Test that when an element is pulled recently, it is not considered the LRU element.
 @pytest.mark.datafiles(DATA_DIR)
-@pytest.mark.xfail()
 def test_recently_pulled_artifact_does_not_expire(cli, datafiles, tmpdir):
     project = str(datafiles)
     element_path = 'elements'
 
     # Create an artifact share (remote cache) in tmpdir/artifactshare
-    # Mock a file system with 12 MB free disk space
+    # Set a 22 MB quota
     with create_artifact_share(os.path.join(str(tmpdir), 'artifactshare'),
-                               min_head_size=int(2e9),
-                               max_head_size=int(2e9),
-                               total_space=int(10e9), free_space=(int(12e6) + int(2e9))) as share:
+                               quota=int(22e6)) as share:
 
         # Configure bst to push to the cache
         cli.configure({
             'artifacts': {'url': share.repo, 'push': True},
         })
 
-        # Create and build 2 elements, each of 5 MB.
+        # Create and build 2 elements, one 5 MB and one 15 MB.
         create_element_size('element1.bst', project, element_path, [], int(5e6))
         result = cli.run(project=project, args=['build', 'element1.bst'])
         result.assert_success()
 
-        create_element_size('element2.bst', project, element_path, [], int(5e6))
+        create_element_size('element2.bst', project, element_path, [], int(15e6))
         result = cli.run(project=project, args=['build', 'element2.bst'])
         result.assert_success()
 
diff --git a/tests/testutils/artifactshare.py b/tests/testutils/artifactshare.py
index 132b6ac4f..7d5faeb66 100644
--- a/tests/testutils/artifactshare.py
+++ b/tests/testutils/artifactshare.py
@@ -20,18 +20,12 @@ from buildstream._protos.buildstream.v2 import artifact_pb2
 #
 # Args:
 #    directory (str): The base temp directory for the test
-#    total_space (int): Mock total disk space on artifact server
-#    free_space (int): Mock free disk space on artifact server
+#    cache_quota (int): Maximum amount of disk space to use
 #    casd (bool): Allow write access via casd
 #
 class ArtifactShare():
 
-    def __init__(self, directory, *,
-                 total_space=None,
-                 free_space=None,
-                 min_head_size=int(2e9),
-                 max_head_size=int(10e9),
-                 casd=False):
+    def __init__(self, directory, *, quota=None, casd=False):
 
         # The working directory for the artifact share (in case it
         # needs to do something outside of its backend's storage folder).
@@ -50,11 +44,7 @@ class ArtifactShare():
 
         self.cas = CASCache(self.repodir, casd=casd)
 
-        self.total_space = total_space
-        self.free_space = free_space
-
-        self.max_head_size = max_head_size
-        self.min_head_size = min_head_size
+        self.quota = quota
 
         q = Queue()
 
@@ -80,15 +70,8 @@ class ArtifactShare():
             pytest_cov.embed.cleanup_on_sigterm()
 
         try:
-            # Optionally mock statvfs
-            if self.total_space:
-                if self.free_space is None:
-                    self.free_space = self.total_space
-                os.statvfs = self._mock_statvfs
-
             with create_server(self.repodir,
-                               max_head_size=self.max_head_size,
-                               min_head_size=self.min_head_size,
+                               quota=self.quota,
                                enable_push=True) as server:
                 port = server.add_insecure_port('localhost:0')
 
@@ -182,30 +165,14 @@ class ArtifactShare():
 
         shutil.rmtree(self.directory)
 
-    def _mock_statvfs(self, _path):
-        repo_size = 0
-        for root, _, files in os.walk(self.repodir):
-            for filename in files:
-                repo_size += os.path.getsize(os.path.join(root, filename))
-
-        return statvfs_result(f_blocks=self.total_space,
-                              f_bfree=self.free_space - repo_size,
-                              f_bavail=self.free_space - repo_size,
-                              f_bsize=1)
-
 
 # create_artifact_share()
 #
 # Create an ArtifactShare for use in a test case
 #
 @contextmanager
-def create_artifact_share(directory, *, total_space=None, free_space=None,
-                          min_head_size=int(2e9),
-                          max_head_size=int(10e9),
-                          casd=False):
-    share = ArtifactShare(directory, total_space=total_space, free_space=free_space,
-                          min_head_size=min_head_size, max_head_size=max_head_size,
-                          casd=casd)
+def create_artifact_share(directory, *, quota=None, casd=False):
+    share = ArtifactShare(directory, quota=quota, casd=casd)
     try:
         yield share
     finally: