<reponame>bpmbank/mrq
from mrq.queue import send_task, Queue
from mrq.task import Task
from mrq.job import Job
from mrq.context import log, connections
import datetime
import time


class RequeueInterruptedJobs(Task):
  """ Requeue jobs that were marked as status=interrupt when a worker got a SIGTERM. """

  def run(self, params):
    return send_task("mrq.basetasks.utils.JobAction", {
      "status": "interrupt",
      "action": "requeue"
    }, sync=True)


class RequeueRetryJobs(Task):
  """ Requeue jobs that were marked as retry. """

  def run(self, params):
    return send_task("mrq.basetasks.utils.JobAction", {
      "status": "retry",
      "dateretry": {"$lte": datetime.datetime.utcnow()},
      "action": "requeue"
    }, sync=True)


class RequeueStartedJobs(Task):
  """ Requeue jobs that were marked as status=started and never finished.

      That may be because the worker got a SIGKILL or was terminated abruptly. The timeout parameter
      of this task is in addition to the task's own timeout.
  """

  def run(self, params):

    additional_timeout = params.get("timeout", 300)

    stats = {
      "requeued": 0,
      "started": 0
    }

    # There shouldn't be that much "started" jobs so we can quite safely iterate over them.
    self.collection = connections.mongodb_jobs.mrq_jobs
    for job_data in self.collection.find({"status": "started"}, fields={"_id": 1, "datestarted": 1, "queue": 1, "path": 1}):
      job = Job(job_data["_id"])
      job.set_data(job_data)

      stats["started"] += 1

      expire_date = datetime.datetime.utcnow() - datetime.timedelta(seconds=job.timeout + additional_timeout)

      if job_data["datestarted"] < expire_date:
        log.debug("Requeueing job %s" % job.id)
        job.requeue()
        stats["requeued"] += 1

    return stats


class RequeueRedisStartedJobs(Task):
  """ Requeue jobs that were started in Redis but not in Mongo.

      They could have been lost by a worker interrupt between
      redis.lpop and mongodb.update
  """

  def run(self, params):

    self.collection = connections.mongodb_jobs.mrq_jobs

    redis_key_started = Queue.redis_key_started()

    stats = {
      "fetched": 0,
      "requeued": 0
    }

    # Fetch all the jobs started more than a minute ago - they should not be in redis:started anymore
    job_ids = connections.redis.zrangebyscore(redis_key_started, "-inf", time.time() - params.get("timeout", 60))

    for job_id in job_ids:

      queue = Job(job_id, start=False, fetch=False).fetch(full_data=True).data["queue"]

      stats["fetched"] += 1

      log.info("Requeueing %s on %s" % (job_id, queue))

      # TODO LUA script & don't rpush if not in zset anymore.
      with connections.redis.pipeline(transaction=True) as pipeline:
        pipeline.zrem(redis_key_started, job_id)
        pipeline.rpush(Queue(queue).redis_key, job_id)
        pipeline.execute()

      stats["requeued"] += 1

    return stats


class RequeueLostJobs(Task):
  """ Requeue jobs that were queued but don't appear in Redis anymore.

      They could have been lost either by a Redis flush
  """

  def run(self, params):

    self.collection = connections.mongodb_jobs.mrq_jobs

    # If there are more than this much items on the queue, we don't try to check if our mongodb
    # jobs are still queued.
    max_queue_items = params.get("max_queue_items", 1000)

    stats = {
      "fetched": 0,
      "requeued": 0
    }

    for job_data in self.collection.find({
      "status": "queued"
    }, fields={"_id": 1, "queue": 1}).sort([("_id", 1)]):

      stats["fetched"] += 1

      queue = Queue(job_data["queue"])
      queue_size = queue.size()
      if queue_size > max_queue_items:
        log.info("Stopping because queue %s has %s items" % (queue, queue_size))
        break

      queue_jobs_ids = set(queue.list_job_ids(limit=max_queue_items + 1))
      if len(queue_jobs_ids) >= max_queue_items:
        log.info("Stopping because queue %s actually had more than %s items" % (queue, len(queue_jobs_ids)))
        break

      if str(job_data["_id"]) in queue_jobs_ids:
        log.info("Stopping because we found job %s in redis" % job_data["_id"])
        break

      # At this point, this job is not on the queue and we're sure the queue is less than max_queue_items
      # We can safely requeue the job.
      log.info("Requeueing %s on %s" % (job_data["_id"], queue.id))

      stats["requeued"] += 1
      job = Job(job_data["_id"])
      job.requeue(queue=job_data["queue"])

    return stats


# class CleanupStoppedWorkers(Task):
#   """ Remove stopped or timeouted workers from mongodb.mrq_workers. """

#   def run(self, params):


from mrq.queue import send_task, Queue
from mrq.task import Task
from mrq.job import Job
from mrq.context import log, connections
import datetime
import time


class RequeueInterruptedJobs(Task):
  """ Requeue jobs that were marked as status=interrupt when a worker got a SIGTERM. """

  def run(self, params):
    return send_task("mrq.basetasks.utils.JobAction", {
      "status": "interrupt",
      "action": "requeue"
    }, sync=True)


class RequeueRetryJobs(Task):
  """ Requeue jobs that were marked as retry. """

  def run(self, params):
    return send_task("mrq.basetasks.utils.JobAction", {
      "status": "retry",
      "dateretry": {"$lte": datetime.datetime.utcnow()},
      "action": "requeue"
    }, sync=True)


class RequeueStartedJobs(Task):
  """ Requeue jobs that were marked as status=started and never finished.

      That may be because the worker got a SIGKILL or was terminated abruptly. The timeout parameter
      of this task is in addition to the task's own timeout.
  """

  def run(self, params):

    additional_timeout = params.get("timeout", 300)

    stats = {
      "requeued": 0,
      "started": 0
    }

    # There shouldn't be that much "started" jobs so we can quite safely iterate over them.
    self.collection = connections.mongodb_jobs.mrq_jobs
    for job_data in self.collection.find({"status": "started"}, fields={"_id": 1, "datestarted": 1, "queue": 1, "path": 1}):
      job = Job(job_data["_id"])
      job.set_data(job_data)

      stats["started"] += 1

      expire_date = datetime.datetime.utcnow() - datetime.timedelta(seconds=job.timeout + additional_timeout)

      if job_data["datestarted"] < expire_date:
        log.debug("Requeueing job %s" % job.id)
        job.requeue()
        stats["requeued"] += 1

    return stats


class RequeueRedisStartedJobs(Task):
  """ Requeue jobs that were started in Redis but not in Mongo.

      They could have been lost by a worker interrupt between
      redis.lpop and mongodb.update
  """

  def run(self, params):

    self.collection = connections.mongodb_jobs.mrq_jobs

    redis_key_started = Queue.redis_key_started()

    stats = {
      "fetched": 0,
      "requeued": 0
    }

    # Fetch all the jobs started more than a minute ago - they should not be in redis:started anymore
    job_ids = connections.redis.zrangebyscore(redis_key_started, "-inf", time.time() - params.get("timeout", 60))

    for job_id in job_ids:

      queue = Job(job_id, start=False, fetch=False).fetch(full_data=True).data["queue"]

      stats["fetched"] += 1

      log.info("Requeueing %s on %s" % (job_id, queue))

      # TODO LUA script & don't rpush if not in zset anymore.
      with connections.redis.pipeline(transaction=True) as pipeline:
        pipeline.zrem(redis_key_started, job_id)
        pipeline.rpush(Queue(queue).redis_key, job_id)
        pipeline.execute()

      stats["requeued"] += 1

    return stats


class RequeueLostJobs(Task):
  """ Requeue jobs that were queued but don't appear in Redis anymore.

      They could have been lost either by a Redis flush
  """

  def run(self, params):

    self.collection = connections.mongodb_jobs.mrq_jobs

    # If there are more than this much items on the queue, we don't try to check if our mongodb
    # jobs are still queued.
    max_queue_items = params.get("max_queue_items", 1000)

    stats = {
      "fetched": 0,
      "requeued": 0
    }

    for job_data in self.collection.find({
      "status": "queued"
    }, fields={"_id": 1, "queue": 1}).sort([("_id", 1)]):

      stats["fetched"] += 1

      queue = Queue(job_data["queue"])
      queue_size = queue.size()
      if queue_size > max_queue_items:
        log.info("Stopping because queue %s has %s items" % (queue, queue_size))
        break

      queue_jobs_ids = set(queue.list_job_ids(limit=max_queue_items + 1))
      if len(queue_jobs_ids) >= max_queue_items:
        log.info("Stopping because queue %s actually had more than %s items" % (queue, len(queue_jobs_ids)))
        break

      if str(job_data["_id"]) in queue_jobs_ids:
        log.info("Stopping because we found job %s in redis" % job_data["_id"])
        break

      # At this point, this job is not on the queue and we're sure the queue is less than max_queue_items
      # We can safely requeue the job.
      log.info("Requeueing %s on %s" % (job_data["_id"], queue.id))

      stats["requeued"] += 1
      job = Job(job_data["_id"])
      job.requeue(queue=job_data["queue"])

    return stats


# class CleanupStoppedWorkers(Task):
#   """ Remove stopped or timeouted workers from mongodb.mrq_workers. """

#   def run(self, params):
Requeue jobs that were marked as status=interrupt when a worker got a SIGTERM. Requeue jobs that were marked as retry. Requeue jobs that were marked as status=started and never finished.        That may be because the worker got a SIGKILL or was terminated abruptly. The timeout parameter       of this task is in addition to the task's own timeout. # There shouldn't be that much "started" jobs so we can quite safely iterate over them. Requeue jobs that were started in Redis but not in Mongo.        They could have been lost by a worker interrupt between       redis.lpop and mongodb.update # Fetch all the jobs started more than a minute ago - they should not be in redis:started anymore # TODO LUA script & don't rpush if not in zset anymore. Requeue jobs that were queued but don't appear in Redis anymore.        They could have been lost either by a Redis flush # If there are more than this much items on the queue, we don't try to check if our mongodb # jobs are still queued. # At this point, this job is not on the queue and we're sure the queue is less than max_queue_items # We can safely requeue the job. # class CleanupStoppedWorkers(Task): #   """ Remove stopped or timeouted workers from mongodb.mrq_workers. """ #   def run(self, params):