buildgrid/_exceptions.py

@@ -52,6 +52,11 @@ class BotError(BgdError):
         super().__init__(message, detail=detail, domain=ErrorDomain.BOT, reason=reason)
 
 
+class CancelledError(BgdError):
+    def __init__(self, message, detail=None, reason=None):
+        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
+
+
 class InvalidArgumentError(BgdError):
     """A bad argument was passed, such as a name which doesn't exist."""
     def __init__(self, message, detail=None, reason=None):

buildgrid/server/execution/instance.py

@@ -21,8 +21,9 @@ An instance of the Remote Execution Service.
 
 import logging
 
-from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError
-from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Action
+from buildgrid._exceptions import CancelledError, FailedPreconditionError, InvalidArgumentError
+from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
+from buildgrid._protos.google.rpc import code_pb2
 
 from ..job import Job
 
@@ -43,7 +44,7 @@ class ExecutionInstance:
         this action.
         """
 
-        action = self._storage.get_message(action_digest, Action)
+        action = self._storage.get_message(action_digest, remote_execution_pb2.Action)
 
         if not action:
             raise FailedPreconditionError("Could not get action from storage.")
@@ -77,4 +78,11 @@ class ExecutionInstance:
         while not operation.done:
             yield operation
             operation = message_queue.get()
+
+        response = remote_execution_pb2.ExecuteResponse()
+        operation.response.Unpack(response)
+
+        if response.status.code == code_pb2.CANCELLED:
+            raise CancelledError(response.status.message)
+
         yield operation

buildgrid/server/execution/service.py

@@ -26,7 +26,7 @@ from functools import partial
 
 import grpc
 
-from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError
+from buildgrid._exceptions import CancelledError, FailedPreconditionError, InvalidArgumentError
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
 from buildgrid._protos.google.longrunning import operations_pb2
 
@@ -67,6 +67,12 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
             context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
             yield operations_pb2.Operation()
 
+        except CancelledError as e:
+            self.logger.error(e)
+            context.set_details(str(e))
+            context.set_code(grpc.StatusCode.CANCELLED)
+            yield operations_pb2.Operation()
+
     def WaitExecution(self, request, context):
         try:
             names = request.name.split("/")

buildgrid/server/job.py

@@ -20,6 +20,7 @@ from enum import Enum
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 from buildgrid._protos.google.longrunning import operations_pb2
+from buildgrid._protos.google.rpc import code_pb2
 
 
 class OperationStage(Enum):
@@ -60,6 +61,7 @@ class Job:
 
         self.__execute_response = None
         self.__operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
+        self.__operation_canceled = False
 
         self.__operation_metadata.action_digest.CopyFrom(action_digest)
         self.__operation_metadata.stage = OperationStage.UNKNOWN.value
@@ -196,6 +198,10 @@ class Job:
             self.__execute_response.cached_result = False
             self.__execute_response.status.CopyFrom(status)
 
+    def cancel_lease(self):
+        if self._lease is not None:
+            self.update_lease_state(LeaseState.CANCELLED)
+
     def update_operation_stage(self, stage):
         """Operates a stage transition for the job's :class:Operation.
 
@@ -219,3 +225,14 @@ class Job:
 
         for queue in self._operation_update_queues:
             queue.put(self._operation)
+
+    def cancel_operation(self):
+        self.__operation_canceled = True
+        if self._lease is not None:
+            self.cancel_lease()
+
+        self.__execute_response = remote_execution_pb2.ExecuteResponse()
+        self.__execute_response.status.code = code_pb2.CANCELLED
+        self.__execute_response.status.message = "The operation was cancelled by the caller."
+
+        self.update_operation_stage(OperationStage.COMPLETED)

buildgrid/server/scheduler.py

@@ -128,3 +128,13 @@ class Scheduler:
     def get_job_operation(self, job_name):
         """Returns the operation associated to job."""
         return self.jobs[job_name].operation
+
+    def cancel_job_operation(self, job_name):
+        """"Cancels the underlying operation of a given job.
+
+        This will also cancel any job's lease that may have been issued.
+
+        Args:
+            job_name (str): name of the job holding the operation to cancel.
+        """
+        self.jobs[job_name].cancel_operation()