Merge pull request #267 from Niccolum/main

find new master every time

Author: sibson

redbeat/schedulers.py

@@ -19,6 +19,7 @@
 from kombu.utils.objects import cached_property
 from kombu.utils.url import maybe_sanitize_url
 from redis.client import StrictRedis
+from redis.sentinel import MasterNotFoundError, Sentinel
 from tenacity import retry, retry_if_exception_type, stop_after_delay, wait_exponential
 
 from .decoder import RedBeatJSONDecoder, RedBeatJSONEncoder, to_timestamp
@@ -51,6 +52,9 @@
     return 1
 """
 
+REDBEAT_REDIS_KEY = "redbeat_redis"
+REDBEAT_SENTINEL_KEY = "redbeat_sentinel"
+
 
 class RetryingConnection:
     """A proxy for the Redis connection that delegates all the calls to
@@ -65,6 +69,7 @@ def __init__(self, retry_period, wrapped_connection):
             retry=(
                 retry_if_exception_type(redis.exceptions.ConnectionError)
                 | retry_if_exception_type(redis.exceptions.TimeoutError)
+                | retry_if_exception_type(MasterNotFoundError)
             ),
             reraise=True,
             wait=wait_exponential(multiplier=1, max=self.RETRY_MAX_WAIT),
@@ -114,9 +119,10 @@ def ensure_conf(app):
 def get_redis(app=None):
     app = app_or_default(app)
     conf = ensure_conf(app)
-    if not hasattr(app, 'redbeat_redis') or app.redbeat_redis is None:
-        redis_options = conf.redbeat_redis_options
-        retry_period = redis_options.get('retry_period')
+    redis_options = conf.redbeat_redis_options
+    retry_period = redis_options.get('retry_period')
+
+    if not hasattr(app, REDBEAT_REDIS_KEY) or getattr(app, REDBEAT_REDIS_KEY) is None:
         if redis_options.get('cluster', False):
             from redis.cluster import RedisCluster
 
@@ -137,9 +143,8 @@ def get_redis(app=None):
                 sentinel_kwargs=redis_options.get('sentinel_kwargs'),
                 **connection_kwargs,
             )
-            connection = sentinel.master_for(
-                redis_options.get('service_name', 'master'), db=redis_options.get('db', 0)
-            )
+            _set_redbeat_connect(app, REDBEAT_SENTINEL_KEY, sentinel, retry_period)
+            connection = None
         elif conf.redis_url.startswith('rediss'):
             ssl_options = {'ssl_cert_reqs': ssl.CERT_REQUIRED}
             if isinstance(conf.redis_use_ssl, dict):
@@ -154,12 +159,24 @@ def get_redis(app=None):
         else:
             connection = StrictRedis.from_url(conf.redis_url, decode_responses=True)
 
-        if retry_period is None:
-            app.redbeat_redis = connection
-        else:
-            app.redbeat_redis = RetryingConnection(retry_period, connection)
+        if connection:
+            _set_redbeat_connect(app, REDBEAT_REDIS_KEY, connection, retry_period)
+
+    if hasattr(app, REDBEAT_SENTINEL_KEY) and isinstance(getattr(app, REDBEAT_SENTINEL_KEY), Sentinel):
+        sentinel = getattr(app, REDBEAT_SENTINEL_KEY)
+        connection = sentinel.master_for(
+            redis_options.get('service_name', 'master'), db=redis_options.get('db', 0)
+        )
+        _set_redbeat_connect(app, REDBEAT_REDIS_KEY, connection, retry_period)
+
+    return getattr(app, REDBEAT_REDIS_KEY)
+
 
-    return app.redbeat_redis
+def _set_redbeat_connect(app, connect_name, connection, retry_period):
+    if retry_period is None:
+        setattr(app, connect_name, connection)
+    else:
+        setattr(app, connect_name, RetryingConnection(retry_period, connection))
 
 
 ADD_ENTRY_ERROR = """\