From bffc0407fc9ae30df11fa0ac97bd057c383aa8eb Mon Sep 17 00:00:00 2001
From: Darius Makovsky <traveltissues@protonmail.com>
Date: Mon, 28 Oct 2019 09:39:14 +0000
Subject: source.py: Add BST_NO_PRESTAGE_KEY

Extend Source API

Add `_stage_into_cas()` private method. Calls `self.stage` on
a `CasBasedDirectory`.

If the source sets BST_NO_PRESTAGE_KEY then the casdir is recreated from
a stored digest and imported directly in `_stage`.
---
 src/buildstream/source.py           | 48 +++++++++++++++++++++++++++++++++----
 src/buildstream/storage/__init__.py |  1 +
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/src/buildstream/source.py b/src/buildstream/source.py
index a6f7535ac..d37b04ae6 100644
--- a/src/buildstream/source.py
+++ b/src/buildstream/source.py
@@ -1,5 +1,6 @@
 #
 #  Copyright (C) 2016 Codethink Limited
+#  Copyright (C) 2019 Bloomberg Finance LP
 #
 #  This program is free software; you can redistribute it and/or
 #  modify it under the terms of the GNU Lesser General Public
@@ -168,11 +169,12 @@ from typing import Iterable, Iterator, Optional, Tuple, TYPE_CHECKING
 from . import _yaml, utils
 from .node import MappingNode
 from .plugin import Plugin
-from .types import Consistency, SourceRef
+from .types import Consistency, SourceRef, Union
 from ._exceptions import BstError, ImplError, ErrorDomain
 from ._loader.metasource import MetaSource
 from ._projectrefs import ProjectRefStorage
 from ._cachekey import generate_key
+from .storage import CasBasedDirectory
 from .storage import FileBasedDirectory
 from .storage.directory import Directory, VirtualDirectoryError
 
@@ -322,6 +324,14 @@ class Source(Plugin):
     *Since: 1.4*
     """
 
+    BST_NO_PRESTAGE_KEY = False
+    """Whether the source will never have a key prior to staging (a pre-stage
+    key). This is true in the case that the source requires staging in order to
+    efficiently generate a unique key.
+
+    *Since: 1.91.1*
+    """
+
     def __init__(self,
                  context: 'Context',
                  project: 'Project',
@@ -359,6 +369,7 @@ class Source(Plugin):
         self.__mirror_directory = None                  # type: Optional[str]
 
         self._configure(self.__config)
+        self.__digest = None
 
     COMMON_CONFIG_KEYS = ['kind', 'directory']
     """Common source config keys
@@ -478,7 +489,7 @@ class Source(Plugin):
         """
         raise ImplError("Source plugin '{}' does not implement fetch()".format(self.get_kind()))
 
-    def stage(self, directory: str) -> None:
+    def stage(self, directory: Union[str, Directory]) -> None:
         """Stage the sources to a directory
 
         Args:
@@ -705,6 +716,23 @@ class Source(Plugin):
     #############################################################
     #            Private Methods used in BuildStream            #
     #############################################################
+    # Stage files at the localpath into the cascache
+    #
+    # Returns:
+    #   the hash of the cas directory
+    def _stage_into_cas(self) -> str:
+        # FIXME: this should not be called for sources with digests already set
+        # since they will already have been staged into the cache. However,
+        # _get_unique_key is sometimes called outside of _generate_key
+        if self.__digest is None:
+            cas_dir = CasBasedDirectory(self._get_context().get_cascache())
+            self.stage(cas_dir)
+            digest = cas_dir._get_digest()
+            self.__digest = digest
+        else:
+            # XXX: an assignment to please mypy
+            digest = self.__digest
+        return digest.hash
 
     # Wrapper around preflight() method
     #
@@ -760,7 +788,14 @@ class Source(Plugin):
     def _stage(self, directory):
         directory = self.__ensure_directory(directory)
 
-        self.stage(directory)
+        if self.BST_NO_PRESTAGE_KEY:
+            # _get_unique_key should be called before _stage
+            assert self.__digest is not None
+            cas_dir = CasBasedDirectory(self._get_context().get_cascache(),
+                                        digest=self.__digest)
+            directory.import_files(cas_dir)
+        else:
+            self.stage(directory)
 
     # Wrapper for init_workspace()
     def _init_workspace(self, directory):
@@ -778,7 +813,10 @@ class Source(Plugin):
     def _get_unique_key(self):
         key = {}
         key['directory'] = self.__directory
-        key['unique'] = self.get_unique_key()  # pylint: disable=assignment-from-no-return
+        if self.BST_NO_PRESTAGE_KEY:
+            key['unique'] = self._stage_into_cas()
+        else:
+            key['unique'] = self.get_unique_key()  # pylint: disable=assignment-from-no-return
         return key
 
     # _project_refs():
@@ -1077,7 +1115,7 @@ class Source(Plugin):
         self.__key = generate_key(keys)
 
         sourcecache = self._get_context().sourcecache
-        if self.get_kind() == 'workspace' and not sourcecache.contains(self):
+        if self.BST_NO_PRESTAGE_KEY and not sourcecache.contains(self):
             sourcecache.commit(self, [])
 
     @property
diff --git a/src/buildstream/storage/__init__.py b/src/buildstream/storage/__init__.py
index 33424ac8d..5571cd807 100644
--- a/src/buildstream/storage/__init__.py
+++ b/src/buildstream/storage/__init__.py
@@ -19,4 +19,5 @@
 #        Jim MacArthur <jim.macarthur@codethink.co.uk>
 
 from ._filebaseddirectory import FileBasedDirectory
+from ._casbaseddirectory import CasBasedDirectory
 from .directory import Directory
-- 
cgit v1.2.1