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()