<commit_before>"""celery.managers"""
from django.db import models
from django.db import connection, transaction
from celery.registry import tasks
from celery.conf import TASK_RESULT_EXPIRES
from datetime import datetime, timedelta
from django.conf import settings
import random

# server_drift can be negative, but timedelta supports addition on
# negative seconds.
SERVER_DRIFT = timedelta(seconds=random.vonmisesvariate(1, 4))


class TableLock(object):
    """Base class for database table locks. Also works as a NOOP lock."""

    def __init__(self, table, type="read"):
        self.table = table
        self.type = type
        self.cursor = None

    def lock_table(self):
        """Lock the table."""
        pass

    def unlock_table(self):
        """Release previously locked tables."""
        pass

    @classmethod
    def acquire(cls, table, type=None):
        """Acquire table lock."""
        lock = cls(table, type)
        lock.lock_table()
        return lock

    def release(self):
        """Release the lock."""
        self.unlock_table()
        if self.cursor:
            self.cursor.close()
            self.cursor = None


class MySQLTableLock(TableLock):
    """Table lock support for MySQL."""

    def lock_table(self):
        """Lock MySQL table."""
        self.cursor = connection.cursor()
        self.cursor.execute("LOCK TABLES %s %s" % (
            self.table, self.type.upper()))

    def unlock_table(self):
        """Unlock MySQL table."""
        self.cursor.execute("UNLOCK TABLES")

TABLE_LOCK_FOR_ENGINE = {"mysql": MySQLTableLock}
table_lock = TABLE_LOCK_FOR_ENGINE.get(settings.DATABASE_ENGINE, TableLock)


class TaskManager(models.Manager):
    """Manager for :class:`celery.models.Task` models."""

    def get_task(self, task_id):
        """Get task meta for task by ``task_id``."""
        task, created = self.get_or_create(task_id=task_id)
        return task

    def is_done(self, task_id):
        """Returns ``True`` if the task was executed successfully."""
        return self.get_task(task_id).status == "DONE"

    def get_all_expired(self):
        """Get all expired task results."""
        return self.filter(date_done__lt=datetime.now() - TASK_RESULT_EXPIRES)

    def delete_expired(self):
        """Delete all expired task results."""
        self.get_all_expired().delete()

    def store_result(self, task_id, result, status, traceback=None,
            exception_retry=True):
        """Store the result and status of a task.

        :param task_id: task id

        :param result: The return value of the task, or an exception
            instance raised by the task.

        :param status: Task status. See
            :meth:`celery.result.AsyncResult.get_status` for a list of
            possible status values.

        :keyword traceback: The traceback at the point of exception (if the
            task failed).

        :keyword exception_retry: If True, we try a single retry with
            transaction rollback on exception
        """
        try:
            task, created = self.get_or_create(task_id=task_id, defaults={
                                                "status": status,
                                                "result": result,
                                                "traceback": traceback})
            if not created:
                task.status = status
                task.result = result
                task.traceback = traceback
                task.save()
        except Exception, exc:
            # depending on the database backend we can get various exceptions.
            # for excample, psycopg2 raises an exception if some operation
            # breaks transaction, and saving task result won't be possible
            # until we rollback transaction
            if exception_retry:
                transaction.rollback_unless_managed()
                self.store_result(task_id, result, status, traceback, False)
            else:
                raise


class PeriodicTaskManager(models.Manager):
    """Manager for :class:`celery.models.PeriodicTask` models."""

    def init_entries(self):
        """Add entries for all registered periodic tasks.

        Should be run at worker start.
        """
        periodic_tasks = tasks.get_all_periodic()
        for task_name in periodic_tasks.keys():
            task_meta, created = self.get_or_create(name=task_name)

    def is_time(self, last_run_at, run_every):
        """Check if if it is time to run the periodic task.

        :param last_run_at: Last time the periodic task was run.
        :param run_every: How often to run the periodic task.

        :rtype bool:

        """
        run_every_drifted = run_every + SERVER_DRIFT
        run_at = last_run_at + run_every_drifted
        if datetime.now() > run_at:
            return True
        return False

    def get_waiting_tasks(self):
        """Get all waiting periodic tasks.

        :returns: list of :class:`celery.models.PeriodicTaskMeta` objects.
        """
        periodic_tasks = tasks.get_all_periodic()
        db_table = self.model._meta.db_table

        # Find all periodic tasks to be run.
        waiting = []
        for task_meta in self.all():
            if task_meta.name in periodic_tasks:
                task = periodic_tasks[task_meta.name]
                run_every = task.run_every
                if self.is_time(task_meta.last_run_at, run_every):
                    # Get the object again to be sure noone else
                    # has already taken care of it.
                    lock = table_lock.acquire(db_table, "write")
                    try:
                        secure = self.get(pk=task_meta.pk)
                        if self.is_time(secure.last_run_at, run_every):
                            secure.last_run_at = datetime.now()
                            secure.save()
                            waiting.append(secure)
                    finally:
                        lock.release()
        return waiting
<commit_msg>Fixed race condition in storing task task in DB<commit_after>"""celery.managers"""
from django.db import models
from django.db import connection, transaction
from celery.registry import tasks
from celery.conf import TASK_RESULT_EXPIRES
from datetime import datetime, timedelta
from django.conf import settings
import random

# server_drift can be negative, but timedelta supports addition on
# negative seconds.
SERVER_DRIFT = timedelta(seconds=random.vonmisesvariate(1, 4))


class TableLock(object):
    """Base class for database table locks. Also works as a NOOP lock."""

    def __init__(self, table, type="read"):
        self.table = table
        self.type = type
        self.cursor = None

    def lock_table(self):
        """Lock the table."""
        pass

    def unlock_table(self):
        """Release previously locked tables."""
        pass

    @classmethod
    def acquire(cls, table, type=None):
        """Acquire table lock."""
        lock = cls(table, type)
        lock.lock_table()
        return lock

    def release(self):
        """Release the lock."""
        self.unlock_table()
        if self.cursor:
            self.cursor.close()
            self.cursor = None


class MySQLTableLock(TableLock):
    """Table lock support for MySQL."""

    def lock_table(self):
        """Lock MySQL table."""
        self.cursor = connection.cursor()
        self.cursor.execute("LOCK TABLES %s %s" % (
            self.table, self.type.upper()))

    def unlock_table(self):
        """Unlock MySQL table."""
        self.cursor.execute("UNLOCK TABLES")

TABLE_LOCK_FOR_ENGINE = {"mysql": MySQLTableLock}
table_lock = TABLE_LOCK_FOR_ENGINE.get(settings.DATABASE_ENGINE, TableLock)


class TaskManager(models.Manager):
    """Manager for :class:`celery.models.Task` models."""

    def get_task(self, task_id, exception_retry_count=1):
        """Get task meta for task by ``task_id``.
        
        :keyword exception_retry_count: How many times to retry with
            transaction rollback on exception. 1 by default: we assume
            the pessimistic case when we get race condition in which
            task is created by other process during get_or_create
        """
        try:
            task, created = self.get_or_create(task_id=task_id)
        except Exception, exc:
            # depending on the database backend we can get various exceptions,
            # so we catch every exception type
            if exception_retry_count > 0:
                transaction.rollback_unless_managed()
                return self.get_task(task_id, exception_retry_count-1)
            else:
                raise
        return task

    def is_done(self, task_id):
        """Returns ``True`` if the task was executed successfully."""
        return self.get_task(task_id).status == "DONE"

    def get_all_expired(self):
        """Get all expired task results."""
        return self.filter(date_done__lt=datetime.now() - TASK_RESULT_EXPIRES)

    def delete_expired(self):
        """Delete all expired task results."""
        self.get_all_expired().delete()

    def store_result(self, task_id, result, status, traceback=None,
            exception_retry_count=2):
        """Store the result and status of a task.

        :param task_id: task id

        :param result: The return value of the task, or an exception
            instance raised by the task.

        :param status: Task status. See
            :meth:`celery.result.AsyncResult.get_status` for a list of
            possible status values.

        :keyword traceback: The traceback at the point of exception (if the
            task failed).

        :keyword exception_retry_count: How many times to retry with
            transaction rollback on exception. 2 by default: we assume
            the pessimistic case when task execution by itself could
            leave broken transaction, and during second try we get
            race condition in which task is created by other process
            during get_or_create
        """
        try:
            task, created = self.get_or_create(task_id=task_id, defaults={
                                                "status": status,
                                                "result": result,
                                                "traceback": traceback})
            if not created:
                task.status = status
                task.result = result
                task.traceback = traceback
                task.save()
        except Exception, exc:
            # depending on the database backend we can get various exceptions.
            # for excample, psycopg2 raises an exception if some operation
            # breaks transaction, and saving task result won't be possible
            # until we rollback transaction
            if exception_retry_count > 0:
                transaction.rollback_unless_managed()
                self.store_result(task_id, result, status, traceback, exception_retry_count-1)
            else:
                raise


class PeriodicTaskManager(models.Manager):
    """Manager for :class:`celery.models.PeriodicTask` models."""

    def init_entries(self):
        """Add entries for all registered periodic tasks.

        Should be run at worker start.
        """
        periodic_tasks = tasks.get_all_periodic()
        for task_name in periodic_tasks.keys():
            task_meta, created = self.get_or_create(name=task_name)

    def is_time(self, last_run_at, run_every):
        """Check if if it is time to run the periodic task.

        :param last_run_at: Last time the periodic task was run.
        :param run_every: How often to run the periodic task.

        :rtype bool:

        """
        run_every_drifted = run_every + SERVER_DRIFT
        run_at = last_run_at + run_every_drifted
        if datetime.now() > run_at:
            return True
        return False

    def get_waiting_tasks(self):
        """Get all waiting periodic tasks.

        :returns: list of :class:`celery.models.PeriodicTaskMeta` objects.
        """
        periodic_tasks = tasks.get_all_periodic()
        db_table = self.model._meta.db_table

        # Find all periodic tasks to be run.
        waiting = []
        for task_meta in self.all():
            if task_meta.name in periodic_tasks:
                task = periodic_tasks[task_meta.name]
                run_every = task.run_every
                if self.is_time(task_meta.last_run_at, run_every):
                    # Get the object again to be sure noone else
                    # has already taken care of it.
                    lock = table_lock.acquire(db_table, "write")
                    try:
                        secure = self.get(pk=task_meta.pk)
                        if self.is_time(secure.last_run_at, run_every):
                            secure.last_run_at = datetime.now()
                            secure.save()
                            waiting.append(secure)
                    finally:
                        lock.release()
        return waiting

{'fraction_non_alphanumeric': 0.08371419865935405, 'fraction_numerical': 0.0010664229128580134, 'mean_word_length': 1.978448275862069, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '762543', 'n_tokens_mistral': 3488, 'n_tokens_neox': 3355, 'n_words': 1069}