| # |
| # 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> |
| |
| from concurrent import futures |
| from enum import Enum |
| import contextlib |
| import logging |
| import os |
| import signal |
| import sys |
| |
| import grpc |
| from google.protobuf.message import DecodeError |
| import click |
| |
| from .._protos.build.bazel.remote.execution.v2 import ( |
| remote_execution_pb2, |
| remote_execution_pb2_grpc, |
| ) |
| from .._protos.google.bytestream import bytestream_pb2_grpc |
| from .._protos.build.buildgrid import local_cas_pb2 |
| from .._protos.buildstream.v2 import ( |
| buildstream_pb2, |
| buildstream_pb2_grpc, |
| artifact_pb2, |
| artifact_pb2_grpc, |
| source_pb2, |
| source_pb2_grpc, |
| ) |
| |
| # Note: We'd ideally like to avoid imports from the core codebase as |
| # much as possible, since we're expecting to eventually split this |
| # module off into its own project. |
| # |
| # Not enough that we'd like to duplicate code, but enough that we want |
| # to make it very obvious what we're using, so in this case we import |
| # the specific methods we'll be using. |
| from ..utils import save_file_atomic, _remove_path_with_parents |
| from .casdprocessmanager import CASDProcessManager |
| |
| |
| # The default limit for gRPC messages is 4 MiB. |
| # Limit payload to 1 MiB to leave sufficient headroom for metadata. |
| _MAX_PAYLOAD_BYTES = 1024 * 1024 |
| |
| |
| # LogLevel(): |
| # |
| # Manage log level choices using click. |
| # |
| class LogLevel(click.Choice): |
| # Levels(): |
| # |
| # Represents the actual buildbox-casd log level. |
| # |
| class Levels(Enum): |
| WARNING = "warning" |
| INFO = "info" |
| TRACE = "trace" |
| |
| def __init__(self): |
| super().__init__([m.lower() for m in LogLevel.Levels._member_names_]) # pylint: disable=no-member |
| |
| def convert(self, value, param, ctx) -> "LogLevel.Levels": |
| return LogLevel.Levels(super().convert(value, param, ctx)) |
| |
| @classmethod |
| def get_logging_equivalent(cls, level) -> int: |
| equivalents = { |
| cls.Levels.WARNING: logging.WARNING, |
| cls.Levels.INFO: logging.INFO, |
| cls.Levels.TRACE: logging.DEBUG, |
| } |
| |
| return equivalents[level] |
| |
| |
| # create_server(): |
| # |
| # Create gRPC CAS artifact server as specified in the Remote Execution API. |
| # |
| # Args: |
| # repo (str): Path to CAS repository |
| # enable_push (bool): Whether to allow blob uploads and artifact updates |
| # index_only (bool): Whether to store CAS blobs or only artifacts |
| # |
| @contextlib.contextmanager |
| def create_server(repo, *, enable_push, quota, index_only, log_level=LogLevel.Levels.WARNING): |
| logger = logging.getLogger("buildstream._cas.casserver") |
| logger.setLevel(LogLevel.get_logging_equivalent(log_level)) |
| handler = logging.StreamHandler(sys.stderr) |
| handler.setFormatter(logging.Formatter(fmt="%(levelname)s: %(funcName)s: %(message)s")) |
| logger.addHandler(handler) |
| |
| casd_manager = CASDProcessManager( |
| os.path.abspath(repo), os.path.join(os.path.abspath(repo), "logs"), log_level, quota, False |
| ) |
| casd_channel = casd_manager.create_channel() |
| |
| try: |
| root = os.path.abspath(repo) |
| sourcedir = os.path.join(root, "source_protos") |
| |
| # Use max_workers default from Python 3.5+ |
| max_workers = (os.cpu_count() or 1) * 5 |
| server = grpc.server(futures.ThreadPoolExecutor(max_workers)) |
| |
| if not index_only: |
| bytestream_pb2_grpc.add_ByteStreamServicer_to_server( |
| _ByteStreamServicer(casd_channel, enable_push=enable_push), server |
| ) |
| |
| remote_execution_pb2_grpc.add_ContentAddressableStorageServicer_to_server( |
| _ContentAddressableStorageServicer(casd_channel, enable_push=enable_push), server |
| ) |
| |
| remote_execution_pb2_grpc.add_CapabilitiesServicer_to_server(_CapabilitiesServicer(), server) |
| |
| buildstream_pb2_grpc.add_ReferenceStorageServicer_to_server( |
| _ReferenceStorageServicer(casd_channel, root, enable_push=enable_push), server |
| ) |
| |
| artifact_pb2_grpc.add_ArtifactServiceServicer_to_server( |
| _ArtifactServicer(casd_channel, root, update_cas=not index_only), server |
| ) |
| |
| source_pb2_grpc.add_SourceServiceServicer_to_server(_SourceServicer(sourcedir), server) |
| |
| # Create up reference storage and artifact capabilities |
| artifact_capabilities = buildstream_pb2.ArtifactCapabilities(allow_updates=enable_push) |
| source_capabilities = buildstream_pb2.SourceCapabilities(allow_updates=enable_push) |
| buildstream_pb2_grpc.add_CapabilitiesServicer_to_server( |
| _BuildStreamCapabilitiesServicer(artifact_capabilities, source_capabilities), server |
| ) |
| |
| yield server |
| |
| finally: |
| casd_channel.close() |
| casd_manager.release_resources() |
| |
| |
| @click.command(short_help="CAS Artifact Server") |
| @click.option("--port", "-p", type=click.INT, required=True, help="Port number") |
| @click.option("--server-key", help="Private server key for TLS (PEM-encoded)") |
| @click.option("--server-cert", help="Public server certificate for TLS (PEM-encoded)") |
| @click.option("--client-certs", help="Public client certificates for TLS (PEM-encoded)") |
| @click.option("--enable-push", is_flag=True, help="Allow clients to upload blobs and update artifact cache") |
| @click.option("--quota", type=click.INT, default=10e9, show_default=True, help="Maximum disk usage in bytes") |
| @click.option( |
| "--index-only", |
| is_flag=True, |
| help='Only provide the BuildStream artifact and source services ("index"), not the CAS ("storage")', |
| ) |
| @click.option("--log-level", type=LogLevel(), help="The log level to launch with", default="warning") |
| @click.argument("repo") |
| def server_main(repo, port, server_key, server_cert, client_certs, enable_push, quota, index_only, log_level): |
| # Handle SIGTERM by calling sys.exit(0), which will raise a SystemExit exception, |
| # properly executing cleanup code in `finally` clauses and context managers. |
| # This is required to terminate buildbox-casd on SIGTERM. |
| signal.signal(signal.SIGTERM, lambda signalnum, frame: sys.exit(0)) |
| |
| with create_server( |
| repo, quota=quota, enable_push=enable_push, index_only=index_only, log_level=log_level |
| ) as server: |
| |
| use_tls = bool(server_key) |
| |
| if bool(server_cert) != use_tls: |
| click.echo("ERROR: --server-key and --server-cert are both required for TLS", err=True) |
| sys.exit(-1) |
| |
| if client_certs and not use_tls: |
| click.echo("ERROR: --client-certs can only be used with --server-key", err=True) |
| sys.exit(-1) |
| |
| if use_tls: |
| # Read public/private key pair |
| with open(server_key, "rb") as f: |
| server_key_bytes = f.read() |
| with open(server_cert, "rb") as f: |
| server_cert_bytes = f.read() |
| |
| if client_certs: |
| with open(client_certs, "rb") as f: |
| client_certs_bytes = f.read() |
| else: |
| client_certs_bytes = None |
| |
| credentials = grpc.ssl_server_credentials( |
| [(server_key_bytes, server_cert_bytes)], |
| root_certificates=client_certs_bytes, |
| require_client_auth=bool(client_certs), |
| ) |
| server.add_secure_port("[::]:{}".format(port), credentials) |
| else: |
| server.add_insecure_port("[::]:{}".format(port)) |
| |
| # Run artifact server |
| server.start() |
| try: |
| while True: |
| signal.pause() |
| finally: |
| server.stop(0) |
| |
| |
| class _ByteStreamServicer(bytestream_pb2_grpc.ByteStreamServicer): |
| def __init__(self, casd, *, enable_push): |
| super().__init__() |
| self.bytestream = casd.get_bytestream() |
| self.enable_push = enable_push |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| |
| def Read(self, request, context): |
| self.logger.debug("Reading %s", request.resource_name) |
| return self.bytestream.Read(request) |
| |
| def Write(self, request_iterator, context): |
| # Note that we can't easily give more information because the |
| # data is stuck in an iterator that will be consumed if read. |
| self.logger.debug("Writing data") |
| return self.bytestream.Write(request_iterator) |
| |
| |
| class _ContentAddressableStorageServicer(remote_execution_pb2_grpc.ContentAddressableStorageServicer): |
| def __init__(self, casd, *, enable_push): |
| super().__init__() |
| self.cas = casd.get_cas() |
| self.enable_push = enable_push |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| |
| def FindMissingBlobs(self, request, context): |
| self.logger.info("Finding '%s'", request.blob_digests) |
| return self.cas.FindMissingBlobs(request) |
| |
| def BatchReadBlobs(self, request, context): |
| self.logger.info("Reading '%s'", request.digests) |
| return self.cas.BatchReadBlobs(request) |
| |
| def BatchUpdateBlobs(self, request, context): |
| self.logger.info("Updating: '%s'", [request.digest for request in request.requests]) |
| return self.cas.BatchUpdateBlobs(request) |
| |
| |
| class _CapabilitiesServicer(remote_execution_pb2_grpc.CapabilitiesServicer): |
| def __init__(self): |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| |
| def GetCapabilities(self, request, context): |
| self.logger.info("Retrieving capabilities") |
| response = remote_execution_pb2.ServerCapabilities() |
| |
| cache_capabilities = response.cache_capabilities |
| cache_capabilities.digest_function.append(remote_execution_pb2.DigestFunction.SHA256) |
| cache_capabilities.action_cache_update_capabilities.update_enabled = False |
| cache_capabilities.max_batch_total_size_bytes = _MAX_PAYLOAD_BYTES |
| cache_capabilities.symlink_absolute_path_strategy = remote_execution_pb2.SymlinkAbsolutePathStrategy.ALLOWED |
| |
| response.deprecated_api_version.major = 2 |
| response.low_api_version.major = 2 |
| response.high_api_version.major = 2 |
| |
| return response |
| |
| |
| class _ReferenceStorageServicer(buildstream_pb2_grpc.ReferenceStorageServicer): |
| def __init__(self, casd, cas_root, *, enable_push): |
| super().__init__() |
| self.cas = casd.get_cas() |
| self.root = cas_root |
| self.enable_push = enable_push |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| self.tmpdir = os.path.join(self.root, "tmp") |
| self.casdir = os.path.join(self.root, "cas") |
| self.refdir = os.path.join(self.casdir, "refs", "heads") |
| os.makedirs(self.tmpdir, exist_ok=True) |
| |
| # ref_path(): |
| # |
| # Get the path to a digest's file. |
| # |
| # Args: |
| # ref - The ref of the digest. |
| # |
| # Returns: |
| # str - The path to the digest's file. |
| # |
| def ref_path(self, ref: str) -> str: |
| return os.path.join(self.refdir, ref) |
| |
| # set_ref(): |
| # |
| # Create or update a ref with a new digest. |
| # |
| # Args: |
| # ref - The ref of the digest. |
| # tree - The digest to write. |
| # |
| def set_ref(self, ref: str, tree): |
| ref_path = self.ref_path(ref) |
| |
| os.makedirs(os.path.dirname(ref_path), exist_ok=True) |
| with save_file_atomic(ref_path, "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 |
| # |
| # Returns: |
| # (Digest): The digest stored in the ref |
| # |
| def resolve_ref(self, ref): |
| ref_path = self.ref_path(ref) |
| |
| with open(ref_path, "rb") as f: |
| os.utime(ref_path) |
| |
| digest = remote_execution_pb2.Digest() |
| digest.ParseFromString(f.read()) |
| return digest |
| |
| def GetReference(self, request, context): |
| self.logger.debug("'%s'", request.key) |
| response = buildstream_pb2.GetReferenceResponse() |
| |
| try: |
| digest = self.resolve_ref(request.key) |
| except FileNotFoundError: |
| with contextlib.suppress(FileNotFoundError): |
| _remove_path_with_parents(self.refdir, request.key) |
| |
| context.set_code(grpc.StatusCode.NOT_FOUND) |
| return response |
| |
| response.digest.hash = digest.hash |
| response.digest.size_bytes = digest.size_bytes |
| |
| return response |
| |
| def UpdateReference(self, request, context): |
| self.logger.debug("%s -> %s", request.keys, request.digest) |
| response = buildstream_pb2.UpdateReferenceResponse() |
| |
| if not self.enable_push: |
| context.set_code(grpc.StatusCode.PERMISSION_DENIED) |
| return response |
| |
| for key in request.keys: |
| self.set_ref(key, request.digest) |
| |
| return response |
| |
| def Status(self, request, context): |
| self.logger.debug("Retrieving status") |
| response = buildstream_pb2.StatusResponse() |
| |
| response.allow_updates = self.enable_push |
| |
| return response |
| |
| |
| class _ArtifactServicer(artifact_pb2_grpc.ArtifactServiceServicer): |
| def __init__(self, casd, root, *, update_cas=True): |
| super().__init__() |
| self.cas = casd.get_cas() |
| self.local_cas = casd.get_local_cas() |
| self.root = root |
| self.artifactdir = os.path.join(root, "artifacts", "refs") |
| self.update_cas = update_cas |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| |
| # object_path(): |
| # |
| # Get the path to an object's file. |
| # |
| # Args: |
| # digest - The digest of the object. |
| # |
| # Returns: |
| # str - The path to the object's file. |
| # |
| def object_path(self, digest) -> str: |
| return os.path.join(self.root, "cas", "objects", digest.hash[:2], digest.hash[2:]) |
| |
| # resolve_digest(): |
| # |
| # Read the directory corresponding to a digest. |
| # |
| # Args: |
| # digest - The digest corresponding to a directory. |
| # |
| # Returns: |
| # remote_execution_pb2.Directory - The directory. |
| # |
| # Raises: |
| # FileNotFoundError - If the digest object doesn't exist. |
| def resolve_digest(self, digest): |
| directory = remote_execution_pb2.Directory() |
| with open(self.object_path(digest), "rb") as f: |
| directory.ParseFromString(f.read()) |
| return directory |
| |
| def GetArtifact(self, request, context): |
| self.logger.info("'%s'", request.cache_key) |
| artifact_path = os.path.join(self.artifactdir, request.cache_key) |
| if not os.path.exists(artifact_path): |
| context.abort(grpc.StatusCode.NOT_FOUND, "Artifact proto not found") |
| |
| artifact = artifact_pb2.Artifact() |
| with open(artifact_path, "rb") as f: |
| artifact.ParseFromString(f.read()) |
| |
| os.utime(artifact_path) |
| |
| # Artifact-only servers will not have blobs on their system, |
| # so we can't reasonably update their mtimes. Instead, we exit |
| # early, and let the CAS server deal with its blobs. |
| # |
| # FIXME: We could try to run FindMissingBlobs on the other |
| # server. This is tricky to do from here, of course, |
| # because we don't know who the other server is, but |
| # the client could be smart about it - but this might |
| # make things slower. |
| # |
| # It needs some more thought... |
| if not self.update_cas: |
| return artifact |
| |
| # Now update mtimes of files present. |
| try: |
| |
| if str(artifact.files): |
| request = local_cas_pb2.FetchTreeRequest() |
| request.root_digest.CopyFrom(artifact.files) |
| request.fetch_file_blobs = True |
| self.local_cas.FetchTree(request) |
| |
| if str(artifact.buildtree): |
| try: |
| request = local_cas_pb2.FetchTreeRequest() |
| request.root_digest.CopyFrom(artifact.buildtree) |
| request.fetch_file_blobs = True |
| self.local_cas.FetchTree(request) |
| except grpc.RpcError as err: |
| # buildtrees might not be there |
| if err.code() != grpc.StatusCode.NOT_FOUND: |
| raise |
| |
| if str(artifact.public_data): |
| request = remote_execution_pb2.FindMissingBlobsRequest() |
| d = request.blob_digests.add() |
| d.CopyFrom(artifact.public_data) |
| self.cas.FindMissingBlobs(request) |
| |
| request = remote_execution_pb2.FindMissingBlobsRequest() |
| for log_file in artifact.logs: |
| d = request.blob_digests.add() |
| d.CopyFrom(log_file.digest) |
| self.cas.FindMissingBlobs(request) |
| |
| except grpc.RpcError as err: |
| if err.code() == grpc.StatusCode.NOT_FOUND: |
| os.unlink(artifact_path) |
| context.abort(grpc.StatusCode.NOT_FOUND, "Artifact files incomplete") |
| else: |
| context.abort(grpc.StatusCode.NOT_FOUND, "Artifact files not valid") |
| |
| return artifact |
| |
| def UpdateArtifact(self, request, context): |
| self.logger.info("'%s' -> '%s'", request.artifact, request.cache_key) |
| artifact = request.artifact |
| |
| if self.update_cas: |
| # Check that the files specified are in the CAS |
| if str(artifact.files): |
| self._check_directory("files", artifact.files, context) |
| |
| # Unset protocol buffers don't evaluated to False but do return empty |
| # strings, hence str() |
| if str(artifact.public_data): |
| self._check_file("public data", artifact.public_data, context) |
| |
| for log_file in artifact.logs: |
| self._check_file("log digest", log_file.digest, context) |
| |
| # Add the artifact proto to the cas |
| artifact_path = os.path.join(self.artifactdir, request.cache_key) |
| os.makedirs(os.path.dirname(artifact_path), exist_ok=True) |
| with save_file_atomic(artifact_path, mode="wb") as f: |
| f.write(artifact.SerializeToString()) |
| |
| return artifact |
| |
| def _check_directory(self, name, digest, context): |
| try: |
| self.resolve_digest(digest) |
| except FileNotFoundError: |
| self.logger.warning("Artifact %s specified but no files found", name) |
| context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Artifact {} specified but no files found".format(name)) |
| except DecodeError: |
| self.logger.warning("Artifact %s specified but directory not found", name) |
| context.abort( |
| grpc.StatusCode.FAILED_PRECONDITION, "Artifact {} specified but directory not found".format(name) |
| ) |
| |
| def _check_file(self, name, digest, context): |
| if not os.path.exists(self.object_path(digest)): |
| context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Artifact {} specified but not found".format(name)) |
| |
| |
| class _BuildStreamCapabilitiesServicer(buildstream_pb2_grpc.CapabilitiesServicer): |
| def __init__(self, artifact_capabilities, source_capabilities): |
| self.artifact_capabilities = artifact_capabilities |
| self.source_capabilities = source_capabilities |
| |
| def GetCapabilities(self, request, context): |
| response = buildstream_pb2.ServerCapabilities() |
| response.artifact_capabilities.CopyFrom(self.artifact_capabilities) |
| response.source_capabilities.CopyFrom(self.source_capabilities) |
| return response |
| |
| |
| class _SourceServicer(source_pb2_grpc.SourceServiceServicer): |
| def __init__(self, sourcedir): |
| self.sourcedir = sourcedir |
| self.logger = logging.getLogger("buildstream._cas.casserver") |
| |
| def GetSource(self, request, context): |
| self.logger.info("'%s'", request.cache_key) |
| try: |
| source_proto = self._get_source(request.cache_key) |
| except FileNotFoundError: |
| context.abort(grpc.StatusCode.NOT_FOUND, "Source not found") |
| except DecodeError: |
| context.abort(grpc.StatusCode.NOT_FOUND, "Sources gives invalid directory") |
| |
| return source_proto |
| |
| def UpdateSource(self, request, context): |
| self.logger.info("'%s' -> '%s'", request.source, request.cache_key) |
| self._set_source(request.cache_key, request.source) |
| return request.source |
| |
| def _get_source(self, cache_key): |
| path = os.path.join(self.sourcedir, cache_key) |
| source_proto = source_pb2.Source() |
| with open(path, "r+b") as f: |
| source_proto.ParseFromString(f.read()) |
| os.utime(path) |
| return source_proto |
| |
| def _set_source(self, cache_key, source_proto): |
| path = os.path.join(self.sourcedir, cache_key) |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| with save_file_atomic(path, "w+b") as f: |
| f.write(source_proto.SerializeToString()) |
