Fixed #30398 -- Added CONN_HEALTH_CHECKS database setting.

The CONN_HEALTH_CHECKS setting can be used to enable database
connection health checks for Django's persistent DB connections.

Thanks Florian Apolloner for reviews.

6 files changed, 216 insertions, 6 deletions

diff --git a/django/db/backends/base/base.py b/django/db/backends/base/base.py
index 14b2623414..58dd6d43bd 100644
--- a/django/db/backends/base/base.py
+++ b/django/db/backends/base/base.py
@@ -92,6 +92,8 @@ class BaseDatabaseWrapper:
         self.close_at = None
         self.closed_in_transaction = False
         self.errors_occurred = False
+        self.health_check_enabled = False
+        self.health_check_done = False
 
         # Thread-safety related attributes.
         self._thread_sharing_lock = threading.Lock()
@@ -210,11 +212,14 @@ class BaseDatabaseWrapper:
         self.savepoint_ids = []
         self.atomic_blocks = []
         self.needs_rollback = False
-        # Reset parameters defining when to close the connection
+        # Reset parameters defining when to close/health-check the connection.
+        self.health_check_enabled = self.settings_dict['CONN_HEALTH_CHECKS']
         max_age = self.settings_dict['CONN_MAX_AGE']
         self.close_at = None if max_age is None else time.monotonic() + max_age
         self.closed_in_transaction = False
         self.errors_occurred = False
+        # New connections are healthy.
+        self.health_check_done = True
         # Establish the connection
         conn_params = self.get_connection_params()
         self.connection = self.get_new_connection(conn_params)
@@ -252,6 +257,7 @@ class BaseDatabaseWrapper:
         return wrapped_cursor
 
     def _cursor(self, name=None):
+        self.close_if_health_check_failed()
         self.ensure_connection()
         with self.wrap_database_errors:
             return self._prepare_cursor(self.create_cursor(name))
@@ -422,6 +428,7 @@ class BaseDatabaseWrapper:
         backends.
         """
         self.validate_no_atomic_block()
+        self.close_if_health_check_failed()
         self.ensure_connection()
 
         start_transaction_under_autocommit = (
@@ -519,12 +526,26 @@ class BaseDatabaseWrapper:
         raise NotImplementedError(
             "subclasses of BaseDatabaseWrapper may require an is_usable() method")
 
+    def close_if_health_check_failed(self):
+        """Close existing connection if it fails a health check."""
+        if (
+            self.connection is None or
+            not self.health_check_enabled or
+            self.health_check_done
+        ):
+            return
+
+        if not self.is_usable():
+            self.close()
+        self.health_check_done = True
+
     def close_if_unusable_or_obsolete(self):
         """
         Close the current connection if unrecoverable errors have occurred
         or if it outlived its maximum age.
         """
         if self.connection is not None:
+            self.health_check_done = False
             # If the application didn't restore the original autocommit setting,
             # don't take chances, drop the connection.
             if self.get_autocommit() != self.settings_dict['AUTOCOMMIT']:
@@ -536,6 +557,7 @@ class BaseDatabaseWrapper:
             if self.errors_occurred:
                 if self.is_usable():
                     self.errors_occurred = False
+                    self.health_check_done = True
                 else:
                     self.close()
                     return
diff --git a/django/db/utils.py b/django/db/utils.py
index f8d56bc418..82498d1df6 100644
--- a/django/db/utils.py
+++ b/django/db/utils.py
@@ -172,6 +172,7 @@ class ConnectionHandler(BaseConnectionHandler):
         if conn['ENGINE'] == 'django.db.backends.' or not conn['ENGINE']:
             conn['ENGINE'] = 'django.db.backends.dummy'
         conn.setdefault('CONN_MAX_AGE', 0)
+        conn.setdefault('CONN_HEALTH_CHECKS', False)
         conn.setdefault('OPTIONS', {})
         conn.setdefault('TIME_ZONE', None)
         for setting in ['NAME', 'USER', 'PASSWORD', 'HOST', 'PORT']:
diff --git a/docs/ref/databases.txt b/docs/ref/databases.txt
index 160b5b405e..d4da6bd637 100644
--- a/docs/ref/databases.txt
+++ b/docs/ref/databases.txt
@@ -62,8 +62,19 @@ At the end of each request, Django closes the connection if it has reached its
 maximum age or if it is in an unrecoverable error state. If any database
 errors have occurred while processing the requests, Django checks whether the
 connection still works, and closes it if it doesn't. Thus, database errors
-affect at most one request; if the connection becomes unusable, the next
-request gets a fresh connection.
+affect at most one request per each application's worker thread; if the
+connection becomes unusable, the next request gets a fresh connection.
+
+Setting :setting:`CONN_HEALTH_CHECKS` to ``True`` can be used to improve the
+robustness of connection reuse and prevent errors when a connection has been
+closed by the database server which is now ready to accept and serve new
+connections, e.g. after database server restart. The health check is performed
+only once per request and only if the database is being accessed during the
+handling of the request.
+
+.. versionchanged:: 4.1
+
+    The :setting:`CONN_HEALTH_CHECKS` setting was added.
 
 Caveats
 ~~~~~~~
diff --git a/docs/ref/settings.txt b/docs/ref/settings.txt
index 54d9c992c9..6c23ba6397 100644
--- a/docs/ref/settings.txt
+++ b/docs/ref/settings.txt
@@ -628,7 +628,25 @@ Default: ``0``
 
 The lifetime of a database connection, as an integer of seconds. Use ``0`` to
 close database connections at the end of each request — Django's historical
-behavior — and ``None`` for unlimited persistent connections.
+behavior — and ``None`` for unlimited :ref:`persistent database connections
+<persistent-database-connections>`.
+
+.. setting:: CONN_HEALTH_CHECKS
+
+``CONN_HEALTH_CHECKS``
+~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 4.1
+
+Default: ``False``
+
+If set to ``True``, existing :ref:`persistent database connections
+<persistent-database-connections>` will be health checked before they are
+reused in each request performing database access. If the health check fails,
+the connection will be re-established without failing the request when the
+connection is no longer usable but the database server is ready to accept and
+serve new connections (e.g. after database server restart closing existing
+connections).
 
 .. setting:: OPTIONS
 
diff --git a/docs/releases/4.1.txt b/docs/releases/4.1.txt
index 34611b4d4d..cc412bebfa 100644
--- a/docs/releases/4.1.txt
+++ b/docs/releases/4.1.txt
@@ -208,6 +208,11 @@ Models
   :class:`~django.db.models.expressions.Window` expression now accepts string
   references to fields and transforms.
 
+* The new :setting:`CONN_HEALTH_CHECKS` setting allows enabling health checks
+  for :ref:`persistent database connections <persistent-database-connections>`
+  in order to reduce the number of failed requests, e.g. after database server
+  restart.
+
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/backends/base/test_base.py b/tests/backends/base/test_base.py
index b07d6d3ceb..b2ebea5829 100644
--- a/tests/backends/base/test_base.py
+++ b/tests/backends/base/test_base.py
@@ -1,8 +1,8 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from django.db import DEFAULT_DB_ALIAS, connection, connections
 from django.db.backends.base.base import BaseDatabaseWrapper
-from django.test import SimpleTestCase, TestCase
+from django.test import SimpleTestCase, TestCase, skipUnlessDBFeature
 
 from ..models import Square
 
@@ -134,3 +134,156 @@ class ExecuteWrapperTests(TestCase):
         self.assertFalse(wrapper.called)
         self.assertEqual(connection.execute_wrappers, [])
         self.assertEqual(connections['other'].execute_wrappers, [])
+
+
+class ConnectionHealthChecksTests(SimpleTestCase):
+    databases = {'default'}
+
+    def setUp(self):
+        # All test cases here need newly configured and created connections.
+        # Use the default db connection for convenience.
+        connection.close()
+        self.addCleanup(connection.close)
+
+    def patch_settings_dict(self, conn_health_checks):
+        self.settings_dict_patcher = patch.dict(connection.settings_dict, {
+            **connection.settings_dict,
+            'CONN_MAX_AGE': None,
+            'CONN_HEALTH_CHECKS': conn_health_checks,
+        })
+        self.settings_dict_patcher.start()
+        self.addCleanup(self.settings_dict_patcher.stop)
+
+    def run_query(self):
+        with connection.cursor() as cursor:
+            cursor.execute('SELECT 42' + connection.features.bare_select_suffix)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_enabled(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        self.assertIs(old_connection, connection.connection)
+
+        # Simulate connection health check failing.
+        with patch.object(connection, 'is_usable', return_value=False) as mocked_is_usable:
+            self.run_query()
+            new_connection = connection.connection
+            # A new connection is established.
+            self.assertIsNot(new_connection, old_connection)
+            # Only one health check per "request" is performed, so the next
+            # query will carry on even if the health check fails. Next query
+            # succeeds because the real connection is healthy and only the
+            # health check failure is mocked.
+            self.run_query()
+            self.assertIs(new_connection, connection.connection)
+        self.assertEqual(mocked_is_usable.call_count, 1)
+
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # The underlying connection is being reused further with health checks
+        # succeeding.
+        self.run_query()
+        self.run_query()
+        self.assertIs(new_connection, connection.connection)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_enabled_errors_occurred(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate errors_occurred.
+        connection.errors_occurred = True
+        # Simulate request_started (the connection is healthy).
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled.
+        self.assertIs(old_connection, connection.connection)
+        # No additional health checks after the one in
+        # close_if_unusable_or_obsolete() are executed during this "request"
+        # when running queries.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_disabled(self):
+        self.patch_settings_dict(conn_health_checks=False)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled (connection is not).
+        self.assertIs(old_connection, connection.connection)
+        # Health checks are not performed.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+            # Health check wasn't performed and the connection is unchanged.
+            self.assertIs(old_connection, connection.connection)
+            self.run_query()
+            # The connection is unchanged after the next query either during
+            # the current "request".
+            self.assertIs(old_connection, connection.connection)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_set_autocommit_health_checks_enabled(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            # Simulate outermost atomic block: changing autocommit for
+            # a connection.
+            connection.set_autocommit(False)
+            self.run_query()
+            connection.commit()
+            connection.set_autocommit(True)
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled.
+        self.assertIs(old_connection, connection.connection)
+
+        # Simulate connection health check failing.
+        with patch.object(connection, 'is_usable', return_value=False) as mocked_is_usable:
+            # Simulate outermost atomic block: changing autocommit for
+            # a connection.
+            connection.set_autocommit(False)
+            new_connection = connection.connection
+            self.assertIsNot(new_connection, old_connection)
+            # Only one health check per "request" is performed, so a query will
+            # carry on even if the health check fails. This query succeeds
+            # because the real connection is healthy and only the health check
+            # failure is mocked.
+            self.run_query()
+            connection.commit()
+            connection.set_autocommit(True)
+            # The connection is unchanged.
+            self.assertIs(new_connection, connection.connection)
+        self.assertEqual(mocked_is_usable.call_count, 1)
+
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # The underlying connection is being reused further with health checks
+        # succeeding.
+        connection.set_autocommit(False)
+        self.run_query()
+        connection.commit()
+        connection.set_autocommit(True)
+        self.assertIs(new_connection, connection.connection)