Add an artifact cache which uses OSTree

4 files changed, 232 insertions, 5 deletions

diff --git a/morphlib/__init__.py b/morphlib/__init__.py
index e2641402..79e829a4 100644
--- a/morphlib/__init__.py
+++ b/morphlib/__init__.py
@@ -73,6 +73,7 @@ import morphology
 import morphloader
 import morphset
 import ostree
+import ostreeartifactcache
 import remoteartifactcache
 import remoterepocache
 import repoaliasresolver
diff --git a/morphlib/builder.py b/morphlib/builder.py
index 3e6e44d2..0c681353 100644
--- a/morphlib/builder.py
+++ b/morphlib/builder.py
@@ -125,11 +125,7 @@ def ldconfig(runcmd, rootdir):  # pragma: no cover
 def download_depends(constituents, lac, rac, metadatas=None):
     for constituent in constituents:
         if not lac.has(constituent):
-            source = rac.get(constituent)
-            target = lac.put(constituent)
-            shutil.copyfileobj(source, target)
-            target.close()
-            source.close()
+            lac.copy_from_remote(constituent, rac)
         if metadatas is not None:
             for metadata in metadatas:
                 if not lac.has_artifact_metadata(constituent, metadata):
diff --git a/morphlib/ostreeartifactcache.py b/morphlib/ostreeartifactcache.py
new file mode 100644
index 00000000..fdb7cb5d
--- /dev/null
+++ b/morphlib/ostreeartifactcache.py
@@ -0,0 +1,229 @@
+# Copyright (C) 2015  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+
+import collections
+import logging
+import os
+import shutil
+import tarfile
+import tempfile
+
+import cliapp
+from gi.repository import GLib
+
+import morphlib
+from morphlib.artifactcachereference import ArtifactCacheReference
+
+class OSTreeArtifactCache(object):
+    """Class to provide the artifact cache API using an OSTree repo."""
+
+    def __init__(self, cachedir):
+        repo_dir = os.path.join(cachedir, 'repo')
+        self.repo = morphlib.ostree.OSTreeRepo(repo_dir)
+        self.cachedir = cachedir
+
+    def _get_file_from_remote(self, artifact, remote, metadata_name=None):
+        if metadata_name:
+            handle = remote.get_artifact_metadata(artifact, metadata_name)
+        else:
+            handle = remote.get(artifact)
+        fd, path = tempfile.mkstemp()
+        with open(path, 'w+') as temp:
+            shutil.copyfileobj(handle, temp)
+        return path
+
+    def _get_artifact_cache_name(self, artifact):
+        logging.debug('LAC: %s' % artifact.basename())
+        cache_key, kind, name = artifact.basename().split('.', 2)
+        suffix = name.split('-')[-1]
+        return '%s-%s' % (cache_key, suffix)
+
+    def put(self, directory, artifact):
+        """Commit the contents of 'directory' to the repo.
+        
+        This uses the artifact name and cache key to create the ref, so the
+        contents of directory should be the contents of the artifact.
+        
+        """
+        ref = self._get_artifact_cache_name(artifact)
+        subject = artifact.name
+        try:
+            logging.debug('Committing %s to artifact cache at %s.' %
+                          (subject, ref))
+            self.repo.commit(subject, directory, ref)
+        except GLib.GError as e:
+            logging.debug('OSTree raised an exception: %s' % e)
+            raise cliapp.AppException('Failed to commit %s to artifact '
+                                      'cache.' % ref)
+
+    def put_non_ostree_artifact(self, artifact, location, metadata_name=None):
+        """Store a single file in the artifact cachedir."""
+        if metadata_name:
+            filename = self._artifact_metadata_filename(artifact,
+                                                        metadata_name)
+        else:
+            filename = self.artifact_filename(artifact)
+        shutil.copy(location, filename)
+        os.remove(location)
+
+    def copy_from_remote(self, artifact, remote):
+        """Get 'artifact' from remote artifact cache and store it locally."""
+        if remote.method == 'tarball':
+            logging.debug('Downloading artifact tarball for %s.' %
+                          artifact.name)
+            location = self._get_file_from_remote(artifact, remote)
+            try:
+                tempdir = tempfile.mkdtemp()
+                with tarfile.open(name=location) as tf:
+                    tf.extractall(path=tempdir)
+                try:
+                    self.put(tempdir, artifact)
+                finally:
+                    os.remove(location)
+                    shutil.rmtree(tempdir)
+            except tarfile.ReadError:
+                # Reading the artifact as a tarball failed, so it must be a
+                # single file (for example a stratum artifact).
+                self.put_non_ostree_artifact(artifact, location)
+
+        elif remote.method == 'ostree':
+            logging.debug('Pulling artifact for %s from remote.' %
+                          artifact.basename())
+            try:
+                ref = self._get_artifact_cache_name(artifact)
+            except Exception:
+                # if we can't split the name properly, we must want metadata
+                a, name = artifact.basename().split('.', 1)
+                location = self._get_file_from_remote(
+                    ArtifactCacheReference(a), remote, name)
+                self.put_non_ostree_artifact(artifact, location, name)
+                return
+
+            if artifact.basename().split('.', 2)[1] == 'stratum':
+                location = self._get_file_from_remote(artifact, remote)
+                self.put_non_ostree_artifact(artifact, location)
+                return
+
+            try:
+                if not self.repo.has_remote(remote.name):
+                    self.repo.add_remote(remote.name, remote.ostree_url)
+                self.repo.pull([ref], remote.name)
+            except GLib.GError as e:
+                logging.debug('OSTree raised an exception: %s' % e)
+                raise cliapp.AppException('Failed to pull %s from remote '
+                                          'cache.' % ref)
+
+    def get(self, artifact, directory=None):
+        """Checkout an artifact from the repo and return its location."""
+        cache_key, kind, name = artifact.basename().split('.', 2)
+        if kind == 'stratum':
+            return self.artifact_filename(artifact)
+        if directory is None:
+            directory = tempfile.mkdtemp()
+        ref = self._get_artifact_cache_name(artifact)
+        try:
+            self.repo.checkout(ref, directory)
+            self.repo.touch_ref(ref)
+        except GLib.GError as e:
+            logging.debug('OSTree raised an exception: %s' % e)
+            raise cliapp.AppException('Failed to checkout %s from artifact '
+                                      'cache.' % ref)
+        return directory
+
+    def list_contents(self):
+        """Return the set of sources cached and related information.
+
+           returns a [(cache_key, set(artifacts), last_used)]
+
+        """
+        CacheInfo = collections.namedtuple('CacheInfo', ('artifacts', 'mtime'))
+        contents = collections.defaultdict(lambda: CacheInfo(set(), 0))
+        for ref in self.repo.list_refs():
+            cachekey = ref[:63]
+            artifact = ref[65:]
+            artifacts, max_mtime = contents[cachekey]
+            artifacts.add(artifact)
+            ref_filename = os.path.join(self.repo.refsdir(), ref)
+            mtime = os.path.getmtime(ref_filename)
+            contents[cachekey] = CacheInfo(artifacts, max(max_mtime, mtime))
+        return ((cache_key, info.artifacts, info.mtime)
+                for cache_key, info in contents.iteritems())
+
+    def remove(self, cachekey):
+        """Remove all artifacts associated with the given cachekey."""
+        for ref in (r for r in self.repo.list_refs()
+                if r.startswith(cachekey)):
+            self.repo.delete_ref(ref)
+
+    def prune(self):
+        """Delete orphaned objects in the repo."""
+        self.repo.prune()
+
+    def has(self, artifact):
+        cachekey, kind, name = artifact.basename().split('.', 2)
+        logging.debug('OSTreeArtifactCache: got %s, %s, %s' %
+                      (cachekey, kind, name))
+        if self._get_artifact_cache_name(artifact) in self.repo.list_refs():
+            self.repo.touch_ref(self._get_artifact_cache_name(artifact))
+            return True
+        if kind == 'stratum' and \
+            self._has_file(self.artifact_filename(artifact)):
+            return True
+        return False
+
+    def get_artifact_metadata(self, artifact, name):
+        filename = self._artifact_metadata_filename(artifact, name)
+        os.utime(filename, None)
+        return open(filename)
+
+    def get_source_metadata_filename(self, source, cachekey, name):
+        return self._source_metadata_filename(source, cachekey, name)
+
+    def get_source_metadata(self, source, cachekey, name):
+        filename = self._source_metadata_filename(source, cachekey, name)
+        os.utime(filename, None)
+        return open(filename)
+
+    def artifact_filename(self, artifact):
+        return os.path.join(self.cachedir, artifact.basename())
+
+    def _artifact_metadata_filename(self, artifact, name):
+        return os.path.join(self.cachedir, artifact.metadata_basename(name))
+
+    def _source_metadata_filename(self, source, cachekey, name):
+        return os.path.join(self.cachedir, '%s.%s' % (cachekey, name))
+
+    def put_artifact_metadata(self, artifact, name):
+        filename = self._artifact_metadata_filename(artifact, name)
+        return morphlib.savefile.SaveFile(filename, mode='w')
+
+    def put_source_metadata(self, source, cachekey, name):
+        filename = self._source_metadata_filename(source, cachekey, name)
+        return morphlib.savefile.SaveFile(filename, mode='w')
+
+    def _has_file(self, filename):
+        if os.path.exists(filename):
+            os.utime(filename, None)
+            return True
+        return False
+
+    def has_artifact_metadata(self, artifact, name):
+        filename = self._artifact_metadata_filename(artifact, name)
+        return self._has_file(filename)
+
+    def has_source_metadata(self, source, cachekey, name):
+        filename = self._source_metadata_filename(source, cachekey, name)
+        return self._has_file(filename)
diff --git a/without-test-modules b/without-test-modules
index c3e7f5a2..ebbcfb6a 100644
--- a/without-test-modules
+++ b/without-test-modules
@@ -55,3 +55,4 @@ distbuild/worker_build_scheduler.py
 # Not unit tested, since it needs a full system branch
 morphlib/buildbranch.py
 morphlib/ostree.py
+morphlib/ostreeartifactcache.py