path: root/src/buildstream/_cas/cascache.py

#
#  Copyright (C) 2018 Codethink Limited
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public
#  License as published by the Free Software Foundation; either
#  version 2 of the License, or (at your option) any later version.
#
#  This library 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
#  Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with this library. If not, see <http://www.gnu.org/licenses/>.
#
#  Authors:
#        Jürg Billeter <juerg.billeter@codethink.co.uk>

import hashlib
import itertools
import os
import stat
import errno
import uuid
import contextlib

import grpc

from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
from .._protos.buildstream.v2 import buildstream_pb2

from .. import utils
from .._exceptions import CASCacheError, LoadError, LoadErrorReason
from .._message import Message, MessageType

from .casremote import BlobNotFound, _CASBatchRead, _CASBatchUpdate

_BUFFER_SIZE = 65536


CACHE_SIZE_FILE = "cache_size"


# CASCacheUsage
#
# A simple object to report the current CAS cache usage details.
#
# Note that this uses the user configured cache quota
# rather than the internal quota with protective headroom
# removed, to provide a more sensible value to display to
# the user.
#
# Args:
#    cas (CASQuota): The CAS cache to get the status of
#
class CASCacheUsage():

    def __init__(self, casquota):
        self.quota_config = casquota._config_cache_quota          # Configured quota
        self.quota_size = casquota._cache_quota_original          # Resolved cache quota in bytes
        self.used_size = casquota.get_cache_size()                # Size used by artifacts in bytes
        self.used_percent = 0                                # Percentage of the quota used
        if self.quota_size is not None:
            self.used_percent = int(self.used_size * 100 / self.quota_size)

    # Formattable into a human readable string
    #
    def __str__(self):
        return "{} / {} ({}%)" \
            .format(utils._pretty_size(self.used_size, dec_places=1),
                    self.quota_config,
                    self.used_percent)


# A CASCache manages a CAS repository as specified in the Remote Execution API.
#
# Args:
#     path (str): The root directory for the CAS repository
#     cache_quota (int): User configured cache quota
#
class CASCache():

    def __init__(self, path):
        self.casdir = os.path.join(path, 'cas')
        self.tmpdir = os.path.join(path, 'tmp')
        os.makedirs(os.path.join(self.casdir, 'refs', 'heads'), exist_ok=True)
        os.makedirs(os.path.join(self.casdir, 'objects'), exist_ok=True)
        os.makedirs(self.tmpdir, exist_ok=True)

        self.__reachable_directory_callbacks = []
        self.__reachable_digest_callbacks = []

    # preflight():
    #
    # Preflight check.
    #
    def preflight(self):
        headdir = os.path.join(self.casdir, 'refs', 'heads')
        objdir = os.path.join(self.casdir, 'objects')
        if not (os.path.isdir(headdir) and os.path.isdir(objdir)):
            raise CASCacheError("CAS repository check failed for '{}'".format(self.casdir))

    # contains():
    #
    # Check whether the specified ref is already available in the local CAS cache.
    #
    # Args:
    #     ref (str): The ref to check
    #
    # Returns: True if the ref is in the cache, False otherwise
    #
    def contains(self, ref):
        refpath = self._refpath(ref)

        # This assumes that the repository doesn't have any dangling pointers
        return os.path.exists(refpath)

    # contains_directory():
    #
    # Check whether the specified directory and subdirecotires are in the cache,
    # i.e non dangling.
    #
    # Args:
    #     digest (Digest): The directory digest to check
    #     with_files (bool): Whether to check files as well
    #
    # Returns: True if the directory is available in the local cache
    #
    def contains_directory(self, digest, *, with_files):
        try:
            directory = remote_execution_pb2.Directory()
            with open(self.objpath(digest), 'rb') as f:
                directory.ParseFromString(f.read())

            # Optionally check presence of files
            if with_files:
                for filenode in directory.files:
                    if not os.path.exists(self.objpath(filenode.digest)):
                        return False

            # Check subdirectories
            for dirnode in directory.directories:
                if not self.contains_directory(dirnode.digest, with_files=with_files):
                    return False

            return True
        except FileNotFoundError:
            return False

    # checkout():
    #
    # Checkout the specified directory digest.
    #
    # Args:
    #     dest (str): The destination path
    #     tree (Digest): The directory digest to extract
    #     can_link (bool): Whether we can create hard links in the destination
    #
    def checkout(self, dest, tree, *, can_link=False):
        os.makedirs(dest, exist_ok=True)

        directory = remote_execution_pb2.Directory()

        with open(self.objpath(tree), 'rb') as f:
            directory.ParseFromString(f.read())

        for filenode in directory.files:
            # regular file, create hardlink
            fullpath = os.path.join(dest, filenode.name)
            if can_link:
                utils.safe_link(self.objpath(filenode.digest), fullpath)
            else:
                utils.safe_copy(self.objpath(filenode.digest), fullpath)

            if filenode.is_executable:
                os.chmod(fullpath, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR |
                         stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH)

        for dirnode in directory.directories:
            fullpath = os.path.join(dest, dirnode.name)
            self.checkout(fullpath, dirnode.digest, can_link=can_link)

        for symlinknode in directory.symlinks:
            # symlink
            fullpath = os.path.join(dest, symlinknode.name)
            os.symlink(symlinknode.target, fullpath)

    # diff():
    #
    # Return a list of files that have been added or modified between
    # the refs described by ref_a and ref_b.
    #
    # Args:
    #     ref_a (str): The first ref
    #     ref_b (str): The second ref
    #     subdir (str): A subdirectory to limit the comparison to
    #
    def diff(self, ref_a, ref_b):
        tree_a = self.resolve_ref(ref_a)
        tree_b = self.resolve_ref(ref_b)

        added = []
        removed = []
        modified = []

        self.diff_trees(tree_a, tree_b, added=added, removed=removed, modified=modified)

        return modified, removed, added

    # pull():
    #
    # Pull a ref from a remote repository.
    #
    # Args:
    #     ref (str): The ref to pull
    #     remote (CASRemote): The remote repository to pull from
    #
    # Returns:
    #   (bool): True if pull was successful, False if ref was not available
    #
    def pull(self, ref, remote):
        try:
            remote.init()

            request = buildstream_pb2.GetReferenceRequest(instance_name=remote.spec.instance_name)
            request.key = ref
            response = remote.ref_storage.GetReference(request)

            tree = response.digest

            # Fetch Directory objects
            self._fetch_directory(remote, tree)

            # Fetch files, excluded_subdirs determined in pullqueue
            required_blobs = self.required_blobs_for_directory(tree)
            missing_blobs = self.local_missing_blobs(required_blobs)
            if missing_blobs:
                self.fetch_blobs(remote, missing_blobs)

            self.set_ref(ref, tree)

            return True
        except grpc.RpcError as e:
            if e.code() != grpc.StatusCode.NOT_FOUND:
                raise CASCacheError("Failed to pull ref {}: {}".format(ref, e)) from e
            else:
                return False
        except BlobNotFound:
            return False

    # pull_tree():
    #
    # Pull a single Tree rather than a ref.
    # Does not update local refs.
    #
    # Args:
    #     remote (CASRemote): The remote to pull from
    #     digest (Digest): The digest of the tree
    #
    def pull_tree(self, remote, digest):
        try:
            remote.init()

            digest = self._fetch_tree(remote, digest)

            return digest

        except grpc.RpcError as e:
            if e.code() != grpc.StatusCode.NOT_FOUND:
                raise

        return None

    # link_ref():
    #
    # Add an alias for an existing ref.
    #
    # Args:
    #     oldref (str): An existing ref
    #     newref (str): A new ref for the same directory
    #
    def link_ref(self, oldref, newref):
        tree = self.resolve_ref(oldref)

        self.set_ref(newref, tree)

    # push():
    #
    # Push committed refs to remote repository.
    #
    # Args:
    #     refs (list): The refs to push
    #     remote (CASRemote): The remote to push to
    #
    # Returns:
    #   (bool): True if any remote was updated, False if no pushes were required
    #
    # Raises:
    #   (CASCacheError): if there was an error
    #
    def push(self, refs, remote):
        skipped_remote = True
        try:
            for ref in refs:
                tree = self.resolve_ref(ref)

                # Check whether ref is already on the server in which case
                # there is no need to push the ref
                try:
                    request = buildstream_pb2.GetReferenceRequest(instance_name=remote.spec.instance_name)
                    request.key = ref
                    response = remote.ref_storage.GetReference(request)

                    if response.digest.hash == tree.hash and response.digest.size_bytes == tree.size_bytes:
                        # ref is already on the server with the same tree
                        continue

                except grpc.RpcError as e:
                    if e.code() != grpc.StatusCode.NOT_FOUND:
                        # Intentionally re-raise RpcError for outer except block.
                        raise

                self._send_directory(remote, tree)

                request = buildstream_pb2.UpdateReferenceRequest(instance_name=remote.spec.instance_name)
                request.keys.append(ref)
                request.digest.CopyFrom(tree)
                remote.ref_storage.UpdateReference(request)

                skipped_remote = False
        except grpc.RpcError as e:
            if e.code() != grpc.StatusCode.RESOURCE_EXHAUSTED:
                raise CASCacheError("Failed to push ref {}: {}".format(refs, e), temporary=True) from e

        return not skipped_remote

    # objpath():
    #
    # Return the path of an object based on its digest.
    #
    # Args:
    #     digest (Digest): The digest of the object
    #
    # Returns:
    #     (str): The path of the object
    #
    def objpath(self, digest):
        return os.path.join(self.casdir, 'objects', digest.hash[:2], digest.hash[2:])

    # add_object():
    #
    # Hash and write object to CAS.
    #
    # Args:
    #     digest (Digest): An optional Digest object to populate
    #     path (str): Path to file to add
    #     buffer (bytes): Byte buffer to add
    #     link_directly (bool): Whether file given by path can be linked
    #
    # Returns:
    #     (Digest): The digest of the added object
    #
    # Either `path` or `buffer` must be passed, but not both.
    #
    def add_object(self, *, digest=None, path=None, buffer=None, link_directly=False):
        # Exactly one of the two parameters has to be specified
        assert (path is None) != (buffer is None)

        if digest is None:
            digest = remote_execution_pb2.Digest()

        try:
            h = hashlib.sha256()
            # Always write out new file to avoid corruption if input file is modified
            with contextlib.ExitStack() as stack:
                if path is not None and link_directly:
                    tmp = stack.enter_context(open(path, 'rb'))
                    for chunk in iter(lambda: tmp.read(_BUFFER_SIZE), b""):
                        h.update(chunk)
                else:
                    tmp = stack.enter_context(self._temporary_object())

                    if path:
                        with open(path, 'rb') as f:
                            for chunk in iter(lambda: f.read(_BUFFER_SIZE), b""):
                                h.update(chunk)
                                tmp.write(chunk)
                    else:
                        h.update(buffer)
                        tmp.write(buffer)

                    tmp.flush()

                digest.hash = h.hexdigest()
                digest.size_bytes = os.fstat(tmp.fileno()).st_size

                # Place file at final location
                objpath = self.objpath(digest)
                os.makedirs(os.path.dirname(objpath), exist_ok=True)
                os.link(tmp.name, objpath)

        except FileExistsError:
            # We can ignore the failed link() if the object is already in the repo.
            pass

        except OSError as e:
            raise CASCacheError("Failed to hash object: {}".format(e)) from e

        return digest

    # set_ref():
    #
    # Create or replace a ref.
    #
    # Args:
    #     ref (str): The name of the ref
    #
    def set_ref(self, ref, tree):
        refpath = self._refpath(ref)
        os.makedirs(os.path.dirname(refpath), exist_ok=True)
        with utils.save_file_atomic(refpath, 'wb', tempdir=self.tmpdir) as f:
            f.write(tree.SerializeToString())

    # resolve_ref():
    #
    # Resolve a ref to a digest.
    #
    # Args:
    #     ref (str): The name of the ref
    #     update_mtime (bool): Whether to update the mtime of the ref
    #
    # Returns:
    #     (Digest): The digest stored in the ref
    #
    def resolve_ref(self, ref, *, update_mtime=False):
        refpath = self._refpath(ref)

        try:
            with open(refpath, 'rb') as f:
                if update_mtime:
                    os.utime(refpath)

                digest = remote_execution_pb2.Digest()
                digest.ParseFromString(f.read())
                return digest

        except FileNotFoundError as e:
            raise CASCacheError("Attempt to access unavailable ref: {}".format(e)) from e

    # update_mtime()
    #
    # Update the mtime of a ref.
    #
    # Args:
    #     ref (str): The ref to update
    #
    def update_mtime(self, ref):
        try:
            os.utime(self._refpath(ref))
        except FileNotFoundError as e:
            raise CASCacheError("Attempt to access unavailable ref: {}".format(e)) from e

    # list_objects():
    #
    # List cached objects in Least Recently Modified (LRM) order.
    #
    # Returns:
    #     (list) - A list of objects and timestamps in LRM order
    #
    def list_objects(self):
        objs = []
        mtimes = []

        for root, _, files in os.walk(os.path.join(self.casdir, 'objects')):
            for filename in files:
                obj_path = os.path.join(root, filename)
                try:
                    mtimes.append(os.path.getmtime(obj_path))
                except FileNotFoundError:
                    pass
                else:
                    objs.append(obj_path)

        # NOTE: Sorted will sort from earliest to latest, thus the
        # first element of this list will be the file modified earliest.
        return sorted(zip(mtimes, objs))

    def clean_up_refs_until(self, time):
        ref_heads = os.path.join(self.casdir, 'refs', 'heads')

        for root, _, files in os.walk(ref_heads):
            for filename in files:
                ref_path = os.path.join(root, filename)
                # Obtain the mtime (the time a file was last modified)
                if os.path.getmtime(ref_path) < time:
                    os.unlink(ref_path)

    # remove():
    #
    # Removes the given symbolic ref from the repo.
    #
    # Args:
    #    ref (str): A symbolic ref
    #    basedir (str): Path of base directory the ref is in, defaults to
    #                   CAS refs heads
    #    defer_prune (bool): Whether to defer pruning to the caller. NOTE:
    #                        The space won't be freed until you manually
    #                        call prune.
    #
    # Returns:
    #    (int|None) The amount of space pruned from the repository in
    #               Bytes, or None if defer_prune is True
    #
    def remove(self, ref, *, basedir=None, defer_prune=False):

        if basedir is None:
            basedir = os.path.join(self.casdir, 'refs', 'heads')
        # Remove cache ref
        self._remove_ref(ref, basedir)

        if not defer_prune:
            pruned = self.prune()
            return pruned

        return None

    # adds callback of iterator over reachable directory digests
    def add_reachable_directories_callback(self, callback):
        self.__reachable_directory_callbacks.append(callback)

    # adds callbacks of iterator over reachable file digests
    def add_reachable_digests_callback(self, callback):
        self.__reachable_digest_callbacks.append(callback)

    # prune():
    #
    # Prune unreachable objects from the repo.
    #
    def prune(self):
        ref_heads = os.path.join(self.casdir, 'refs', 'heads')

        pruned = 0
        reachable = set()

        # Check which objects are reachable
        for root, _, files in os.walk(ref_heads):
            for filename in files:
                ref_path = os.path.join(root, filename)
                ref = os.path.relpath(ref_path, ref_heads)

                tree = self.resolve_ref(ref)
                self._reachable_refs_dir(reachable, tree)

        # check callback directory digests that are reachable
        for digest_callback in self.__reachable_directory_callbacks:
            for digest in digest_callback():
                self._reachable_refs_dir(reachable, digest)

        # check callback file digests that are reachable
        for digest_callback in self.__reachable_digest_callbacks:
            for digest in digest_callback():
                reachable.add(digest.hash)

        # Prune unreachable objects
        for root, _, files in os.walk(os.path.join(self.casdir, 'objects')):
            for filename in files:
                objhash = os.path.basename(root) + filename
                if objhash not in reachable:
                    obj_path = os.path.join(root, filename)
                    pruned += os.stat(obj_path).st_size
                    os.unlink(obj_path)

        return pruned

    def update_tree_mtime(self, tree):
        reachable = set()
        self._reachable_refs_dir(reachable, tree, update_mtime=True)

    # remote_missing_blobs_for_directory():
    #
    # Determine which blobs of a directory tree are missing on the remote.
    #
    # Args:
    #     digest (Digest): The directory digest
    #
    # Returns: List of missing Digest objects
    #
    def remote_missing_blobs_for_directory(self, remote, digest):
        required_blobs = self.required_blobs_for_directory(digest)

        return self.remote_missing_blobs(remote, required_blobs)

    # remote_missing_blobs():
    #
    # Determine which blobs are missing on the remote.
    #
    # Args:
    #     blobs ([Digest]): List of directory digests to check
    #
    # Returns: List of missing Digest objects
    #
    def remote_missing_blobs(self, remote, blobs):
        missing_blobs = dict()
        # Limit size of FindMissingBlobs request
        for required_blobs_group in _grouper(iter(blobs), 512):
            request = remote_execution_pb2.FindMissingBlobsRequest(instance_name=remote.spec.instance_name)

            for required_digest in required_blobs_group:
                d = request.blob_digests.add()
                d.CopyFrom(required_digest)

            response = remote.cas.FindMissingBlobs(request)
            for missing_digest in response.missing_blob_digests:
                d = remote_execution_pb2.Digest()
                d.CopyFrom(missing_digest)
                missing_blobs[d.hash] = d

        return missing_blobs.values()

    # local_missing_blobs():
    #
    # Check local cache for missing blobs.
    #
    # Args:
    #    digests (list): The Digests of blobs to check
    #
    # Returns: Missing Digest objects
    #
    def local_missing_blobs(self, digests):
        missing_blobs = []
        for digest in digests:
            objpath = self.objpath(digest)
            if not os.path.exists(objpath):
                missing_blobs.append(digest)
        return missing_blobs

    # required_blobs_for_directory():
    #
    # Generator that returns the Digests of all blobs in the tree specified by
    # the Digest of the toplevel Directory object.
    #
    def required_blobs_for_directory(self, directory_digest, *, excluded_subdirs=None):
        if not excluded_subdirs:
            excluded_subdirs = []

        # parse directory, and recursively add blobs

        yield directory_digest

        directory = remote_execution_pb2.Directory()

        with open(self.objpath(directory_digest), 'rb') as f:
            directory.ParseFromString(f.read())

        for filenode in directory.files:
            yield filenode.digest

        for dirnode in directory.directories:
            if dirnode.name not in excluded_subdirs:
                yield from self.required_blobs_for_directory(dirnode.digest)

    def diff_trees(self, tree_a, tree_b, *, added, removed, modified, path=""):
        dir_a = remote_execution_pb2.Directory()
        dir_b = remote_execution_pb2.Directory()

        if tree_a:
            with open(self.objpath(tree_a), 'rb') as f:
                dir_a.ParseFromString(f.read())
        if tree_b:
            with open(self.objpath(tree_b), 'rb') as f:
                dir_b.ParseFromString(f.read())

        a = 0
        b = 0
        while a < len(dir_a.files) or b < len(dir_b.files):
            if b < len(dir_b.files) and (a >= len(dir_a.files) or
                                         dir_a.files[a].name > dir_b.files[b].name):
                added.append(os.path.join(path, dir_b.files[b].name))
                b += 1
            elif a < len(dir_a.files) and (b >= len(dir_b.files) or
                                           dir_b.files[b].name > dir_a.files[a].name):
                removed.append(os.path.join(path, dir_a.files[a].name))
                a += 1
            else:
                # File exists in both directories
                if dir_a.files[a].digest.hash != dir_b.files[b].digest.hash:
                    modified.append(os.path.join(path, dir_a.files[a].name))
                a += 1
                b += 1

        a = 0
        b = 0
        while a < len(dir_a.directories) or b < len(dir_b.directories):
            if b < len(dir_b.directories) and (a >= len(dir_a.directories) or
                                               dir_a.directories[a].name > dir_b.directories[b].name):
                self.diff_trees(None, dir_b.directories[b].digest,
                                added=added, removed=removed, modified=modified,
                                path=os.path.join(path, dir_b.directories[b].name))
                b += 1
            elif a < len(dir_a.directories) and (b >= len(dir_b.directories) or
                                                 dir_b.directories[b].name > dir_a.directories[a].name):
                self.diff_trees(dir_a.directories[a].digest, None,
                                added=added, removed=removed, modified=modified,
                                path=os.path.join(path, dir_a.directories[a].name))
                a += 1
            else:
                # Subdirectory exists in both directories
                if dir_a.directories[a].digest.hash != dir_b.directories[b].digest.hash:
                    self.diff_trees(dir_a.directories[a].digest, dir_b.directories[b].digest,
                                    added=added, removed=removed, modified=modified,
                                    path=os.path.join(path, dir_a.directories[a].name))
                a += 1
                b += 1

    ################################################
    #             Local Private Methods            #
    ################################################

    def _refpath(self, ref):
        return os.path.join(self.casdir, 'refs', 'heads', ref)

    # _remove_ref()
    #
    # Removes a ref.
    #
    # This also takes care of pruning away directories which can
    # be removed after having removed the given ref.
    #
    # Args:
    #    ref (str): The ref to remove
    #    basedir (str): Path of base directory the ref is in
    #
    # Raises:
    #    (CASCacheError): If the ref didnt exist, or a system error
    #                     occurred while removing it
    #
    def _remove_ref(self, ref, basedir):

        # Remove the ref itself
        refpath = os.path.join(basedir, ref)

        try:
            os.unlink(refpath)
        except FileNotFoundError as e:
            raise CASCacheError("Could not find ref '{}'".format(ref)) from e

        # Now remove any leading directories

        components = list(os.path.split(ref))
        while components:
            components.pop()
            refdir = os.path.join(basedir, *components)

            # Break out once we reach the base
            if refdir == basedir:
                break

            try:
                os.rmdir(refdir)
            except FileNotFoundError:
                # The parent directory did not exist, but it's
                # parent directory might still be ready to prune
                pass
            except OSError as e:
                if e.errno == errno.ENOTEMPTY:
                    # The parent directory was not empty, so we
                    # cannot prune directories beyond this point
                    break

                # Something went wrong here
                raise CASCacheError("System error while removing ref '{}': {}".format(ref, e)) from e

    def _get_subdir(self, tree, subdir):
        head, name = os.path.split(subdir)
        if head:
            tree = self._get_subdir(tree, head)

        directory = remote_execution_pb2.Directory()

        with open(self.objpath(tree), 'rb') as f:
            directory.ParseFromString(f.read())

        for dirnode in directory.directories:
            if dirnode.name == name:
                return dirnode.digest

        raise CASCacheError("Subdirectory {} not found".format(name))

    def _reachable_refs_dir(self, reachable, tree, update_mtime=False, check_exists=False):
        if tree.hash in reachable:
            return
        try:
            if update_mtime:
                os.utime(self.objpath(tree))

            reachable.add(tree.hash)

            directory = remote_execution_pb2.Directory()

            with open(self.objpath(tree), 'rb') as f:
                directory.ParseFromString(f.read())

        except FileNotFoundError:
            # Just exit early if the file doesn't exist
            return

        for filenode in directory.files:
            if update_mtime:
                os.utime(self.objpath(filenode.digest))
            if check_exists:
                if not os.path.exists(self.objpath(filenode.digest)):
                    raise FileNotFoundError
            reachable.add(filenode.digest.hash)

        for dirnode in directory.directories:
            self._reachable_refs_dir(reachable, dirnode.digest, update_mtime=update_mtime, check_exists=check_exists)

    # _temporary_object():
    #
    # Returns:
    #     (file): A file object to a named temporary file.
    #
    # Create a named temporary file with 0o0644 access rights.
    @contextlib.contextmanager
    def _temporary_object(self):
        with utils._tempnamedfile(dir=self.tmpdir) as f:
            os.chmod(f.name,
                     stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)
            yield f

    # _ensure_blob():
    #
    # Fetch and add blob if it's not already local.
    #
    # Args:
    #     remote (Remote): The remote to use.
    #     digest (Digest): Digest object for the blob to fetch.
    #
    # Returns:
    #     (str): The path of the object
    #
    def _ensure_blob(self, remote, digest):
        objpath = self.objpath(digest)
        if os.path.exists(objpath):
            # already in local repository
            return objpath

        with self._temporary_object() as f:
            remote._fetch_blob(digest, f)

            added_digest = self.add_object(path=f.name, link_directly=True)
            assert added_digest.hash == digest.hash

        return objpath

    def _batch_download_complete(self, batch, *, missing_blobs=None):
        for digest, data in batch.send(missing_blobs=missing_blobs):
            with self._temporary_object() as f:
                f.write(data)
                f.flush()

                added_digest = self.add_object(path=f.name, link_directly=True)
                assert added_digest.hash == digest.hash

    # Helper function for _fetch_directory().
    def _fetch_directory_batch(self, remote, batch, fetch_queue, fetch_next_queue):
        self._batch_download_complete(batch)

        # All previously scheduled directories are now locally available,
        # move them to the processing queue.
        fetch_queue.extend(fetch_next_queue)
        fetch_next_queue.clear()
        return _CASBatchRead(remote)

    # Helper function for _fetch_directory().
    def _fetch_directory_node(self, remote, digest, batch, fetch_queue, fetch_next_queue, *, recursive=False):
        in_local_cache = os.path.exists(self.objpath(digest))

        if in_local_cache:
            # Skip download, already in local cache.
            pass
        elif (digest.size_bytes >= remote.max_batch_total_size_bytes or
              not remote.batch_read_supported):
            # Too large for batch request, download in independent request.
            self._ensure_blob(remote, digest)
            in_local_cache = True
        else:
            if not batch.add(digest):
                # Not enough space left in batch request.
                # Complete pending batch first.
                batch = self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)
                batch.add(digest)

        if recursive:
            if in_local_cache:
                # Add directory to processing queue.
                fetch_queue.append(digest)
            else:
                # Directory will be available after completing pending batch.
                # Add directory to deferred processing queue.
                fetch_next_queue.append(digest)

        return batch

    # _fetch_directory():
    #
    # Fetches remote directory and adds it to content addressable store.
    #
    # This recursively fetches directory objects but doesn't fetch any
    # files.
    #
    # Args:
    #     remote (Remote): The remote to use.
    #     dir_digest (Digest): Digest object for the directory to fetch.
    #
    def _fetch_directory(self, remote, dir_digest):
        # TODO Use GetTree() if the server supports it

        fetch_queue = [dir_digest]
        fetch_next_queue = []
        batch = _CASBatchRead(remote)

        while len(fetch_queue) + len(fetch_next_queue) > 0:
            if not fetch_queue:
                batch = self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)

            dir_digest = fetch_queue.pop(0)

            objpath = self._ensure_blob(remote, dir_digest)

            directory = remote_execution_pb2.Directory()
            with open(objpath, 'rb') as f:
                directory.ParseFromString(f.read())

            for dirnode in directory.directories:
                batch = self._fetch_directory_node(remote, dirnode.digest, batch,
                                                   fetch_queue, fetch_next_queue, recursive=True)

        # Fetch final batch
        self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)

    def _fetch_tree(self, remote, digest):
        # download but do not store the Tree object
        with utils._tempnamedfile(dir=self.tmpdir) as out:
            remote._fetch_blob(digest, out)

            tree = remote_execution_pb2.Tree()

            with open(out.name, 'rb') as f:
                tree.ParseFromString(f.read())

            tree.children.extend([tree.root])
            for directory in tree.children:
                dirbuffer = directory.SerializeToString()
                dirdigest = self.add_object(buffer=dirbuffer)
                assert dirdigest.size_bytes == len(dirbuffer)

        return dirdigest

    # fetch_blobs():
    #
    # Fetch blobs from remote CAS. Returns missing blobs that could not be fetched.
    #
    # Args:
    #    remote (CASRemote): The remote repository to fetch from
    #    digests (list): The Digests of blobs to fetch
    #
    # Returns: The Digests of the blobs that were not available on the remote CAS
    #
    def fetch_blobs(self, remote, digests):
        missing_blobs = []

        batch = _CASBatchRead(remote)

        for digest in digests:
            if (digest.size_bytes >= remote.max_batch_total_size_bytes or
                    not remote.batch_read_supported):
                # Too large for batch request, download in independent request.
                try:
                    self._ensure_blob(remote, digest)
                except grpc.RpcError as e:
                    if e.code() == grpc.StatusCode.NOT_FOUND:
                        missing_blobs.append(digest)
                    else:
                        raise CASCacheError("Failed to fetch blob: {}".format(e)) from e
            else:
                if not batch.add(digest):
                    # Not enough space left in batch request.
                    # Complete pending batch first.
                    self._batch_download_complete(batch, missing_blobs=missing_blobs)

                    batch = _CASBatchRead(remote)
                    batch.add(digest)

        # Complete last pending batch
        self._batch_download_complete(batch, missing_blobs=missing_blobs)

        return missing_blobs

    # send_blobs():
    #
    # Upload blobs to remote CAS.
    #
    # Args:
    #    remote (CASRemote): The remote repository to upload to
    #    digests (list): The Digests of Blobs to upload
    #
    def send_blobs(self, remote, digests, u_uid=uuid.uuid4()):
        batch = _CASBatchUpdate(remote)

        for digest in digests:
            with open(self.objpath(digest), 'rb') as f:
                assert os.fstat(f.fileno()).st_size == digest.size_bytes

                if (digest.size_bytes >= remote.max_batch_total_size_bytes or
                        not remote.batch_update_supported):
                    # Too large for batch request, upload in independent request.
                    remote._send_blob(digest, f, u_uid=u_uid)
                else:
                    if not batch.add(digest, f):
                        # Not enough space left in batch request.
                        # Complete pending batch first.
                        batch.send()
                        batch = _CASBatchUpdate(remote)
                        batch.add(digest, f)

        # Send final batch
        batch.send()

    def _send_directory(self, remote, digest, u_uid=uuid.uuid4()):
        missing_blobs = self.remote_missing_blobs_for_directory(remote, digest)

        # Upload any blobs missing on the server
        self.send_blobs(remote, missing_blobs, u_uid)


class CASQuota:
    def __init__(self, context):
        self.context = context
        self.cas = context.get_cascache()
        self.casdir = self.cas.casdir
        self._config_cache_quota = context.config_cache_quota
        self._config_cache_quota_string = context.config_cache_quota_string
        self._cache_size = None               # The current cache size, sometimes it's an estimate
        self._cache_quota = None              # The cache quota
        self._cache_quota_original = None     # The cache quota as specified by the user, in bytes
        self._cache_quota_headroom = None     # The headroom in bytes before reaching the quota or full disk
        self._cache_lower_threshold = None    # The target cache size for a cleanup
        self.available_space = None

        self._message = context.message

        self._remove_callbacks = []   # Callbacks to remove unrequired refs and their remove method
        self._list_refs_callbacks = []  # Callbacks to all refs

        self._calculate_cache_quota()

    # compute_cache_size()
    #
    # Computes the real artifact cache size.
    #
    # Returns:
    #    (int): The size of the artifact cache.
    #
    def compute_cache_size(self):
        self._cache_size = utils._get_dir_size(self.casdir)
        return self._cache_size

    # get_cache_size()
    #
    # Fetches the cached size of the cache, this is sometimes
    # an estimate and periodically adjusted to the real size
    # when a cache size calculation job runs.
    #
    # When it is an estimate, the value is either correct, or
    # it is greater than the actual cache size.
    #
    # Returns:
    #     (int) An approximation of the artifact cache size, in bytes.
    #
    def get_cache_size(self):

        # If we don't currently have an estimate, figure out the real cache size.
        if self._cache_size is None:
            stored_size = self._read_cache_size()
            if stored_size is not None:
                self._cache_size = stored_size
            else:
                self.compute_cache_size()

        return self._cache_size

    # set_cache_size()
    #
    # Forcefully set the overall cache size.
    #
    # This is used to update the size in the main process after
    # having calculated in a cleanup or a cache size calculation job.
    #
    # Args:
    #     cache_size (int): The size to set.
    #     write_to_disk (bool): Whether to write the value to disk.
    #
    def set_cache_size(self, cache_size, *, write_to_disk=True):

        assert cache_size is not None

        self._cache_size = cache_size
        if write_to_disk:
            self._write_cache_size(self._cache_size)

    # full()
    #
    # Checks if the artifact cache is full, either
    # because the user configured quota has been exceeded
    # or because the underlying disk is almost full.
    #
    # Returns:
    #    (bool): True if the artifact cache is full
    #
    def full(self):

        if self.get_cache_size() > self._cache_quota:
            return True

        _, volume_avail = self._get_cache_volume_size()
        if volume_avail < self._cache_quota_headroom:
            return True

        return False

    # add_remove_callbacks()
    #
    # This adds tuples of iterators over unrequired objects (currently
    # artifacts and source refs), and a callback to remove them.
    #
    # Args:
    #    callback (iter(unrequired), remove): tuple of iterator and remove
    #        method associated.
    #
    def add_remove_callbacks(self, list_unrequired, remove_method):
        self._remove_callbacks.append((list_unrequired, remove_method))

    def add_list_refs_callback(self, list_callback):
        self._list_refs_callbacks.append(list_callback)

    ################################################
    #             Local Private Methods            #
    ################################################

    # _read_cache_size()
    #
    # Reads and returns the size of the artifact cache that's stored in the
    # cache's size file
    #
    # Returns:
    #    (int): The size of the artifact cache, as recorded in the file
    #
    def _read_cache_size(self):
        size_file_path = os.path.join(self.casdir, CACHE_SIZE_FILE)

        if not os.path.exists(size_file_path):
            return None

        with open(size_file_path, "r") as f:
            size = f.read()

        try:
            num_size = int(size)
        except ValueError as e:
            raise CASCacheError("Size '{}' parsed from '{}' was not an integer".format(
                size, size_file_path)) from e

        return num_size

    # _write_cache_size()
    #
    # Writes the given size of the artifact to the cache's size file
    #
    # Args:
    #    size (int): The size of the artifact cache to record
    #
    def _write_cache_size(self, size):
        assert isinstance(size, int)
        size_file_path = os.path.join(self.casdir, CACHE_SIZE_FILE)
        with utils.save_file_atomic(size_file_path, "w", tempdir=self.cas.tmpdir) as f:
            f.write(str(size))

    # _get_cache_volume_size()
    #
    # Get the available space and total space for the volume on
    # which the artifact cache is located.
    #
    # Returns:
    #    (int): The total number of bytes on the volume
    #    (int): The number of available bytes on the volume
    #
    # NOTE: We use this stub to allow the test cases
    #       to override what an artifact cache thinks
    #       about it's disk size and available bytes.
    #
    def _get_cache_volume_size(self):
        return utils._get_volume_size(self.casdir)

    # _calculate_cache_quota()
    #
    # Calculates and sets the cache quota and lower threshold based on the
    # quota set in Context.
    # It checks that the quota is both a valid expression, and that there is
    # enough disk space to satisfy that quota
    #
    def _calculate_cache_quota(self):
        # Headroom intended to give BuildStream a bit of leeway.
        # This acts as the minimum size of cache_quota and also
        # is taken from the user requested cache_quota.
        #
        if 'BST_TEST_SUITE' in os.environ:
            self._cache_quota_headroom = 0
        else:
            self._cache_quota_headroom = 2e9

        total_size, available_space = self._get_cache_volume_size()
        cache_size = self.get_cache_size()
        self.available_space = available_space

        # Ensure system has enough storage for the cache_quota
        #
        # If cache_quota is none, set it to the maximum it could possibly be.
        #
        # Also check that cache_quota is at least as large as our headroom.
        #
        cache_quota = self._config_cache_quota
        if cache_quota is None:
            # The user has set no limit, so we may take all the space.
            cache_quota = min(cache_size + available_space, total_size)
        if cache_quota < self._cache_quota_headroom:  # Check minimum
            raise LoadError(
                LoadErrorReason.INVALID_DATA,
                "Invalid cache quota ({}): BuildStream requires a minimum cache quota of {}.".format(
                    utils._pretty_size(cache_quota),
                    utils._pretty_size(self._cache_quota_headroom)))
        elif cache_quota > total_size:
            # A quota greater than the total disk size is certianly an error
            raise CASCacheError("Your system does not have enough available " +
                                "space to support the cache quota specified.",
                                detail=("You have specified a quota of {quota} total disk space.\n" +
                                        "The filesystem containing {local_cache_path} only " +
                                        "has {total_size} total disk space.")
                                .format(
                                    quota=self._config_cache_quota,
                                    local_cache_path=self.casdir,
                                    total_size=utils._pretty_size(total_size)),
                                reason='insufficient-storage-for-quota')

        elif cache_quota > cache_size + available_space:
            # The quota does not fit in the available space, this is a warning
            if '%' in self._config_cache_quota_string:
                available = (available_space / total_size) * 100
                available = '{}% of total disk space'.format(round(available, 1))
            else:
                available = utils._pretty_size(available_space)

            self._message(Message(
                None,
                MessageType.WARN,
                "Your system does not have enough available " +
                "space to support the cache quota specified.",
                detail=("You have specified a quota of {quota} total disk space.\n" +
                        "The filesystem containing {local_cache_path} only " +
                        "has {available_size} available.")
                .format(quota=self._config_cache_quota,
                        local_cache_path=self.casdir,
                        available_size=available)))

        # Place a slight headroom (2e9 (2GB) on the cache_quota) into
        # cache_quota to try and avoid exceptions.
        #
        # Of course, we might still end up running out during a build
        # if we end up writing more than 2G, but hey, this stuff is
        # already really fuzzy.
        #
        self._cache_quota_original = cache_quota
        self._cache_quota = cache_quota - self._cache_quota_headroom
        self._cache_lower_threshold = self._cache_quota / 2

    # clean():
    #
    # Clean the artifact cache as much as possible.
    #
    # Args:
    #    progress (callable): A callback to call when a ref is removed
    #
    # Returns:
    #    (int): The size of the cache after having cleaned up
    #
    def clean(self, progress=None):
        context = self.context

        # Some accumulative statistics
        removed_ref_count = 0
        space_saved = 0

        total_refs = 0
        for refs in self._list_refs_callbacks:
            total_refs += len(list(refs()))

        # Start off with an announcement with as much info as possible
        volume_size, volume_avail = self._get_cache_volume_size()
        self._message(Message(
            None, MessageType.STATUS, "Starting cache cleanup",
            detail=("Elements required by the current build plan:\n" + "{}\n" +
                    "User specified quota: {} ({})\n" +
                    "Cache usage: {}\n" +
                    "Cache volume: {} total, {} available")
            .format(
                total_refs,
                context.config_cache_quota,
                utils._pretty_size(self._cache_quota, dec_places=2),
                utils._pretty_size(self.get_cache_size(), dec_places=2),
                utils._pretty_size(volume_size, dec_places=2),
                utils._pretty_size(volume_avail, dec_places=2))))

        # Do a real computation of the cache size once, just in case
        self.compute_cache_size()
        usage = CASCacheUsage(self)
        self._message(Message(None, MessageType.STATUS,
                              "Cache usage recomputed: {}".format(usage)))

        # Collect digests and their remove method
        all_unrequired_refs = []
        for (unrequired_refs, remove) in self._remove_callbacks:
            for (mtime, ref) in unrequired_refs():
                all_unrequired_refs.append((mtime, ref, remove))

        # Pair refs and their remove method sorted in time order
        all_unrequired_refs = [(ref, remove) for (_, ref, remove) in sorted(all_unrequired_refs)]

        # Go through unrequired refs and remove them, oldest first
        made_space = False
        for (ref, remove) in all_unrequired_refs:
            size = remove(ref)
            removed_ref_count += 1
            space_saved += size

            self._message(Message(
                None, MessageType.STATUS,
                "Freed {: <7} {}".format(
                    utils._pretty_size(size, dec_places=2),
                    ref)))

            self.set_cache_size(self._cache_size - size)

            # User callback
            #
            # Currently this process is fairly slow, but we should
            # think about throttling this progress() callback if this
            # becomes too intense.
            if progress:
                progress()

            if self.get_cache_size() < self._cache_lower_threshold:
                made_space = True
                break

        if not made_space and self.full():
            # If too many artifacts are required, and we therefore
            # can't remove them, we have to abort the build.
            #
            # FIXME: Asking the user what to do may be neater
            #
            default_conf = os.path.join(os.environ['XDG_CONFIG_HOME'],
                                        'buildstream.conf')
            detail = ("Aborted after removing {} refs and saving {} disk space.\n"
                      "The remaining {} in the cache is required by the {} references in your build plan\n\n"
                      "There is not enough space to complete the build.\n"
                      "Please increase the cache-quota in {} and/or make more disk space."
                      .format(removed_ref_count,
                              utils._pretty_size(space_saved, dec_places=2),
                              utils._pretty_size(self.get_cache_size(), dec_places=2),
                              total_refs,
                              (context.config_origin or default_conf)))

            raise CASCacheError("Cache too full. Aborting.",
                                detail=detail,
                                reason="cache-too-full")

        # Informational message about the side effects of the cleanup
        self._message(Message(
            None, MessageType.INFO, "Cleanup completed",
            detail=("Removed {} refs and saving {} disk space.\n" +
                    "Cache usage is now: {}")
            .format(removed_ref_count,
                    utils._pretty_size(space_saved, dec_places=2),
                    utils._pretty_size(self.get_cache_size(), dec_places=2))))

        return self.get_cache_size()


def _grouper(iterable, n):
    while True:
        try:
            current = next(iterable)
        except StopIteration:
            return
        yield itertools.chain([current], itertools.islice(iterable, n - 1))