buildgrid/_app/commands/cmd_operation.py

@@ -155,6 +155,19 @@ def status(context, operation_name, json):
         click.echo(json_format.MessageToJson(operation))
 
 
+@cli.command('cancel', short_help="Cancel an operation.")
+@click.argument('operation-name', nargs=1, type=click.STRING, required=True)
+@pass_context
+def cancel(context, operation_name):
+    context.logger.info("Cancelling an operation...")
+    stub = operations_pb2_grpc.OperationsStub(context.channel)
+
+    request = operations_pb2.CancelOperationRequest(name=operation_name)
+
+    stub.CancelOperation(request)
+    context.logger.info("Operation cancelled: [{}]".format(request))
+
+
 @cli.command('list', short_help="List operations.")
 @click.option('--json', is_flag=True, show_default=True,
               help="Print operations list in JSON format.")

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 FailedPreconditionError, InvalidArgumentError, CancelledError
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
 from buildgrid._protos.google.longrunning import operations_pb2
 
@@ -76,6 +76,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("/")
@@ -106,6 +112,12 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
             context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
             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 _get_instance(self, name):
         try:
             return self._instances[name]

buildgrid/server/job.py

@@ -19,6 +19,7 @@ import uuid
 from google.protobuf import timestamp_pb2
 
 from buildgrid._enums import LeaseState, OperationStage
+from buildgrid._exceptions import CancelledError
 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
@@ -37,10 +38,14 @@ class Job:
 
         self.__execute_response = None
         self.__operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
+
         self.__queued_timestamp = timestamp_pb2.Timestamp()
         self.__worker_start_timestamp = timestamp_pb2.Timestamp()
         self.__worker_completed_timestamp = timestamp_pb2.Timestamp()
 
+        self.__operation_cancelled = False
+        self.__lease_cancelled = False
+
         self.__operation_metadata.action_digest.CopyFrom(action_digest)
         self.__operation_metadata.stage = OperationStage.UNKNOWN.value
 
@@ -131,7 +136,9 @@ class Job:
         Only one :class:`Lease` can be emitted for a given job. This method
         should only be used once, any furhter calls are ignored.
         """
-        if self._lease is not None:
+        if self.__operation_cancelled:
+            return None
+        elif self._lease is not None:
             return None
 
         self._lease = bots_pb2.Lease()
@@ -189,6 +196,15 @@ class Job:
             self.__execute_response.cached_result = False
             self.__execute_response.status.CopyFrom(status)
 
+    def cancel_lease(self):
+        """Triggers a job's :class:Lease cancellation.
+
+        This will not cancel the job's :class:Operation.
+        """
+        self.__lease_cancelled = True
+        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.
 
@@ -214,3 +230,20 @@ class Job:
 
         for queue in self._operation_update_queues:
             queue.put(self._operation)
+
+    def cancel_operation(self):
+        """Triggers a job's :class:Operation cancellation.
+
+        This will also cancel any job's :class:Lease that may have been issued.
+        """
+        self.__operation_cancelled = 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 = "Operation cancelled by client."
+
+        self.update_operation_stage(OperationStage.COMPLETED)
+
+        raise CancelledError("Operation cancelled: {}".format(self._name))

buildgrid/server/operations/instance.py

@@ -64,6 +64,13 @@ class OperationsInstance:
         except KeyError:
             raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
 
+    def cancel_operation(self, name):
+        try:
+            self._scheduler.cancel_job_operation(name)
+
+        except KeyError:
+            raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
+
     def register_message_client(self, name, queue):
         try:
             self._scheduler.register_client(name, queue)
@@ -84,7 +91,3 @@ class OperationsInstance:
             yield operation
             operation = message_queue.get()
         yield operation
-
-    def cancel_operation(self, name):
-        # TODO: Cancel leases
-        raise NotImplementedError("Cancelled operations not supported")

buildgrid/server/operations/service.py

@@ -25,7 +25,7 @@ import grpc
 
 from google.protobuf.empty_pb2 import Empty
 
-from buildgrid._exceptions import InvalidArgumentError
+from buildgrid._exceptions import CancelledError, InvalidArgumentError
 from buildgrid._protos.google.longrunning import operations_pb2_grpc, operations_pb2
 
 
@@ -112,10 +112,10 @@ class OperationsService(operations_pb2_grpc.OperationsServicer):
             operation_name = self._parse_operation_name(name)
             instance.cancel_operation(operation_name)
 
-        except NotImplementedError as e:
+        except CancelledError as e:
             self.logger.error(e)
             context.set_details(str(e))
-            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_code(grpc.StatusCode.CANCELLED)
 
         except InvalidArgumentError as e:
             self.logger.error(e)

buildgrid/server/scheduler.py

@@ -94,7 +94,10 @@ class Scheduler:
         # For now, one lease at a time:
         lease = job.create_lease()
 
-        return [lease]
+        if lease:
+            return [lease]
+
+        return None
 
     def update_job_lease_state(self, job_name, lease_state, lease_status=None, lease_result=None):
         """Requests a state transition for a job's current :class:Lease.
@@ -133,3 +136,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()

tests/integration/operations_service.py

@@ -24,6 +24,7 @@ import grpc
 from grpc._server import _Context
 import pytest
 
+from buildgrid._enums import OperationStage
 from buildgrid._exceptions import InvalidArgumentError
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from buildgrid._protos.google.longrunning import operations_pb2
@@ -236,12 +237,26 @@ def test_delete_operation_fail(instance, context):
     context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
 
 
-def test_cancel_operation(instance, context):
+def test_cancel_operation(instance, controller, execute_request, context):
+    response_execute = controller.execution_instance.execute(execute_request.action_digest,
+                                                             execute_request.skip_cache_lookup)
+
     request = operations_pb2.CancelOperationRequest()
-    request.name = "{}/{}".format(instance_name, "runner")
+    request.name = "{}/{}".format(instance_name, response_execute.name)
+
     instance.CancelOperation(request, context)
 
-    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_code.assert_called_once_with(grpc.StatusCode.CANCELLED)
+
+    request = operations_pb2.ListOperationsRequest(name=instance_name)
+    response = instance.ListOperations(request, context)
+
+    assert len(response.operations) is 1
+
+    for operation in response.operations:
+        operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
+        operation.metadata.Unpack(operation_metadata)
+        assert operation_metadata.stage == OperationStage.COMPLETED.value
 
 
 def test_cancel_operation_blank(blank_instance, context):
@@ -249,7 +264,7 @@ def test_cancel_operation_blank(blank_instance, context):
     request.name = "runner"
     blank_instance.CancelOperation(request, context)
 
-    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
 
 
 def test_cancel_operation_instance_fail(instance, context):