def cancel(task_id):
    """
    Cancel the task that is represented by the given task_id. This method cancels only the task
    with given task_id, not the spawned tasks. This also updates task's state to 'canceled'.

    :param task_id: The ID of the task you wish to cancel
    :type  task_id: basestring

    :raises MissingResource: if a task with given task_id does not exist
    :raises PulpCodedException: if given task is already in a complete state
    """
    task_status = TaskStatusManager.find_by_task_id(task_id)
    if task_status is None:
        raise MissingResource(task_id)
    if task_status['state'] in constants.CALL_COMPLETE_STATES:
        # If the task is already done, just stop
        msg = _('Task [%(task_id)s] already in a completed state: %(state)s')
        logger.info(msg % {'task_id': task_id, 'state': task_status['state']})
        return
    controller.revoke(task_id, terminate=True)
    TaskStatus.get_collection().find_and_modify(
        {'task_id': task_id, 'state': {'$nin': constants.CALL_COMPLETE_STATES}},
        {'$set': {'state': constants.CALL_CANCELED_STATE}})
    msg = _('Task canceled: %(task_id)s.')
    msg = msg % {'task_id': task_id}
    logger.info(msg)

    def apply_async(self, *args, **kwargs):
        """
        A wrapper around the Celery apply_async method. It allows us to accept a few more
        parameters than Celery does for our own purposes, listed below. It also allows us
        to create and update task status which can be used to track status of this task
        during it's lifetime.

        :param queue:       The queue that the task has been placed into (optional, defaults to
                            the general Celery queue.)
        :type  queue:       basestring
        :param tags:        A list of tags (strings) to place onto the task, used for searching for
                            tasks by tag
        :type  tags:        list
        :return:            An AsyncResult instance as returned by Celery's apply_async
        :rtype:             celery.result.AsyncResult
        """
        queue = kwargs.get('queue', defaults.NAMESPACES['CELERY']['DEFAULT_QUEUE'].default)
        tags = kwargs.pop('tags', [])
        async_result = super(Task, self).apply_async(*args, **kwargs)
        async_result.tags = tags

        # Create a new task status with the task id and tags.
        # To avoid the race condition where __call__ method below is called before
        # this change is propagated to all db nodes, using an 'upsert' here and setting
        # the task state to 'waiting' only on an insert.
        TaskStatus.get_collection().update(
            {'task_id': async_result.id},
            {'$setOnInsert': {'state':dispatch_constants.CALL_WAITING_STATE},
             '$set': {'queue': queue, 'tags': tags}},
            upsert=True)
        return async_result

 def __call__(self, *args, **kwargs):
     """
     This overrides CeleryTask's __call__() method. We use this method
     for task state tracking of Pulp tasks.
     """
     # Add task_id to the task context, so that agent and plugins have access to the task id.
     # There are a few other attributes in the context as defined by old dispatch system.
     # These are unused right now. These should be removed when we cleanup the dispatch folder
     # after the migration to celery is complete.
     task_context = dispatch_factory.context()
     task_context.call_request_id = self.request.id
     # Check task status and skip running the task if task state is 'canceled'.
     task_status = TaskStatusManager.find_by_task_id(task_id=self.request.id)
     if task_status and task_status['state'] == dispatch_constants.CALL_CANCELED_STATE:
         logger.debug("Task cancel received for task-id : [%s]" % self.request.id)
         return
     # Update start_time and set the task state to 'running' for asynchronous tasks.
     # Skip updating status for eagerly executed tasks, since we don't want to track
     # synchronous tasks in our database.
     if not self.request.called_directly:
         # Using 'upsert' to avoid a possible race condition described in the apply_async method
         # above.
         TaskStatus.get_collection().update(
             {'task_id': self.request.id},
             {'$set': {'state': dispatch_constants.CALL_RUNNING_STATE,
                       'start_time': dateutils.now_utc_timestamp()}},
             upsert=True)
     # Run the actual task
     logger.debug("Running task : [%s]" % self.request.id)
     return super(Task, self).__call__(*args, **kwargs)

    def test_task_status_update_fires_notification(self, mock_send):
        """
        Test that update_one() also fires a notification.
        """
        task_id = self.get_random_uuid()
        worker_name = 'special_worker_name'
        tags = ['test-tag1', 'test-tag2']
        state = 'waiting'
        ts = TaskStatus(task_id, worker_name, tags, state)
        ts.save()
        # ensure event was fired for save()
        mock_send.assert_called_once_with(ts, routing_key="tasks.%s" % task_id)
        now = datetime.now(dateutils.utc_tz())
        start_time = dateutils.format_iso8601_datetime(now)
        delta = {'start_time': start_time,
                 'state': 'running',
                 'progress_report': {'report-id': 'my-progress'}}

        self.assertEquals(len(mock_send.call_args_list), 1)
        TaskStatus.objects(task_id=task_id).update_one(
            set__start_time=delta['start_time'], set__state=delta['state'],
            set__progress_report=delta['progress_report'])

        # ensure event was fired for update_one()
        self.assertEquals(len(mock_send.call_args_list), 2)
        mock_send.assert_called_with(ts, routing_key="tasks.%s" % task_id)

    def failed(self, reply):
        """
        Notification (reply) indicating an RMI failed.
        This information used to update the task status.
        :param reply: A failure reply object.
        :type reply: gofer.rmi.async.Failed
        """
        _logger.info(_('Task RMI (failed): %(r)s'), {'r': reply})

        call_context = dict(reply.data)
        action = call_context.get('action')
        task_id = call_context['task_id']
        traceback = reply.xstate['trace']
        finished = reply.timestamp
        if not finished:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)

        TaskStatus.objects(task_id=task_id).update_one(set__finish_time=finished,
                                                       set__state=constants.CALL_ERROR_STATE,
                                                       set__traceback=traceback)

        if action == 'bind':
            ReplyHandler._bind_failed(task_id, call_context)
            return
        if action == 'unbind':
            ReplyHandler._unbind_failed(task_id, call_context)
            return

    def test_task_status_update(self):
        """
        Tests the successful operation of task status update.
        """
        task_id = self.get_random_uuid()
        worker_name = 'special_worker_name'
        tags = ['test-tag1', 'test-tag2']
        state = 'waiting'
        TaskStatus(task_id, worker_name, tags, state).save()
        now = datetime.now(dateutils.utc_tz())
        start_time = dateutils.format_iso8601_datetime(now)
        delta = {'start_time': start_time,
                 'state': 'running',
                 'progress_report': {'report-id': 'my-progress'}}

        TaskStatus.objects(task_id=task_id).update_one(
            set__start_time=delta['start_time'], set__state=delta['state'],
            set__progress_report=delta['progress_report'])

        task_status = TaskStatus.objects(task_id=task_id).first()
        self.assertEqual(task_status['start_time'], delta['start_time'])
        # Make sure that parse_iso8601_datetime is able to parse the start_time without errors
        dateutils.parse_iso8601_datetime(task_status['start_time'])
        self.assertEqual(task_status['state'], delta['state'])
        self.assertEqual(task_status['progress_report'], delta['progress_report'])
        self.assertEqual(task_status['worker_name'], worker_name)

    def succeeded(self, reply):
        """
        Notification (reply) indicating an RMI succeeded.
        This information is relayed to the task coordinator.
        :param reply: A successful reply object.
        :type reply: gofer.rmi.async.Succeeded
        """
        _logger.info(_('Task RMI (succeeded): %(r)s'), {'r': reply})

        call_context = dict(reply.data)
        action = call_context.get('action')
        task_id = call_context['task_id']
        result = dict(reply.retval)
        finished = reply.timestamp
        if not finished:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)

        TaskStatus.objects(task_id=task_id).update_one(set__finish_time=finished,
                                                       set__state=constants.CALL_FINISHED_STATE,
                                                       set__result=result)
        if action == 'bind':
            if result['succeeded']:
                ReplyHandler._bind_succeeded(task_id, call_context)
            else:
                ReplyHandler._bind_failed(task_id, call_context)
            return
        if action == 'unbind':
            if result['succeeded']:
                ReplyHandler._unbind_succeeded(call_context)
            else:
                ReplyHandler._unbind_failed(task_id, call_context)
            return

def cancel(task_id):
    """
    Cancel the task that is represented by the given task_id. This method cancels only the task
    with given task_id, not the spawned tasks. This also updates task's state to 'canceled'.

    :param task_id: The ID of the task you wish to cancel
    :type  task_id: basestring

    :raises MissingResource: if a task with given task_id does not exist
    :raises PulpCodedException: if given task is already in a complete state
    """
    try:
        task_status = TaskStatus.objects.get(task_id=task_id)
    except DoesNotExist:
        raise MissingResource(task_id)
    if task_status['state'] in constants.CALL_COMPLETE_STATES:
        # If the task is already done, just stop
        msg = _('Task [%(task_id)s] already in a completed state: %(state)s')
        _logger.info(msg % {'task_id': task_id, 'state': task_status['state']})
        return
    controller.revoke(task_id, terminate=True)
    TaskStatus.objects(task_id=task_id, state__nin=constants.CALL_COMPLETE_STATES).\
        update_one(set__state=constants.CALL_CANCELED_STATE)
    msg = _('Task canceled: %(task_id)s.')
    msg = msg % {'task_id': task_id}
    _logger.info(msg)

 def __call__(self, *args, **kwargs):
     """
     This overrides CeleryTask's __call__() method. We use this method
     for task state tracking of Pulp tasks.
     """
     # Check task status and skip running the task if task state is 'canceled'.
     task_status = TaskStatusManager.find_by_task_id(task_id=self.request.id)
     if task_status and task_status['state'] == dispatch_constants.CALL_CANCELED_STATE:
         logger.debug("Task cancel received for task-id : [%s]" % self.request.id)
         return
     # Update start_time and set the task state to 'running' for asynchronous tasks.
     # Skip updating status for eagerly executed tasks, since we don't want to track
     # synchronous tasks in our database.
     if not self.request.called_directly:
         now = datetime.now(dateutils.utc_tz())
         start_time = dateutils.format_iso8601_datetime(now)
         # Using 'upsert' to avoid a possible race condition described in the apply_async method
         # above.
         TaskStatus.get_collection().update(
             {'task_id': self.request.id},
             {'$set': {'state': dispatch_constants.CALL_RUNNING_STATE,
                       'start_time': start_time}},
             upsert=True)
     # Run the actual task
     logger.debug("Running task : [%s]" % self.request.id)
     return super(Task, self).__call__(*args, **kwargs)

    def set_progress(self, status):
        """
        Informs the server of the current state of the publish operation. The
        contents of the status is dependent on how the distributor
        implementation chooses to divide up the publish process.

        @param status: contains arbitrary data to describe the state of the
               publish; the contents may contain whatever information is relevant
               to the distributor implementation so long as it is serializable
        """

        if self.task_id is None:
            # not running within a task
            return

        try:
            self.progress_report[self.report_id] = status
            TaskStatus.objects(task_id=self.task_id).update_one(set__progress_report=self.progress_report)
        except Exception, e:
            logger.exception('Exception from server setting progress for report [%s]' % self.report_id)
            try:
                logger.error('Progress value: %s' % str(status))
            except Exception:
                # Best effort to print this, but if its that grossly unserializable
                # the log will tank and we don't want that exception to bubble up
                pass
            raise self.exception_class(e), None, sys.exc_info()[2]

    def apply_async(self, *args, **kwargs):
        """
        A wrapper around the Celery apply_async method. It allows us to accept a few more
        parameters than Celery does for our own purposes, listed below. It also allows us
        to create and update task status which can be used to track status of this task
        during it's lifetime.

        :param queue:       The queue that the task has been placed into (optional, defaults to
                            the general Celery queue.)
        :type  queue:       basestring
        :param tags:        A list of tags (strings) to place onto the task, used for searching for
                            tasks by tag
        :type  tags:        list
        :return:            An AsyncResult instance as returned by Celery's apply_async
        :rtype:             celery.result.AsyncResult
        """
        routing_key = kwargs.get('routing_key',
                                 defaults.NAMESPACES['CELERY']['DEFAULT_ROUTING_KEY'].default)
        tags = kwargs.pop('tags', [])

        async_result = super(Task, self).apply_async(*args, **kwargs)
        async_result.tags = tags

        # Create a new task status with the task id and tags.
        task_status = TaskStatus(
            task_id=async_result.id, task_type=self.name,
            state=constants.CALL_WAITING_STATE, worker_name=routing_key, tags=tags)
        # To avoid the race condition where __call__ method below is called before
        # this change is propagated to all db nodes, using an 'upsert' here and setting
        # the task state to 'waiting' only on an insert.
        task_status.save_with_set_on_insert(fields_to_set_on_insert=['state', 'start_time'])
        return async_result

    def update_task_status(task_id, delta):
        """
        Updates status of the task with given task id. Only the following
        fields may be updated through this call:
        * state
        * result
        * traceback
        * start_time
        * finish_time
        * error
        * spawned_tasks
        * progress_report
        Other fields found in delta will be ignored.

        :param task_id: identity of the task this status corresponds to
        :type  task_id: basetring
        :param delta: list of attributes and their new values to change
        :type  delta: dict
        :return: updated serialized task status
        :rtype:  dict
        :raise MissingResource: if there is no task status corresponding to the given task_id
        """

        task_status = TaskStatus.get_collection().find_one({'task_id': task_id})
        if task_status is None:
            raise MissingResource(task_id)

        updatable_attributes = ['state', 'result', 'traceback', 'start_time', 'finish_time',
                                'error', 'spawned_tasks', 'progress_report']
        for key, value in delta.items():
            if key in updatable_attributes:
                task_status[key] = value

        TaskStatus.get_collection().save(task_status, safe=True)
        return task_status

    def test_set_accepted(self):
        task_id = self.get_random_uuid()
        TaskStatus(task_id, state=constants.CALL_WAITING_STATE).save()

        TaskStatus.objects(task_id=task_id, state=constants.CALL_WAITING_STATE).\
            update_one(set__state=constants.CALL_ACCEPTED_STATE)
        task_status = TaskStatus.objects.get(task_id=task_id)
        self.assertTrue(task_status['state'], constants.CALL_ACCEPTED_STATE)

 def progress(self, reply):
     """
     Notification (reply) indicating an RMI has reported status.
     This information is relayed to the task coordinator.
     :param reply: A progress reply object.
     :type reply: gofer.rmi.async.Progress
     """
     call_context = dict(reply.data)
     task_id = call_context['task_id']
     TaskStatus.objects(task_id=task_id).update_one(set__progress_report=reply.details)

 def GET(self):
     valid_filters = ['tag']
     filters = self.filters(valid_filters)
     tags = filters.get('tag', [])
     if tags:
         raw_tasks = TaskStatus.objects(tags__all=tags)
     else:
         raw_tasks = TaskStatus.objects()
     serialized_task_statuses = [task_serializer(task) for task in raw_tasks]
     return self.ok(serialized_task_statuses)

    def test_save_task_status_fires_notification(self, mock_send):
        """
        Test that saving a TaskStatus fires an event notification.
        """
        task_id = self.get_random_uuid()

        ts = TaskStatus(task_id)
        ts.save()

        mock_send.assert_called_once_with(ts, routing_key="tasks.%s" % task_id)

    def test_data_call(self, mock_current_task):
        mock_current_task.request.id = 'fake_id'
        fake_task_status = TaskStatus('fake_id')
        fake_task_status.save()
        data = {'event_type': 'test_type',
                'payload': 'test_payload',
                'call_report': fake_task_status}

        event = event_data.Event(data['event_type'], data['payload'])

        self.assertEqual(data, event.data())

    def test_save_update_with_set_on_insert(self):
        """
        Test the save method with set on insert arguments when the object is already in the
        database.
        """
        task_id = str(uuid4())
        worker_name = 'worker_name'
        tags = ['tag_1', 'tag_2']
        state = constants.CALL_ACCEPTED_STATE
        spawned_tasks = ['foo']
        error = {'error': 'some_error'}
        progress_report = {'what do we want?': 'progress!', 'when do we want it?': 'now!'}
        task_type = 'some.task'
        old_start_time = start_time = datetime.now()
        finish_time = start_time + timedelta(minutes=5)
        start_time = dateutils.format_iso8601_datetime(start_time)
        finish_time = dateutils.format_iso8601_datetime(finish_time)
        result = None
        ts = TaskStatus(
            task_id, worker_name, tags, state, spawned_tasks=spawned_tasks, error=error,
            progress_report=progress_report, task_type=task_type, start_time=start_time,
            finish_time=finish_time, result=result)
        # Put the object in the database, and then change some of it settings.
        ts.save()
        new_worker_name = 'a different_worker'
        new_state = constants.CALL_SUSPENDED_STATE
        new_start_time = old_start_time + timedelta(minutes=10)
        new_start_time = dateutils.format_iso8601_datetime(new_start_time)
        ts.worker_name = new_worker_name
        ts.state = new_state
        ts.start_time = new_start_time

        # This should update the worker_name on ts in the database, but should not update the state
        # or start_time
        ts.save_with_set_on_insert(fields_to_set_on_insert=['state', 'start_time'])

        ts = TaskStatus.objects()
        # There should only be one TaskStatus in the db
        self.assertEqual(len(ts), 1)
        ts = ts[0]
        # Make sure all the attributes are correct
        self.assertEqual(ts['task_id'], task_id)
        # Queue should have been updated
        self.assertEqual(ts['worker_name'], new_worker_name)
        self.assertEqual(ts['tags'], tags)
        # state should not have been updated
        self.assertEqual(ts['state'], state)
        self.assertEqual(ts['error'], error)
        self.assertEqual(ts['spawned_tasks'], spawned_tasks)
        self.assertEqual(ts['progress_report'], progress_report)
        self.assertEqual(ts['task_type'], task_type)
        # start_time should not have been updated
        self.assertEqual(ts['start_time'], start_time)
        self.assertEqual(ts['finish_time'], finish_time)
        self.assertEqual(ts['result'], result)
        # These are always None
        self.assertEqual(ts['traceback'], None)
        self.assertEqual(ts['exception'], None)

 def accepted(self, reply):
     """
     Notification that an RMI has started executing in the agent.
     The task status is updated in the pulp DB.
     :param reply: A status reply object.
     :type reply: gofer.rmi.async.Accepted
     """
     _logger.debug(_('Task RMI (accepted): %(r)s'), {'r': reply})
     call_context = dict(reply.data)
     task_id = call_context['task_id']
     TaskStatus.objects(task_id=task_id, state=constants.CALL_WAITING_STATE).\
         update_one(set__state=constants.CALL_ACCEPTED_STATE)

 def test_illegal_multi_arg(self):
     """
     Test that we receive an exception if we try to use the 'multi' kwarg
     """
     task_id = self.get_random_uuid()
     worker_name = 'special_worker_name'
     tags = ['test-tag1', 'test-tag2']
     state = 'waiting'
     ts = TaskStatus(task_id, worker_name, tags, state)
     ts.save()
     self.assertRaises(NotImplementedError, TaskStatus.objects(task_id=task_id).update_one,
                       multi=True)