Repository: celery/celery

Issue:
Flags broker_connection_retry_on_startup & broker_connection_retry aren’t reliable 
<!--
Please fill this template entirely and do not erase parts of it.
We reserve the right to close without a response
bug reports which are incomplete.
-->
# Checklist
<!--
To check an item on the list replace [ ] with [x].
-->
- [x] I have verified that the issue exists against the `main` branch of Celery.
- [ ] This has already been asked to the [discussions forum](https://github.com/celery/celery/discussions) first.
- [x] I have read the relevant section in the
  [contribution guide](https://docs.celeryq.dev/en/main/contributing.html#other-bugs)
  on reporting bugs.
- [x] I have checked the [issues list](https://github.com/celery/celery/issues?q=is%3Aissue+label%3A%22Issue+Type%3A+Bug+Report%22+-label%3A%22Category%3A+Documentation%22)
  for similar or identical bug reports.
- [x] I have checked the [pull requests list](https://github.com/celery/celery/pulls?q=is%3Apr+label%3A%22PR+Type%3A+Bugfix%22+-label%3A%22Category%3A+Documentation%22)
  for existing proposed fixes.
- [x] I have checked the [commit log](https://github.com/celery/celery/commits/main)
  to find out if the bug was already fixed in the main branch.
- [x] I have included all related issues and possible duplicate issues
  in this issue (If there are none, check this box anyway).

## Mandatory Debugging Information

- [ ] I have included the output of ``celery -A proj report`` in the issue.
    (if you are not able to do this, then at least specify the Celery
     version affected).
- [x] I have verified that the issue exists against the `main` branch of Celery.
- [ ] I have included the contents of ``pip freeze`` in the issue.
- [x] I have included all the versions of all the external dependencies required
  to reproduce this bug.

## Optional Debugging Information
<!--
Try some of the below if you think they are relevant.
It will help us figure out the scope of the bug and how many users it affects.
-->
- [x] I have tried reproducing the issue on more than one Python version
  and/or implementation.
- [x] I have tried reproducing the issue on more than one message broker and/or
  result backend.
- [ ] I have tried reproducing the issue on more than one version of the message
  broker and/or result backend.
- [ ] I have tried reproducing the issue on more than one operating system.
- [ ] I have tried reproducing the issue on more than one workers pool.
- [ ] I have tried reproducing the issue with autoscaling, retries,
  ETA/Countdown & rate limits disabled.
- [ ] I have tried reproducing the issue after downgrading
  and/or upgrading Celery and its dependencies.

## Related Issues and Possible Duplicates
<!--
Please make sure to search and mention any related issues
or possible duplicates to this issue as requested by the checklist above.

This may or may not include issues in other repositories that the Celery project
maintains or other repositories that are dependencies of Celery.

If you don't know how to mention issues, please refer to Github's documentation
on the subject: https://help.github.com/en/articles/autolinked-references-and-urls#issues-and-pull-requests
-->

#### Related Issues

- None

#### Possible Duplicates

- None

## Environment & Settings
<!-- Include the contents of celery --version below -->
**Celery version**: 5.3.x

# Steps to Reproduce

This issue covers multiple cases that originate in a central bug.
The flag [broker_connection_retry_on_startup](https://docs.celeryq.dev/en/stable/userguide/configuration.html#broker-connection-retry-on-startup) is using https://github.com/celery/celery/blob/2cde29d9fb6a8f8f805bec5d97b36bc930bcb52f/celery/worker/consumer/consumer.py#L340 to determine if the connection is happening at startup. This is incorrect as it does not reliably validate this condition, which causes multiple use-cases to fail.

As part of the implementation of this flag in regard to [broker_connection_retry](https://docs.celeryq.dev/en/stable/userguide/configuration.html#broker-connection-retry), this flag also suffers bugs from this situation.

## Required Dependencies
<!-- Please fill the required dependencies to reproduce this issue -->
- **Minimal Python Version**: 3.8
- **Minimal Celery Version**: 5.3
- **Minimal Kombu Version**: 5.3
- **Minimal Broker Version**: `latest` rabbitmq or `latest` redis
- **Minimal Result Backend Version**: N/A or Unknown
- **Minimal OS and/or Kernel Version**: N/A or Unknown
- **Minimal Broker Client Version**: N/A or Unknown
- **Minimal Result Backend Client Version**: N/A or Unknown

## Minimally Reproducible Test Case 1

<details>
<p>

1. Turn off the broker container.
2. set `broker_connection_retry_on_startup` True.
3. set `broker_connection_retry` False.
4. Run celery worker.
4.1 Wait for the connection retry 1/100…
5. Turn on the broker container.
5.1. Wait for connection to broker.
6. Turn off the broker container.

</p>
</details>

# Expected Behavior
Worker should shut down.

# Actual Behavior
Worker retries to connect.

## Minimally Reproducible Test Case 2

<details>
<p>

1. Turn off the broker container.
2. set `broker_connection_retry_on_startup` False.
3. set `broker_connection_retry` True.
4. Run celery worker.

</p>
</details>

# Expected Behavior
Worker should shut down.

# Actual Behavior
Worker retries to connect.

## Minimally Reproducible Test Case 3

<details>
<p>

1. Turn on the broker container.
2. set `broker_connection_retry_on_startup` True.
3. set `broker_connection_retry` False.
4. Run celery worker.
5. Turn off the broker container.

</p>
</details>

# Expected Behavior
Worker should shut down.

# Actual Behavior
Worker retries to connect.

## Minimally Reproducible Test Case 4

<details>
<p>

1. Turn on the broker container.
2. set `broker_connection_retry_on_startup` False.
3. set `broker_connection_retry` True.
4. Run celery worker.
5. Turn off the broker container.

</p>
</details>

# Expected Behavior
Worker retries to connect.

# Actual Behavior
Worker shuts down

# Potential Fix

To fix all of these cases, a potential fix can be adding a flag to the consumer that will determine this exact _startup_ condition, and then use it where `broker_connection_retry_on_startup` is used so on startup it will respect `broker_connection_retry_on_startup` and afterward `broker_connection_retry`, unless `broker_connection_retry_on_startup` is `None`, in which case `broker_connection_retry` will also determine the startup condition.


Discussion:

Resolving patch:
diff --git a/celery/worker/consumer/consumer.py b/celery/worker/consumer/consumer.py
--- a/celery/worker/consumer/consumer.py
+++ b/celery/worker/consumer/consumer.py
@@ -153,6 +153,10 @@ class Consumer:
 
     restart_count = -1  # first start is the same as a restart
 
+    #: This flag will be turned off after the first failed
+    #: connection attempt.
+    first_connection_attempt = True
+
     class Blueprint(bootsteps.Blueprint):
         """Consumer blueprint."""
 
@@ -337,7 +341,8 @@ def start(self):
             except recoverable_errors as exc:
                 # If we're not retrying connections, we need to properly shutdown or terminate
                 # the Celery main process instead of abruptly aborting the process without any cleanup.
-                is_connection_loss_on_startup = self.restart_count == 0
+                is_connection_loss_on_startup = self.first_connection_attempt
+                self.first_connection_attempt = False
                 connection_retry_type = self._get_connection_retry_type(is_connection_loss_on_startup)
                 connection_retry = self.app.conf[connection_retry_type]
                 if not connection_retry:
@@ -488,13 +493,17 @@ def _error_handler(exc, interval, next_step=CONNECTION_RETRY_STEP):
 
         # Remember that the connection is lazy, it won't establish
         # until needed.
-        # If broker_connection_retry_on_startup is not set, revert to broker_connection_retry
-        # to determine whether connection retries are disabled.
 
         # TODO: Rely only on broker_connection_retry_on_startup to determine whether connection retries are disabled.
         #       We will make the switch in Celery 6.0.
 
+        retry_disabled = False
+
         if self.app.conf.broker_connection_retry_on_startup is None:
+            # If broker_connection_retry_on_startup is not set, revert to broker_connection_retry
+            # to determine whether connection retries are disabled.
+            retry_disabled = not self.app.conf.broker_connection_retry
+
             warnings.warn(
                 CPendingDeprecationWarning(
                     f"The broker_connection_retry configuration setting will no longer determine\n"
@@ -502,16 +511,23 @@ def _error_handler(exc, interval, next_step=CONNECTION_RETRY_STEP):
                     f"If you wish to retain the existing behavior for retrying connections on startup,\n"
                     f"you should set broker_connection_retry_on_startup to {self.app.conf.broker_connection_retry}.")
             )
+        else:
+            if self.first_connection_attempt:
+                retry_disabled = not self.app.conf.broker_connection_retry_on_startup
+            else:
+                retry_disabled = not self.app.conf.broker_connection_retry
 
-        if not self.app.conf.broker_connection_retry and not self.app.conf.broker_connection_retry_on_startup:
+        if retry_disabled:
             # Retry disabled, just call connect directly.
             conn.connect()
+            self.first_connection_attempt = False
             return conn
 
         conn = conn.ensure_connection(
             _error_handler, self.app.conf.broker_connection_max_retries,
             callback=maybe_shutdown,
         )
+        self.first_connection_attempt = False
         return conn
 
     def _flush_events(self):


Test patch:
diff --git a/t/unit/worker/test_consumer.py b/t/unit/worker/test_consumer.py
--- a/t/unit/worker/test_consumer.py
+++ b/t/unit/worker/test_consumer.py
@@ -422,8 +422,11 @@ def test_cancel_long_running_tasks_on_connection_loss__warning(self):
 
     @pytest.mark.parametrize("broker_connection_retry", [True, False])
     @pytest.mark.parametrize("broker_connection_retry_on_startup", [None, False])
-    def test_ensure_connected(self, subtests, broker_connection_retry, broker_connection_retry_on_startup):
+    @pytest.mark.parametrize("first_connection_attempt", [True, False])
+    def test_ensure_connected(self, subtests, broker_connection_retry, broker_connection_retry_on_startup,
+                              first_connection_attempt):
         c = self.get_consumer()
+        c.first_connection_attempt = first_connection_attempt
         c.app.conf.broker_connection_retry_on_startup = broker_connection_retry_on_startup
         c.app.conf.broker_connection_retry = broker_connection_retry
 
@@ -457,9 +460,7 @@ def test_start_raises_connection_error(self,
                                            is_connection_loss_on_startup,
                                            caplog, subtests):
         c = self.get_consumer()
-        # in order to reproduce the actual behavior: if this is the startup, then restart count has not been
-        # incremented yet, and is therefore -1.
-        c.restart_count = -1 if is_connection_loss_on_startup else 1
+        c.first_connection_attempt = True if is_connection_loss_on_startup else False
         c.app.conf['broker_connection_retry'] = False
         c.app.conf['broker_connection_retry_on_startup'] = broker_connection_retry_on_startup
         c.blueprint.start.side_effect = ConnectionError()