def test_queue(self):
q = PersistentPriorityQueue()
# Medium priority tasks
obj1 = {'test':'medium_priority 1'}
q.put(obj1, priority=10)
obj2 = {'test':'medium_priority 2'}
q.put(obj2, priority=10)
# High priority tasks
obj3 = {'test':'medium_priority 3'}
q.put(obj3, priority=1)
obj4 = {'test':'medium_priority 4'}
q.put(obj4, priority=1)
order = [obj3, obj4, obj1, obj2]
for o in order:
self.assertEqual(o, q.get())
def __init__(self, filename=None):
self.filename = filename
available_queue_state = None
wait_queue_state = None
if self.filename is not None:
try:
with open(filename, 'rb') as input_file:
(available_queue_state, wait_queue_state) = pickle.load(input_file)
except IOError:
# Not an error, just first run where we don't have a queue.
pass
except EOFError:
# The file we're trying to read from is probably broken
print 'Persistence error'
self.available_queue = PersistentPriorityQueue(state=available_queue_state)
# queue for tasks that have a delay
self.wait_queue = PersistentPriorityQueue(state=wait_queue_state)
class DelayPersistentPriorityQueue(object):
def __init__(self, filename=None):
self.filename = filename
available_queue_state = None
wait_queue_state = None
if self.filename is not None:
try:
with open(filename, 'rb') as input_file:
(available_queue_state, wait_queue_state) = pickle.load(input_file)
except IOError:
# Not an error, just first run where we don't have a queue.
pass
except EOFError:
# The file we're trying to read from is probably broken
print 'Persistence error'
self.available_queue = PersistentPriorityQueue(state=available_queue_state)
# queue for tasks that have a delay
self.wait_queue = PersistentPriorityQueue(state=wait_queue_state)
def put(self, message, priority=0, delay=0):
if delay == 0:
# No delay, message is immediately available
self.put_in_avail(message, priority)
else:
self.put_in_wait(message, delay, priority)
# Update available queue to avoid backlog upon get.
self.move_newly_available_messages()
def put_in_wait(self, message, delay, priority=0):
data = (priority, message)
available_at = datetime.datetime.now() + datetime.timedelta(seconds=delay)
self.wait_queue.put(data, priority=available_at)
def put_in_avail(self, message, priority):
self.available_queue.put(message, priority=priority)
def get(self):
self.move_newly_available_messages()
return self.available_queue.get()
def move_newly_available_messages(self):
'''
Moves messages that should be available from the
wait queue to the available queue
Called upon both gets and puts to avoid backlogs.
'''
now = datetime.datetime.now()
while (self.wait_queue.next_priority() is not None
and self.wait_queue.next_priority() < now):
(priority, message) = self.wait_queue.get()
self.put_in_avail(message, priority)
def get_state(self):
'''
State of this queue is simply the state of the two queues
that it's comprised of
'''
return (self.available_queue.get_state(), self.wait_queue.get_state())
def persist(self):
if self.filename is not None:
with open(self.filename, 'wb') as output_file:
pickle.dump(self.get_state(), output_file)