-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathqueue.py
More file actions
134 lines (106 loc) · 4.08 KB
/
Copy pathqueue.py
File metadata and controls
134 lines (106 loc) · 4.08 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""Mongo backed queue model"""
import datetime
from mongokit import Document, ObjectId, OperationFailure
from pymongo.errors import InvalidId as PymongoInvalidId
from httpqueue.model.errors import InvalidId as PQInvalidId
# The prefix added to collections that belong to priority queues.
PRIORITY_QUEUE_PREFIX = 'pq_'
class TaskDoc(Document):
"""Document representing a task and its associated metadata."""
structure = {
'priority': datetime.datetime,
'task': None,
'in_progress': bool,
'expire_time': datetime.datetime,
'pending_life': int,
}
required_fields = ['priority', 'task']
default_values = {'in_progress': False, 'pending_life': 5 * 60}
indexes = [
{
'fields': ['priority', 'in_progress'],
},
{
'fields': ['expire_time'],
},
]
use_dot_notation = True
dot_notation_warning = True
class PriorityQueue(object):
"""A priority queue implementation using MongoDB
:param connection: A connection to MongoDB.
:param q_name: The name of the system wide queue to back this
object with.
"""
def __init__(self, connection, q_name):
self.name = PRIORITY_QUEUE_PREFIX + q_name
self.con = connection
self.con.register([TaskDoc])
@property
def db(self):
"FIXME: Should be passed in some way or other."
return self.con.test
@property
def collection(self):
return getattr(self.db, self.name)
def push(self, priority, obj):
pq = self.collection.TaskDoc()
pq.priority = priority
pq.task = obj
pq.save()
return pq['_id']
def pop(self):
try:
rv = self.db.command(
"findandmodify", self.name,
query={
"in_progress": False,
"priority": {"$lte" : datetime.datetime.utcnow() }},
sort={"priority" : 1},
update={"$set": {"in_progress" : True}},
)['value']
except OperationFailure:
return None
# Set the expiration date before returning the value.
expiration = self._calculate_expiration_time(rv['pending_life'])
self.collection.update({'_id': self._parse_object_id(rv['_id'])},
{'$set': {'expire_time': expiration}})
return rv
def ack(self, id):
"""Drop a task with the given id that has already been picked up."""
rv = self.collection.remove({
'_id': self._parse_object_id(id), 'in_progress': True}, safe=True)
if rv['n'] is 0:
raise KeyError
def cancel(self, id):
"""Drop a task with the given id given that it has not yet been popped.
"""
rv = self.collection.remove({
'_id': self._parse_object_id(id), 'in_progress': False}, safe=True)
if rv['n'] is 0:
raise KeyError
def restore_pending(self):
"""Restore all pending tasks whose pending-life has expired.
They will be returned to the state they were in prior to being
POPed. This should be run as frequently as is feasable so
that jobs can be picked back up by a worker quickly.
"""
self.collection.update(
{'in_progress': True,
'expire_time': {'$lte': datetime.datetime.utcnow()}},
{'$unset': {'expire_time': 1},
'$set': {'in_progress': False}})
def _parse_object_id(self, id):
"""Return the object id or raise an error."""
try:
return ObjectId(id)
except PymongoInvalidId as ex:
raise PQInvalidId(ex)
def _calculate_expiration_time(self, pending_life):
"""Calculate the expiration date based on the pending life."""
return datetime.datetime.utcnow() + \
datetime.timedelta(seconds=pending_life)
def list_queues(db):
"""Return a list of the queues available on :param db:."""
return [name[len(PRIORITY_QUEUE_PREFIX):] for name in db.collection_names()
if name.startswith(PRIORITY_QUEUE_PREFIX)]