job.py: Implement operation cancellation

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,17 @@ class Job:
 
         for queue in self._operation_update_queues:
             queue.put(self._operation)
+
+    def cancel_operation(self):
+        self.__operation_canceled = True
+
+        self.cancel_lease()
+
+        response = remote_execution_pb2.ExecuteResponse()
+        response.status.code = code_pb2.CANCELLED
+        response.status.message = "The operation was cancelled by the caller."
+
+        self._operation.response.Pack(response)
+        self._operation.done = True
+
+        self.update_operation_stage(OperationStage.COMPLETED)