[buildstream] 02/02: cachequeryqueue.py: support querying sources when artifacts are cached

[ak...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

akitouni pushed a commit to branch abderrahim/cache-query
in repository https://gitbox.apache.org/repos/asf/buildstream.git

commit a920598d045eb8a7e6ef859632075cb74b98870b
Author: Abderrahim Kitouni <ak...@gnome.org>
AuthorDate: Fri Sep 9 08:55:39 2022 +0100

    cachequeryqueue.py: support querying sources when artifacts are cached
---
 src/buildstream/_scheduler/queues/cachequeryqueue.py | 19 +++++++++++++++++--
 src/buildstream/_stream.py                           |  7 ++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/buildstream/_scheduler/queues/cachequeryqueue.py b/src/buildstream/_scheduler/queues/cachequeryqueue.py
index b650a91b8..0b1b8b501 100644
--- a/src/buildstream/_scheduler/queues/cachequeryqueue.py
+++ b/src/buildstream/_scheduler/queues/cachequeryqueue.py
@@ -28,18 +28,27 @@ class CacheQueryQueue(Queue):
     complete_name = "Cache queried"
     resources = [ResourceType.PROCESS, ResourceType.CACHE]
 
-    def __init__(self, scheduler, *, sources=False):
+    def __init__(self, scheduler, *, sources=False, sources_if_cached=False):
         super().__init__(scheduler)
 
         self._sources = sources
+        self._sources_if_cached = sources_if_cached
 
     def get_process_func(self):
-        if not self._sources:
+        if self._sources_if_cached:
+            return CacheQueryQueue._query_artifacts_and_sources
+        elif not self._sources:
             return CacheQueryQueue._query_artifacts_or_sources
         else:
             return CacheQueryQueue._query_sources
 
     def status(self, element):
+        if element._can_query_cache():
+            # Cache status already available.
+            # This is the case for artifact elements, which load the
+            # artifact early on.
+            return QueueStatus.SKIP
+
         if not element._get_cache_key(strength=_KeyStrength.WEAK):
             # Strict and weak cache keys are unavailable if the element or
             # a dependency has an unresolved source
@@ -61,6 +70,12 @@ class CacheQueryQueue(Queue):
         if not element._can_query_cache() or not element._cached_success():
             element._query_source_cache()
 
+    @staticmethod
+    def _query_artifacts_and_sources(element):
+        element.warn("querying artifact and sources")
+        element._load_artifact(pull=False)
+        element._query_source_cache()
+
     @staticmethod
     def _query_sources(element):
         element._query_source_cache()
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index efad29313..07f87822f 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -209,7 +209,12 @@ class Stream:
             if self._context.remote_cache_spec:
                 # Parallelize cache queries if a remote cache is configured
                 self._reset()
-                self._add_queue(CacheQueryQueue(self._scheduler, sources=only_sources), track=True)
+                self._add_queue(
+                    CacheQueryQueue(
+                        self._scheduler, sources=only_sources, sources_if_cached=sources_of_cached_elements
+                    ),
+                    track=True,
+                )
                 self._enqueue_plan(plan)
                 self._run()
             else: