Redis: deliver rotating StreamingCredentialProvider tokens to long-lived BRPOP and pub/sub connections (#2563)

* Implement fix to flush tokens onto the BRPOP and Pub/sub LISTEN connections

* robustify

* changes made in response to review

---------

Co-authored-by: Asif Saif Uddin {"Auvi":"অভি"} <auvipy@gmail.com>

Author: ThomasAitken

docs/reference/kombu.transport.redis.rst

@@ -35,6 +35,57 @@
     .. versionadded:: 5.7.0
     Supports Queue TTL
 
+    Streaming credentials / automatic re-authentication
+    ----------------------------------------------------
+    .. versionadded:: 5.7.0
+
+    The transport supports rotating credentials supplied by a
+    ``redis.credentials.StreamingCredentialProvider`` (for example the
+    Microsoft Entra ID provider from ``redis-entraid``, or an AWS ElastiCache
+    IAM provider).  Configure it like any other credential provider:
+
+    .. code-block:: python
+
+        app.conf.broker_transport_options = {
+            "credential_provider": my_streaming_credential_provider,
+        }
+
+    Such providers emit a fresh authentication token in the background before
+    the current one expires.  redis-py delivers these tokens to *pooled*
+    connections when they are released back to the pool, but the Redis
+    transport keeps two connections busy for the entire lifetime of the
+    worker — the ``BRPOP`` connection (used to consume regular queues) and the
+    pub/sub ``LISTEN`` connection (used to consume fanout queues) — so they are
+    never released and would otherwise never receive a rotated token.  The
+    broker then severs them once the original credentials expire (e.g. AWS
+    ElastiCache with IAM auth enforces a hard 12-hour connection limit),
+    causing redelivered messages, interrupted in-flight tasks and brief worker
+    unavailability.
+
+    To avoid this, the transport periodically flushes any pending token onto
+    those long-lived connections from the event loop:
+
+    * the ``BRPOP`` connection is re-authenticated in place with an ``AUTH``
+      command, sent only when no blocking pop is in flight;
+    * the pub/sub ``LISTEN`` connection cannot process ``AUTH`` while
+      subscribed under RESP2, so it is transparently reconnected (and
+      re-subscribed) to pick up the new credentials.  Under RESP3, redis-py
+      re-authenticates pub/sub connections itself and the transport leaves
+      them untouched.
+
+    How often the flush runs is controlled by the ``reauth_check_interval``
+    transport option (seconds, default ``10``):
+
+    .. code-block:: python
+
+        app.conf.broker_transport_options = {
+            "credential_provider": my_streaming_credential_provider,
+            "reauth_check_interval": 10,
+        }
+
+    When no streaming credential provider is configured this machinery is a
+    cheap no-op, so it is always safe to leave enabled.
+
     Queue arguments
     ---------------
     The following queue argument is supported. Pass it per-queue via

kombu/transport/redis.py

@@ -48,6 +48,11 @@
 * ``queue_order_strategy``
 * ``max_connections``
 * ``health_check_interval``
+* ``reauth_check_interval``: (int) How often, in seconds, to flush pending
+  streaming re-authentication tokens (emitted by a
+  ``redis.credentials.StreamingCredentialProvider`` such as the Entra ID /
+  IAM providers) onto the long-lived BRPOP and pub/sub connections held by
+  the transport. Defaults to ``10``. See :meth:`Channel.maybe_reauth`.
 * ``retry_on_timeout``
 * ``priority_steps``
 * ``client_name``: (str) The name to use when connecting to Redis server.
@@ -113,6 +118,11 @@
 
 DEFAULT_HEALTH_CHECK_INTERVAL = 25
 
+#: How often (in seconds) to flush pending streaming re-authentication
+#: tokens onto the long-lived BRPOP and pub/sub connections.  See
+#: :meth:`Channel.maybe_reauth` for why this is needed.
+DEFAULT_REAUTH_CHECK_INTERVAL = 10
+
 PRIORITY_STEPS = [0, 3, 6, 9]
 
 error_classes_t = namedtuple('error_classes_t', (
@@ -600,6 +610,20 @@ def maybe_check_subclient_health(self):
                     )
                     return
 
+    def maybe_reauth(self):
+        for channel in self._channels:
+            try:
+                channel.maybe_reauth()
+            except channel.connection_errors:
+                # Connection is broken; skip this cycle and retry next tick.
+                # Stop iterating to avoid repeated exceptions/log spam when
+                # the broker is down (channels share the broken connection).
+                logger.debug(
+                    'maybe_reauth: connection error, '
+                    'will retry on next cycle', exc_info=True
+                )
+                return
+
     def on_readable(self, fileno):
         chan_type = self._fd_to_chan.get(fileno)
         if chan_type is None:
@@ -1033,13 +1057,154 @@ def _brpop_read(self, **options):
                 raise Empty()
         finally:
             self._in_poll = None
+            # The BRPOP connection is now idle (its reply has been fully
+            # consumed and no new BRPOP has been issued yet): this is the
+            # safe moment to flush any pending streaming re-auth token.
+            self._flush_brpop_reauth()
 
     def _poll_error(self, type, **options):
         if type == 'LISTEN':
             self.subclient.parse_response()
         else:
             self.client.parse_response(self.client.connection, type)
 
+    @staticmethod
+    def _pending_reauth_token(connection):
+        """Return a streaming re-auth token stored on ``connection``, if any.
+
+        redis-py's :class:`~redis.event.RegisterReAuthForPooledConnections`
+        listener calls ``Connection.set_re_auth_token`` for every *in-use*
+        pooled connection whenever a
+        :class:`~redis.credentials.StreamingCredentialProvider` emits a fresh
+        token.  The token is stored on the connection and only turned into an
+        actual ``AUTH`` command when the connection is released back to the
+        pool.  We key off this attribute so that we do not have to couple to
+        the credential provider itself; it is only ever set when streaming
+        re-authentication is in effect.
+        """
+        if connection is None:
+            return None
+        return getattr(connection, '_re_auth_token', None)
+
+    @staticmethod
+    def _pubsub_reauth_handled_by_redis(connection):
+        """Whether redis-py re-authenticates ``connection`` itself.
+
+        redis-py only re-authenticates *subscribed* connections in place when
+        the RESP3 protocol has been negotiated (see
+        ``redis.event.RegisterReAuthForPubSub``); a RESP2 subscriber
+        connection cannot process an ``AUTH`` command at all.  When RESP3 is in
+        use we therefore leave the pub/sub connection to redis-py and avoid
+        interfering with it.
+        """
+        get_protocol = getattr(connection, 'get_protocol', None)
+        if get_protocol is not None:
+            protocol = get_protocol()
+        else:
+            protocol = getattr(connection, 'protocol', 2)
+        return str(protocol) == '3'
+
+    def maybe_reauth(self):
+        """Flush pending streaming re-auth tokens onto long-lived connections.
+
+        The transport holds two connections for the lifetime of the worker
+        that are never released back to the pool: the ``BRPOP`` connection
+        (used to consume from ordinary queues) and the pub/sub ``LISTEN``
+        connection (used to consume from fanout queues).  Because they are
+        never released, redis-py's release-triggered re-authentication never
+        fires for them, so a streaming credential provider's rotated tokens
+        never reach them and the broker eventually severs the connections when
+        the original credentials expire (e.g. the 12h limit imposed by AWS
+        ElastiCache with IAM auth).
+
+        This is called periodically from the event loop (a single thread), so
+        it can safely flush the tokens without racing the socket reads.
+        """
+        self._flush_brpop_reauth()
+        self._flush_listen_reauth()
+
+    def _flush_brpop_reauth(self):
+        """Send a pending re-auth token's ``AUTH`` on the BRPOP connection.
+
+        Only safe to do while no ``BRPOP`` command is in flight, otherwise the
+        ``AUTH`` reply would interleave with the blocking pop's reply.  We are
+        called both from the periodic timer and from :meth:`_brpop_read` (right
+        after a reply has been fully consumed), so the token is flushed at the
+        first idle moment after it is emitted.
+        """
+        if self._in_poll:
+            # A BRPOP is outstanding; retry at the next idle opportunity.
+            return
+        client = self.__dict__.get('client')  # only if property is cached
+        connection = getattr(client, 'connection', None)
+        if self._pending_reauth_token(connection) is None:
+            return
+        if getattr(connection, '_sock', None) is None:
+            # The socket is already down (e.g. a BRPOP connection error just
+            # disconnected it, or a previous flush failed).  ``re_auth`` would
+            # transparently reconnect a *fresh* socket via ``send_command``,
+            # but behind the poller's back: ``_register_BRPOP`` would then skip
+            # re-registering it (its ``_chan_to_sock`` entry survives the
+            # disconnect and ``_sock`` is no longer ``None``), so the new fd is
+            # never handed to the event loop and the channel silently stalls.
+            # Leave it to ``_register_BRPOP`` to reconnect *and* re-register;
+            # ``on_connect`` re-authenticates with the current credentials.
+            return
+        try:
+            connection.re_auth()
+        except self.connection_errors + (self.ResponseError,):
+            # The AUTH failed: a network/auth error (ConnectionError,
+            # AuthenticationError, ...) or a rejected token surfacing as a
+            # ResponseError.  Drop the socket so the next poll reconnects and
+            # authenticates with fresh credentials from the credential
+            # provider.  Also clear the stored token: ``re_auth`` only clears
+            # it on success, and once the socket is dropped the reconnect
+            # applies the current credentials via ``on_connect`` — re-sending
+            # this same (possibly expired/rejected) token in place would just
+            # fail again on every tick.  This runs from ``_brpop_read``'s
+            # ``finally`` block, so it must never raise.
+            warning('Redis streaming re-auth failed on BRPOP connection; '
+                    'reconnecting', exc_info=True)
+            try:
+                connection.set_re_auth_token(None)
+            except AttributeError:
+                pass
+            try:
+                connection.disconnect()
+            except self.connection_errors:
+                pass
+
+    def _flush_listen_reauth(self):
+        """Refresh credentials on the long-lived pub/sub (LISTEN) connection.
+
+        A subscribed RESP2 connection cannot process an ``AUTH`` command, so
+        the stored re-auth token can never be flushed in place.  Instead we
+        drop the connection; the poller reconnects and re-subscribes on the
+        next tick, authenticating with the current credentials from the
+        credential provider.  Fanout delivery is best-effort, so the brief
+        reconnect is far less disruptive than a broker-forced disconnect.
+
+        Under RESP3, redis-py re-authenticates pub/sub connections itself, so
+        we leave those untouched.
+        """
+        subclient = self.__dict__.get('subclient')  # only if property cached
+        connection = getattr(subclient, 'connection', None)
+        if self._pending_reauth_token(connection) is None:
+            return
+        if self._pubsub_reauth_handled_by_redis(connection):
+            return
+        logger.info('Refreshing Redis pub/sub connection to apply rotated '
+                    'streaming credentials')
+        # Clear the stored token first so we do not reconnect again on the
+        # next cycle, then drop the socket.  The poller re-registers the
+        # LISTEN connection and re-subscribes on the next tick.
+        try:
+            connection.set_re_auth_token(None)
+        except AttributeError:
+            pass
+        self._in_listen = None
+        connection.disconnect()
+
     def _get(self, queue):
         with self.conn_or_acquire() as client:
             for pri in self.priority_steps:
@@ -1536,7 +1701,8 @@ def on_poll_start():
         # registering new ones. Without this, each reconnect accumulates
         # an extra entry in hub.timer._queue; they all fire against the
         # same cycle and can crash the event loop during reconnect.
-        for attr in ('_restore_messages_tref', '_subclient_health_tref'):
+        for attr in ('_restore_messages_tref', '_subclient_health_tref',
+                     '_reauth_tref'):
             old_tref = getattr(cycle, attr, None)
             if old_tref is not None:
                 old_tref.cancel()
@@ -1552,6 +1718,14 @@ def on_poll_start():
             health_check_interval,
             cycle.maybe_check_subclient_health
         )
+        reauth_check_interval = connection.client.transport_options.get(
+            'reauth_check_interval',
+            DEFAULT_REAUTH_CHECK_INTERVAL
+        )
+        cycle._reauth_tref = loop.call_repeatedly(
+            reauth_check_interval,
+            cycle.maybe_reauth
+        )
 
     def on_readable(self, fileno):
         """Handle AIO event for one of our file descriptors."""