From 91cb7a47a129d8714e3e937230fd36ac897a6a04 Mon Sep 17 00:00:00 2001
From: Alexander Zhukov <azhukov@atlassian.com>
Date: Mon, 24 Apr 2017 18:58:31 +0300
Subject: [PATCH] Allow using scheduler as a context manager

---
 rq_scheduler/scheduler.py           | 14 ++++++++++++--
 rq_scheduler/scripts/rqscheduler.py |  4 ++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/rq_scheduler/scheduler.py b/rq_scheduler/scheduler.py
index 78bab4e..509b6f0 100644
--- a/rq_scheduler/scheduler.py
+++ b/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()
diff --git a/rq_scheduler/scripts/rqscheduler.py b/rq_scheduler/scripts/rqscheduler.py
index 0d71ffd..ad6b425 100755
--- a/rq_scheduler/scripts/rqscheduler.py
+++ b/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()