common.lib.queue

  1"""
  2A job queue, to divide work over the workers
  3"""
  4import time
  5import json
  6
  7from common.lib.job import Job
  8import psycopg2
  9
 10
 11class JobQueue:
 12	"""
 13	A simple job queue
 14
 15	Jobs are basically database records. The job has some information that a worker
 16	can use to do its job. The job queue is shared between workers so that nothing
 17	is done twice.
 18	"""
 19	db = None
 20	log = None
 21
 22	def __init__(self, logger, database):
 23		"""
 24		Set up database handler
 25		"""
 26		self.log = logger
 27
 28		self.db = database
 29
 30	def get_job(self, jobtype, timestamp=-1, restrict_claimable=True):
 31		"""
 32		Get job of a specific type
 33
 34		Returns a job's data. The `details` column is parsed JSON, and can thus contain all
 35		kinds of data.
 36
 37		:param string jobtype:  Job type
 38		:param int timestamp:  Find jobs that may be claimed after this timestamp. If set to
 39							   a negative value (default), any job with a "claim after" time
 40							   earlier than the current time is selected.
 41		:param bool restrict_claimable:  Only return jobs that may be claimed
 42		according to their parameters
 43		:return dict: Job data, or `None` if no job was found
 44		"""
 45		if timestamp < 0:
 46			timestamp = int(time.time())
 47
 48		# select the number of jobs of the same type that have been queued for
 49		# longer than the job as well
 50		replacements = [jobtype]
 51		query = (
 52			"SELECT main_queue.* FROM jobs AS main_queue"
 53			"        WHERE main_queue.jobtype = %s"
 54		)
 55
 56		if restrict_claimable:
 57			# claimability is determined through various timestamps
 58			query += (
 59			"          AND main_queue.timestamp_claimed = 0"
 60			"          AND main_queue.timestamp_after < %s"
 61			"          AND (main_queue.interval = 0 OR main_queue.timestamp_lastclaimed + main_queue.interval < %s)"
 62			)
 63			replacements.append(timestamp)
 64			replacements.append(timestamp)
 65
 66		query += "    ORDER BY main_queue.timestamp ASC LIMIT 1;"
 67
 68		job = self.db.fetchone(query, tuple(replacements))
 69
 70		return Job.get_by_data(job, database=self.db) if job else None
 71
 72	def get_all_jobs(self, jobtype="*", remote_id=False, restrict_claimable=True):
 73		"""
 74		Get all unclaimed (and claimable) jobs
 75
 76		:param string jobtype:  Type of job, "*" for all types
 77		:param string remote_id:  Remote ID, takes precedence over `jobtype`
 78		:param bool restrict_claimable:  Only return jobs that may be claimed
 79		according to their parameters
 80		:return list:
 81		"""
 82		replacements = []
 83		if remote_id:
 84			filter = "WHERE remote_id = %s"
 85			replacements = [remote_id]
 86		elif jobtype != "*":
 87			filter = "WHERE jobtype = %s"
 88			replacements = [jobtype]
 89		else:
 90			filter = "WHERE jobtype != ''"
 91
 92		query = "SELECT * FROM jobs %s" % filter
 93
 94		if restrict_claimable:
 95			query += ("        AND timestamp_claimed = 0"
 96					  "              AND timestamp_after < %s"
 97					  "              AND (interval = 0 OR timestamp_lastclaimed + interval < %s)")
 98
 99			now = int(time.time())
100			replacements.append(now)
101			replacements.append(now)
102
103		query += "         ORDER BY timestamp ASC"
104
105		jobs = self.db.fetchall(query, replacements)
106
107		return [Job.get_by_data(job, self.db) for job in jobs if job]
108
109	def get_job_count(self, jobtype="*"):
110		"""
111		Get total number of jobs
112
113		:param jobtype:  Type of jobs to count. Default (`*`) counts all jobs.
114		:return int:  Number of jobs
115		"""
116		if jobtype == "*":
117			count = self.db.fetchone("SELECT COUNT(*) FROM jobs;", ())
118		else:
119			count = self.db.fetchone("SELECT COUNT(*) FROM jobs WHERE jobtype = %s;", (jobtype,))
120
121		return int(count["count"])
122
123	def add_job(self, jobtype, details=None, remote_id=0, claim_after=0, interval=0):
124		"""
125		Add a new job to the queue
126
127		There can only be one job for any combination of job type and remote id. If a job
128		already exists for the given combination, no new job is added.
129
130		:param jobtype:  Job type
131		:param details:  Job details - may be empty, will be stored as JSON
132		:param remote_id:  Remote ID of object to work on. For example, a post or thread ID
133		:param claim_after:  Absolute timestamp after which job may be claimed
134		:param interval:  If this is not zero, the job is made a repeating job,
135		                  which will be repeated at most every `interval` seconds.
136
137		:return Job: A job that matches the input type and remote ID. This may
138		             be a newly added job or an existing that matched the same
139		             combination (which is required to be unique, so no new job
140		             with those parameters could be queued, and the old one is
141		             just as valid).
142		"""
143		data = {
144			"jobtype": jobtype,
145			"details": json.dumps(details),
146			"timestamp": int(time.time()),
147			"timestamp_claimed": 0,
148			"timestamp_lastclaimed": 0,
149			"remote_id": remote_id,
150			"timestamp_after": claim_after,
151			"interval": interval,
152			"attempts": 0
153		}
154
155		self.db.insert("jobs", data, safe=True, constraints=("jobtype", "remote_id"))
156
157		return Job.get_by_data(data, database=self.db)
158
159	def release_all(self):
160		"""
161		Release all jobs
162
163		All claimed jobs are released. This is useful to run when the backend is restarted.
164		"""
165		self.db.execute("UPDATE jobs SET timestamp_claimed = 0")
166
167	def get_place_in_queue(self, job):
168		"""
169		What is the place of this job in the queue?
170
171		:param Job job:  Job to get place in queue for
172
173		:return int: Place in queue. 0 means the job is currently being
174		processed; 1+ means the job is queued, with 1 corresponding to the
175		front of the queue.
176		"""
177		if job.data["timestamp_claimed"] > 0:
178			return 0
179
180		all_queued = self.get_all_jobs(jobtype=job.data["jobtype"])
181		our_timestamp = job.data["timestamp"]
182		return len(
183			[queued_job for queued_job in all_queued if queued_job.data["timestamp"] < our_timestamp or queued_job.data["timestamp_claimed"] > 0])