Allow using scheduler as a context manager

rq_scheduler/scheduler.py

@@ -28,6 +28,18 @@ def __init__(self, queue_name='default', interval=60, connection=None):
         self.log = logger
         self._lock_acquired = False
 
+    def __enter__(self):
+        while True:
+            try:
+                self.register_birth()
+                return self
+            except ValueError:  # assume register_birth failed
+                self.log.info('Waiting for registering birth...')
+                time.sleep(self._interval)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.register_death()
+
     def register_birth(self):
         if self.connection.exists(self.scheduler_key) and \
                 not self.connection.hexists(self.scheduler_key, 'death'):
@@ -364,7 +376,6 @@ def run(self, burst=False):
         """
         self.log.info('Running RQ scheduler...')
 
-        self.register_birth()
         self._install_signal_handlers()
 
         try:
@@ -384,4 +395,3 @@ def run(self, burst=False):
                 time.sleep(self._interval - (time.time() - start_time))
         finally:
             self.remove_lock()
-            self.register_death()

rq_scheduler/scripts/rqscheduler.py

@@ -49,8 +49,8 @@ def main():
         level = 'INFO'
     setup_loghandlers(level)
 
-    scheduler = Scheduler(connection=connection, interval=args.interval)
-    scheduler.run(burst=args.burst)
+    with Scheduler(connection=connection, interval=args.interval) as scheduler:
+        scheduler.run(burst=args.burst)
 
 if __name__ == '__main__':
     main()