diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py index 68fd4b610..61a84b694 100644 --- a/src/buildstream/_cas/cascache.py +++ b/src/buildstream/_cas/cascache.py @@ -127,7 +127,13 @@ def contains_directory(self, digest): request.fetch_file_blobs = not self._remote_cache try: - local_cas.FetchTree(request) + response_future = local_cas.FetchTree.future(request) + try: + response_future.result() + except: + response_future.cancel() + raise + if not self._remote_cache: return True except grpc.RpcError as e: @@ -141,7 +147,13 @@ def contains_directory(self, digest): request = local_cas_pb2.UploadTreeRequest() request.root_digest.CopyFrom(digest) try: - local_cas.UploadTree(request) + response_future = local_cas.UploadTree.future(request) + try: + response_future.result() + except: + response_future.cancel() + raise + return True except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: @@ -231,7 +243,12 @@ def ensure_tree(self, tree): request.root_digest.CopyFrom(tree) request.fetch_file_blobs = True - local_cas.FetchTree(request) + response_future = local_cas.FetchTree.future(request) + try: + response_future.result() + except: + response_future.cancel() + raise # fetch_directory(): # @@ -252,7 +269,13 @@ def fetch_directory(self, remote, dir_digest): request.fetch_file_blobs = False try: - local_cas.FetchTree(request) + response_future = local_cas.FetchTree.future(request) + try: + response_future.result() + except: + response_future.cancel() + raise + except grpc.RpcError as e: if e.code() == grpc.StatusCode.NOT_FOUND: raise BlobNotFound( @@ -536,7 +559,13 @@ def missing_blobs(self, blobs, *, remote=None): d.CopyFrom(required_digest) try: - response = cas.FindMissingBlobs(request) + response_future = cas.FindMissingBlobs.future(request) + try: + response = response_future.result() + except: + response_future.cancel() + raise + except grpc.RpcError as e: if e.code() == grpc.StatusCode.INVALID_ARGUMENT and e.details().startswith("Invalid instance name"): raise CASCacheError("Unsupported buildbox-casd version: FindMissingBlobs failed") from e @@ -566,7 +595,12 @@ def required_blobs_for_directory(self, directory_digest, *, excluded_subdirs=Non request.root_digest.CopyFrom(directory_digest) request.fetch_file_blobs = False - local_cas.FetchTree(request) + response_future = local_cas.FetchTree.future(request) + try: + response_future.result() + except: + response_future.cancel() + raise # parse directory, and recursively add blobs diff --git a/src/buildstream/_cas/casremote.py b/src/buildstream/_cas/casremote.py index 3fefb9d47..f04d8d7be 100644 --- a/src/buildstream/_cas/casremote.py +++ b/src/buildstream/_cas/casremote.py @@ -92,7 +92,13 @@ def send(self, *, missing_blobs=None): local_cas = self._remote.casd.get_local_cas() for request in self._requests: - batch_response = local_cas.FetchMissingBlobs(request) + batch_response_future = local_cas.FetchMissingBlobs.future(request) + + try: + batch_response = batch_response_future.result() + except: + batch_response_future.cancel() + raise for response in batch_response.responses: if response.status.code == code_pb2.NOT_FOUND: @@ -146,7 +152,13 @@ def send(self): local_cas = self._remote.casd.get_local_cas() for request in self._requests: - batch_response = local_cas.UploadMissingBlobs(request) + batch_response_future = local_cas.UploadMissingBlobs.future(request) + + try: + batch_response = batch_response_future.result() + except: + batch_response_future.cancel() + raise for response in batch_response.responses: if response.status.code != code_pb2.OK: