TEMP: check that cache size only written when expectedaevri/assert_no_unexpected_size_writes

4 files changed, 50 insertions, 4 deletions

diff --git a/buildstream/_cas/cascache.py b/buildstream/_cas/cascache.py
index 04a09299a..ca7a370b2 100644
--- a/buildstream/_cas/cascache.py
+++ b/buildstream/_cas/cascache.py
@@ -42,6 +42,32 @@ _BUFFER_SIZE = 65536
 
 CACHE_SIZE_FILE = "cache_size"
 
+# Note that we will only allow writing in the main process, keep track of the
+# allowed pid.
+_CACHE_SIZE_WRITE_ALLOWED_PID = os.getpid()
+_CACHE_SIZE_WRITE_ALLOWED = False
+
+
+def disallow_cache_size_write_context_pid():
+    global _CACHE_SIZE_WRITE_ALLOWED_PID  # pylint: disable=global-statement
+    _CACHE_SIZE_WRITE_ALLOWED_PID = None
+
+
+def allow_cache_size_write_context_pid():
+    global _CACHE_SIZE_WRITE_ALLOWED_PID  # pylint: disable=global-statement
+    _CACHE_SIZE_WRITE_ALLOWED_PID = os.getpid()
+
+
+@contextlib.contextmanager
+def allow_cache_size_write_context(reset_pid=False):
+    global _CACHE_SIZE_WRITE_ALLOWED  # pylint: disable=global-statement
+    assert _CACHE_SIZE_WRITE_ALLOWED is not None
+    _CACHE_SIZE_WRITE_ALLOWED = True
+    try:
+        yield
+    finally:
+        _CACHE_SIZE_WRITE_ALLOWED = False
+
 
 # CASCacheUsage
 #
@@ -1186,6 +1212,8 @@ class CASQuota:
     #
     def _write_cache_size(self, size):
         assert isinstance(size, int)
+        assert _CACHE_SIZE_WRITE_ALLOWED_PID == os.getpid()
+        assert _CACHE_SIZE_WRITE_ALLOWED
         size_file_path = os.path.join(self.casdir, CACHE_SIZE_FILE)
         with utils.save_file_atomic(size_file_path, "w") as f:
             f.write(str(size))
diff --git a/buildstream/_scheduler/jobs/cachesizejob.py b/buildstream/_scheduler/jobs/cachesizejob.py
index 5f27b7fc1..a371680d8 100644
--- a/buildstream/_scheduler/jobs/cachesizejob.py
+++ b/buildstream/_scheduler/jobs/cachesizejob.py
@@ -32,7 +32,9 @@ class CacheSizeJob(Job):
 
     def parent_complete(self, status, result):
         if status == JobStatus.OK:
-            self._casquota.set_cache_size(result)
+            from ..._cas.cascache import allow_cache_size_write_context
+            with allow_cache_size_write_context():
+                self._casquota.set_cache_size(result)
 
         if self._complete_cb:
             self._complete_cb(status, result)
diff --git a/buildstream/_scheduler/jobs/cleanupjob.py b/buildstream/_scheduler/jobs/cleanupjob.py
index 9610d53f8..58e37493b 100644
--- a/buildstream/_scheduler/jobs/cleanupjob.py
+++ b/buildstream/_scheduler/jobs/cleanupjob.py
@@ -32,20 +32,34 @@ class CleanupJob(Job):
         def progress():
             self.send_message('update-cache-size',
                               self._casquota.get_cache_size())
-        return self._casquota.clean(progress)
+
+        from ..._cas.cascache import allow_cache_size_write_context
+        from ..._cas.cascache import allow_cache_size_write_context_pid
+        from ..._cas.cascache import disallow_cache_size_write_context_pid
+        allow_cache_size_write_context_pid()
+        with allow_cache_size_write_context():
+            result = self._casquota.clean(progress)
+        disallow_cache_size_write_context_pid()
+        return result
 
     def handle_message(self, message_type, message):
         # Update the cache size in the main process as we go,
         # this provides better feedback in the UI.
         if message_type == 'update-cache-size':
             self._casquota.set_cache_size(message)
+            self._casquota.set_cache_size(message, write_to_disk=False)
+            from ..._cas.cascache import allow_cache_size_write_context
+            with allow_cache_size_write_context():
+                self._casquota.set_cache_size(message)
             return True
 
         return False
 
     def parent_complete(self, status, result):
         if status == JobStatus.OK:
-            self._casquota.set_cache_size(result)
+            from ..._cas.cascache import allow_cache_size_write_context
+            with allow_cache_size_write_context():
+                self._casquota.set_cache_size(result)
 
         if self._complete_cb:
             self._complete_cb(status, result)
diff --git a/buildstream/_scheduler/queues/buildqueue.py b/buildstream/_scheduler/queues/buildqueue.py
index 60ec19ff4..977933b4f 100644
--- a/buildstream/_scheduler/queues/buildqueue.py
+++ b/buildstream/_scheduler/queues/buildqueue.py
@@ -117,4 +117,6 @@ class BuildQueue(Queue):
         #        failed build also grows the artifact cache size.
         #
         if status == JobStatus.OK:
-            self._check_cache_size(job, element, result)
+            from ..._cas.cascache import allow_cache_size_write_context
+            with allow_cache_size_write_context():
+                self._check_cache_size(job, element, result)