added the first parts of delayed tasks

Author: Matt George

pyres/__init__.py

@@ -3,7 +3,8 @@
 from redis import Redis
 import pyres.json_parser as json
 
-import types
+#import types
+import time, datetime
 
 def my_import(name):
     mod = __import__(name)    
@@ -179,6 +180,45 @@ def close(self):
         """
         self.redis.disconnect()
 
+    def enqueue_at(self, timestamp, klass, *args):
+        class_name = '%s.%s' % (klass.__module__, klass.__name__)
+        self.delayed_push(timestamp, {'class':class_name,'queue': klass.queue, 'args':args})
+    
+    def delayed_push(self, timestamp, item):
+        key = int(time.mktime(timestamp.timetuple()))
+        self.redis.push('resque:delayed:%s' % key, ResQ.encode(item))
+        self.redis.zadd('resque:delayed_queue_schedule', key, key)
+    
+    def delayed_queue_peek(self, start, count):
+        return [int(item) for item in self.redis.zrange('resque:delayed_queue_schedule', start, start+count)]
+    
+    def delayed_queue_schedule_size(self):
+        return self.redis.zcard('resque:delayed_queue_schedule')
+    
+    def delayed_timestamp_size(self, timestamp):
+        key = int(time.mktime(timestamp.timetuple()))
+        return self.redis.llen("resque:delayed:%i" % key)
+    
+    def next_delayed_timestamp(self):
+        key = int(time.mktime(datetime.datetime.now().timetuple()))
+        array = self.redis.zrangebyscore('resque:delayed_queue_schedule', '-inf', key)
+        timestamp = None
+        if array:
+            timestamp = array[0]
+        return timestamp
+    
+    def next_item_for_timestamp(self, timestamp):
+        #key = int(time.mktime(timestamp.timetuple()))
+        key = "resque:delayed:%s" % timestamp
+        ret = self.redis.pop(key)
+        item = None
+        if ret:
+            item = ResQ.decode(ret)
+        if self.redis.llen(key) == 0:
+            self.redis.delete(key)
+            self.redis.zrem('resque:delayed_queue_schedule', timestamp)
+        return item
+    
     @classmethod
     def encode(cls, item):
         return json.dumps(item)

pyres/scheduler.py

@@ -0,0 +1,61 @@
+import signal
+import time
+
+from pyres import ResQ
+
+class Scheduler(object):
+    
+    def __init__(self, server="localhost:6379", password=None):
+        self._shutdown = False
+        if isinstance(server,basestring):
+            self.resq = ResQ(server=server, password=password)
+        elif isinstance(server, ResQ):
+            self.resq = server
+        else:
+            raise Exception("Bad server argument")
+    
+    def register_signal_handlers(self):
+        signal.signal(signal.SIGTERM, self.schedule_shutdown)
+        signal.signal(signal.SIGINT, self.schedule_shutdown)
+        signal.signal(signal.SIGQUIT, self.schedule_shutdown)
+    
+    def schedule_shutdown(self):
+        self._shutdown = True
+    
+    def run(self):
+        self.register_signal_handlers()
+        #self.load_schedule()
+        while True:
+            if self._shutdown is True:
+                break
+            self.handle_delayed_items()
+            time.sleep(5)
+    
+    def next_timestamp(self):
+        while True
+            timestamp = self.resq.next_delayed_timestamp()
+            if timestamp:
+                yield timestamp
+            else:
+                break
+            
+    
+    def next_item(self, timestamp):
+        while True:
+            item = self.resq.next_item_for_timestamp(timestamp)
+            if item:
+                yield item
+            else:
+                break
+    
+    def handle_delayed_items(self):
+        for timestamp in self.next_timestamp():
+            for item in self.next_item(timestamp):
+                print 'queueing item'
+                klass = item['class']
+                queue = item['queue']
+                args = item['args']
+                self.resq.enqueue_from_string(klass, queue, args)
+            
+        
+    

tests/test_schedule.py

@@ -0,0 +1,62 @@
+from tests import PyResTests, Basic, TestProcess, ErrorObject
+from pyres import ResQ
+from pyres.job import Job
+from pyres.scheduler import Scheduler
+import os
+import datetime
+import time
+class ScheduleTests(PyResTests):
+    def test_enqueue_at(self):
+        d = datetime.datetime.now() + datetime.timedelta(days=1)
+        d2 = d + datetime.timedelta(days=1)
+        key = int(time.mktime(d.timetuple()))
+        key2 = int(time.mktime(d2.timetuple()))
+        self.resq.enqueue_at(d, Basic,"test1")
+        self.resq.enqueue_at(d, Basic,"test2")
+        assert self.redis.llen("resque:delayed:%s" % key) == 2
+        assert len(self.redis.zrange('resque:delayed_queue_schedule',0,20)) == 1
+        self.resq.enqueue_at(d2, Basic,"test1")
+        assert self.redis.llen("resque:delayed:%s" % key2) == 1
+        assert len(self.redis.zrange('resque:delayed_queue_schedule',0,20)) == 2
+    
+    def test_delayed_queue_schedule_size(self):
+        d = datetime.datetime.now() + datetime.timedelta(days=1)
+        d2 = d + datetime.timedelta(days=1)
+        key = int(time.mktime(d.timetuple()))
+        key2 = int(time.mktime(d2.timetuple()))
+        self.resq.enqueue_at(d, Basic,"test1")
+        self.resq.enqueue_at(d2, Basic,"test1")
+        assert self.resq.delayed_queue_schedule_size() == 2
+    
+    def test_delayed_timestamp_size(self):
+        d = datetime.datetime.now() + datetime.timedelta(days=1)
+        d2 = d + datetime.timedelta(days=1)
+        key = int(time.mktime(d.timetuple()))
+        key2 = int(time.mktime(d2.timetuple()))
+        self.resq.enqueue_at(d, Basic,"test1")
+        assert self.resq.delayed_timestamp_size(d) == 1
+        self.resq.enqueue_at(d, Basic,"test1")
+        assert self.resq.delayed_timestamp_size(d) == 2
+    
+    def test_next_delayed_timestamp(self):
+        d = datetime.datetime.now() + datetime.timedelta(days=-1)
+        d2 = d + datetime.timedelta(days=-2)
+        key = int(time.mktime(d.timetuple()))
+        key2 = int(time.mktime(d2.timetuple()))
+        self.resq.enqueue_at(d, Basic,"test1")
+        self.resq.enqueue_at(d2, Basic,"test1")
+        item = self.resq.next_delayed_timestamp()
+        assert  item == key2
+    
+    def test_next_item_for_timestamp(self):
+        d = datetime.datetime.now() + datetime.timedelta(days=-1)
+        d2 = d + datetime.timedelta(days=-2)
+        #key = int(time.mktime(d.timetuple()))
+        #key2 = int(time.mktime(d2.timetuple()))
+        self.resq.enqueue_at(d, Basic,"test1")
+        self.resq.enqueue_at(d2, Basic,"test1")
+        timestamp = self.resq.next_delayed_timestamp()
+        item = self.resq.next_item_for_timestamp(timestamp)
+        assert isinstance(item, dict)
+        assert self.redis.zcard('resque:delayed_queue_schedule') == 1
+