Harden Redis batch publishing lifecycle

Author: Shivam

docs/reference/kombu.exceptions.rst

@@ -15,4 +15,4 @@
     .. autoexception:: LimitExceeded
     .. autoexception:: ConnectionLimitExceeded
     .. autoexception:: ChannelLimitExceeded
-
+    .. autoexception:: BatchPublishError

docs/reference/kombu.rst

@@ -162,7 +162,9 @@
         .. autoattribute:: auto_declare
         .. autoattribute:: on_return
         .. autoattribute:: connection
+        .. autoattribute:: supports_batch_publish
 
+        .. automethod:: batch
         .. automethod:: declare
         .. automethod:: maybe_declare
         .. automethod:: publish

docs/userguide/producers.rst

@@ -107,6 +107,8 @@ empty string, and set the routing key to be the name of the queue:
 Batch publishing
 ----------------
 
+.. versionadded:: 5.7
+
 :meth:`~kombu.Producer.batch` groups normal :meth:`~kombu.Producer.publish`
 calls so a supporting transport can send their final broker operations
 together:
@@ -172,7 +174,11 @@ The standard Redis, Redis TLS, and Redis Sentinel transports share this batch
 implementation. Standard Redis is covered by integration tests. TLS and
 Sentinel use the same channel implementation but are not covered by
 real-service batch integration tests. Redis Cluster is not currently an
-upstream Kombu transport and is not supported or tested by this API.
+upstream Kombu transport and is not supported or tested by this API. Redis
+publishing also remains immediate when the underlying connection pool enables
+automatic retries, including ``retry_on_timeout`` or a custom retry policy,
+because replaying a pipeline after an ambiguous failure can publish duplicate
+messages.
 
 Serialization
 =============

kombu/messaging.py

@@ -22,41 +22,25 @@
 DEFAULT_BATCH_SIZE = 1000
 
 
-class _ImmediatePublishBatch:
-    """Transport-neutral batch that retains immediate publication."""
-
-    def __init__(self, channel):
-        self.channel = channel
-
-    def publish(self, message, **kwargs):
-        return self.channel.basic_publish(message, **kwargs)
-
-    def flush(self):
-        """Flush buffered messages (there are none)."""
-
-    def discard(self):
-        """Discard buffered messages (there are none)."""
-
-    def close(self):
-        """Release batch resources (there are none)."""
-
-
 class _ProducerBatchState:
     """State shared by nested batch contexts in one producer thread."""
 
     def __init__(self, max_size):
         self.max_size = max_size
         self.aborted = False
         self.session = None
+        self.immediate = False
 
     def _get_session(self, channel):
+        if self.immediate:
+            return None
         if self.session is None:
             create_batch = getattr(channel, 'create_publish_batch', None)
-            self.session = (
-                create_batch(max_size=self.max_size)
-                if create_batch is not None
-                else _ImmediatePublishBatch(channel)
-            )
+            supports_batch = getattr(channel, 'supports_batch_publish', True)
+            if create_batch is None or not supports_batch:
+                self.immediate = True
+                return None
+            self.session = create_batch(max_size=self.max_size)
         elif self.session.channel is not channel:
             raise RuntimeError(
                 'Producer channel changed while a publish batch was active',
@@ -66,7 +50,10 @@ def _get_session(self, channel):
     def publish(self, channel, message, **kwargs):
         if self.aborted:
             raise RuntimeError('Cannot publish through an aborted batch')
-        return self._get_session(channel).publish(message, **kwargs)
+        session = self._get_session(channel)
+        if session is None:
+            return channel.basic_publish(message, **kwargs)
+        return session.publish(message, **kwargs)
 
     def flush(self):
         if self.aborted:
@@ -116,19 +103,33 @@ def __exit__(self, exc_type, exc_value, traceback):
         if state is None:
             return None
 
+        cleanup_error = None
         try:
-            if exc_type is not None:
-                state.abort()
-
             if self.is_outermost:
                 del self.producer._batch_local.current
+
+            if exc_type is not None:
                 try:
-                    if exc_type is None and not state.aborted:
+                    state.abort()
+                except BaseException as exc:
+                    cleanup_error = exc
+
+            if self.is_outermost:
+                if exc_type is None and not state.aborted:
+                    try:
                         state.flush()
-                finally:
+                    except BaseException as exc:
+                        cleanup_error = exc
+                try:
                     state.close()
+                except BaseException as exc:
+                    if cleanup_error is None:
+                        cleanup_error = exc
         finally:
             self.active = False
+
+        if exc_type is None and cleanup_error is not None:
+            raise cleanup_error
         return None
 
     def flush(self):
@@ -233,9 +234,23 @@ def supports_batch_publish(self):
         if connection is None:
             return False
         try:
-            return connection.transport.implements.batch_publish
+            transport_support = connection.transport.implements.batch_publish
         except AttributeError:
             return False
+        channel = self._channel
+        if channel is not None and not isinstance(channel, ChannelPromise):
+            configured_support = getattr(
+                channel,
+                'supports_batch_publish',
+                transport_support,
+            )
+        else:
+            configured_support = getattr(
+                connection.transport,
+                'supports_batch_publish',
+                transport_support,
+            )
+        return bool(transport_support and configured_support)
 
     def batch(self, max_size=DEFAULT_BATCH_SIZE):
         """Group normal :meth:`publish` calls into transport-owned batches.

kombu/transport/redis.py

@@ -693,6 +693,15 @@ def fds(self):
         return self._fd_to_chan
 
 
+def _batch_publish_retry_safe(connection_kwargs):
+    """Return whether Redis will execute a pipeline without retrying it."""
+    return not (
+        connection_kwargs.get('retry_on_timeout')
+        or connection_kwargs.get('retry_on_error')
+        or connection_kwargs.get('retry') is not None
+    )
+
+
 class PublishBatch:
     """Buffer Redis publish commands in a non-transactional pipeline."""
 
@@ -832,6 +841,14 @@ class Channel(virtual.Channel):
     max_connections = 10
     health_check_interval = DEFAULT_HEALTH_CHECK_INTERVAL
     client_name = None
+
+    @property
+    def supports_batch_publish(self):
+        """Return whether pipelines can run without ambiguous write replay."""
+        return (
+            not self.retry_on_timeout
+            and _batch_publish_retry_safe(self.pool.connection_kwargs)
+        )
     #: Transport option to disable fanout keyprefix.
     #: Can also be string, in which case it changes the default
     #: prefix ('/{db}.') into to something else.  The prefix must
@@ -1752,6 +1769,17 @@ class Transport(virtual.Transport):
         exchange_type=frozenset(['direct', 'topic', 'fanout']),
     )
 
+    @property
+    def supports_batch_publish(self):
+        """Return whether configured timeout handling permits safe batching."""
+        if not _batch_publish_retry_safe(self.client.transport_options):
+            return False
+        hostname = self.client.hostname or ''
+        if '://' not in hostname:
+            return True
+        *_, query = _parse_url(hostname)
+        return _batch_publish_retry_safe(query)
+
     if redis:
         connection_errors, channel_errors = get_redis_error_classes()
 

t/integration/test_redis.py

@@ -222,6 +222,57 @@ def callback(body, message):
 @pytest.mark.flaky(reruns=5, reruns_delay=2)
 class test_RedisPublishBatch:
 
+    def test_retry_on_timeout_keeps_publication_immediate(self, connection):
+        connection.transport_options = {
+            **connection.transport_options,
+            'retry_on_timeout': True,
+        }
+        queue = kombu.Queue('batch_retry_on_timeout_queue')
+
+        with connection as conn:
+            with conn.channel() as channel:
+                producer = kombu.Producer(channel, serializer='json')
+
+                assert producer.supports_batch_publish is False
+                with producer.batch():
+                    producer.publish(
+                        {'delivery': 'immediate'},
+                        exchange='',
+                        routing_key=queue.name,
+                        declare=[queue],
+                    )
+                    message = queue(channel).get(no_ack=True)
+
+        assert message.payload == {'delivery': 'immediate'}
+
+    def test_manual_flush_sends_and_batch_continues(self, connection):
+        queue = kombu.Queue('batch_manual_flush_queue')
+
+        with connection as conn:
+            with conn.channel() as channel:
+                producer = kombu.Producer(channel, serializer='json')
+                bound_queue = queue(channel)
+
+                with producer.batch() as batch:
+                    producer.publish(
+                        {'position': 'first'},
+                        exchange='',
+                        routing_key=queue.name,
+                        declare=[queue],
+                    )
+                    batch.flush()
+                    first = bound_queue.get(no_ack=True)
+                    producer.publish(
+                        {'position': 'second'},
+                        exchange='',
+                        routing_key=queue.name,
+                    )
+
+                second = bound_queue.get(no_ack=True)
+
+        assert first.payload == {'position': 'first'}
+        assert second.payload == {'position': 'second'}
+
     def test_direct_priority_and_fifo(self, connection):
         exchange = kombu.Exchange('batch_direct_exchange', type='direct')
         queue = kombu.Queue(