schedule: Allow arguments for scheduled functions, lock before modifying heap

Signed-off-by: James McCoy <jamessan@users.sourceforge.net>
This commit is contained in:
James McCoy 2012-09-08 16:02:30 -04:00
parent 6f925e5f7a
commit 45329fbdce

View File

@ -34,6 +34,7 @@ Supybot driver.
import time import time
import heapq import heapq
from threading import Lock
import supybot.log as log import supybot.log as log
import supybot.world as world import supybot.world as world
@ -61,10 +62,12 @@ class Schedule(drivers.IrcDriver):
self.schedule = [] self.schedule = []
self.events = {} self.events = {}
self.counter = 0 self.counter = 0
self.lock = Lock()
def reset(self): def reset(self):
self.events.clear() with self.lock:
self.schedule[:] = [] self.events.clear()
self.schedule[:] = []
# We don't reset the counter here because if someone has held an id of # We don't reset the counter here because if someone has held an id of
# one of the nuked events, we don't want him removing new events with # one of the nuked events, we don't want him removing new events with
# his old id. # his old id.
@ -72,7 +75,7 @@ class Schedule(drivers.IrcDriver):
def name(self): def name(self):
return 'Schedule' return 'Schedule'
def addEvent(self, f, t, name=None): def addEvent(self, f, t, name=None, args=[], kwargs={}):
"""Schedules an event f to run at time t. """Schedules an event f to run at time t.
name must be hashable and not an int. name must be hashable and not an int.
@ -82,20 +85,22 @@ class Schedule(drivers.IrcDriver):
self.counter += 1 self.counter += 1
assert name not in self.events, \ assert name not in self.events, \
'An event with the same name has already been scheduled.' 'An event with the same name has already been scheduled.'
self.events[name] = f with self.lock:
heapq.heappush(self.schedule, mytuple((t, name))) self.events[name] = f
heapq.heappush(self.schedule, mytuple((t, name, args, kwargs)))
return name return name
def removeEvent(self, name): def removeEvent(self, name):
"""Removes the event with the given name from the schedule.""" """Removes the event with the given name from the schedule."""
f = self.events.pop(name) f = self.events.pop(name)
self.schedule = [(t, n) for (t, n) in self.schedule if n != name]
# We must heapify here because the heap property may not be preserved # We must heapify here because the heap property may not be preserved
# by the above list comprehension. We could, conceivably, just mark # by the above list comprehension. We could, conceivably, just mark
# the elements of the heap as removed and ignore them when we heappop, # the elements of the heap as removed and ignore them when we heappop,
# but that would only save a constant factor (we're already linear for # but that would only save a constant factor (we're already linear for
# the listcomp) so I'm not worried about it right now. # the listcomp) so I'm not worried about it right now.
heapq.heapify(self.schedule) with self.lock:
self.schedule = [x for x in self.schedule if x[1] != name]
heapq.heapify(self.schedule)
return f return f
def rescheduleEvent(self, name, t): def rescheduleEvent(self, name, t):
@ -123,11 +128,12 @@ class Schedule(drivers.IrcDriver):
'why do we continue to live?') 'why do we continue to live?')
time.sleep(1) # We're the only driver; let's pause to think. time.sleep(1) # We're the only driver; let's pause to think.
while self.schedule and self.schedule[0][0] < time.time(): while self.schedule and self.schedule[0][0] < time.time():
(t, name) = heapq.heappop(self.schedule) with self.lock:
f = self.events[name] (t, name, args, kwargs) = heapq.heappop(self.schedule)
f = self.events[name]
del self.events[name] del self.events[name]
try: try:
f() f(*args, **kwargs)
except Exception, e: except Exception, e:
log.exception('Uncaught exception in scheduled function:') log.exception('Uncaught exception in scheduled function:')