vm: make async operation tasks abortable

circle/vm/operations.py

@@ -9,7 +9,9 @@ from django.utils.translation import ugettext_lazy as _
 from celery.exceptions import TimeLimitExceeded
 
 from common.operations import Operation, register_operation
-from .tasks.local_tasks import async_instance_operation, async_node_operation
+from .tasks.local_tasks import (
+    abortable_async_instance_operation, abortable_async_node_operation,
+)
 from .models import (
     Instance, InstanceActivity, InstanceTemplate, Interface, Node,
     NodeActivity,
@@ -21,7 +23,7 @@ logger = getLogger(__name__)
 
 class InstanceOperation(Operation):
     acl_level = 'owner'
-    async_operation = async_instance_operation
+    async_operation = abortable_async_instance_operation
     host_cls = Instance
 
     def __init__(self, instance):
@@ -440,7 +442,7 @@ register_operation(WakeUpOperation)
 
 
 class NodeOperation(Operation):
-    async_operation = async_node_operation
+    async_operation = abortable_async_node_operation
     host_cls = Node
 
     def __init__(self, node):

circle/vm/tasks/local_tasks.py

+from celery.contrib.abortable import AbortableTask
 from manager.mancelery import celery
 
 
-@celery.task
-def async_instance_operation(operation_id, instance_pk, activity_pk, allargs,
-                             auxargs):
+@celery.task(base=AbortableTask, bind=True)
+def abortable_async_instance_operation(task, operation_id, instance_pk,
+                                       activity_pk, allargs, auxargs):
     from vm.models import Instance, InstanceActivity
     instance = Instance.objects.get(pk=instance_pk)
     operation = getattr(instance, operation_id)
     activity = InstanceActivity.objects.get(pk=activity_pk)
 
     # save async task UUID to activity
-    activity.task_uuid = async_instance_operation.request.id
+    activity.task_uuid = task.request.id
     activity.save()
 
     allargs['activity'] = activity
+    allargs['task'] = task
 
     return operation._exec_op(allargs, auxargs)
 
 
-@celery.task
-def async_node_operation(operation_id, node_pk, activity_pk, allargs, auxargs):
+@celery.task(base=AbortableTask, bind=True)
+def abortable_async_node_operation(task, operation_id, node_pk, activity_pk,
+                                   allargs, auxargs):
     from vm.models import Node, NodeActivity
     node = Node.objects.get(pk=node_pk)
     operation = getattr(node, operation_id)
     activity = NodeActivity.objects.get(pk=activity_pk)
 
     # save async task UUID to activity
-    activity.task_uuid = async_node_operation.request.id
+    activity.task_uuid = task.request.id
     activity.save()
 
     allargs['activity'] = activity
+    allargs['task'] = task
 
     return operation._exec_op(allargs, auxargs)