connection_pool: Add additional health check for pool instances

Allow `ConnectionPool` to exclude instances until an application running on top
of Tarantool has completed its post-startup initialization. The built-in
health check only confirms that Tarantool is reachable and `box.info.status ==
"running"`. The application may still be initializing its own components after
Tarantool starts, so an additional health check prevents requests from being
routed to the instance until the application reports that it is ready.

Author: ThCompiler

CHANGELOG.md

@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
- 
+
+- Support additional health check for pool instances (PR #348).
+
 ### Changed
 
 ### Fixed

tarantool/connection_pool.py

@@ -398,7 +398,8 @@ def __init__(self,
                  connection_timeout=CONNECTION_TIMEOUT,
                  strategy_class=RoundRobinStrategy,
                  refresh_delay=POOL_REFRESH_DELAY,
-                 fetch_schema=True):
+                 fetch_schema=True,
+                 additional_health_check=None):
         """
         :param addrs: List of dictionaries describing server addresses:
 
@@ -476,6 +477,17 @@ def __init__(self,
         :param fetch_schema: Refer to
             :paramref:`~tarantool.Connection.params.fetch_schema`.
 
+        :param additional_health_check: An additional health check for each
+            pool instance. It is called with the instance connection after
+            the built-in checks succeed.
+            Return :attr:`Status.UNHEALTHY` to exclude the
+            instance from request routing until the next successful health
+            check; return :attr:`Status.HEALTHY` to keep it available. The
+            check is performed when the pool connects and during periodic
+            state refreshes. It should be fast and must not raise exceptions;
+            the pool does not handle callback errors.
+        :type additional_health_check: :obj:`typing.Callable[[Connection], Status]`, optional
+
         :raise: :exc:`~tarantool.error.ConfigurationError`,
             :class:`~tarantool.Connection` exceptions
 
@@ -497,6 +509,8 @@ def __init__(self,
             new_addrs.append(new_addr)
         self.addrs = new_addrs
 
+        self.additional_health_check = additional_health_check
+
         # Create connections
         self.pool = {}
         self.refresh_delay = refresh_delay
@@ -603,6 +617,12 @@ def _get_new_state(self, unit):
             warn(msg, PoolTopologyWarning)
             return InstanceState(Status.UNHEALTHY)
 
+        if self.additional_health_check is not None:
+            status = self.additional_health_check(conn)
+
+            if status == Status.UNHEALTHY:
+                return InstanceState(Status.UNHEALTHY)
+
         return InstanceState(Status.HEALTHY, read_only)
 
     def _refresh_state(self, key):

test/suites/test_pool.py

@@ -594,14 +594,24 @@ def test_17_instance_bootstrap_error_does_not_kill_refresh(self):
         warnings.simplefilter('ignore', category=PoolTolopogyWarning)
 
         self.set_cluster_ro([False, True, True, True, True])
+        callback_calls = 0
+        target_port = self.addrs[0]['port']
+
+        def additional_health_check(conn):
+            nonlocal callback_calls
+            if conn.port == target_port:
+                callback_calls += 1
+            return Status.HEALTHY
 
         self.pool = tarantool.ConnectionPool(
             addrs=self.addrs,
             user='test',
             password='test',
-            refresh_delay=0.2)
+            refresh_delay=0.2,
+            additional_health_check=additional_health_check)
 
         self.pool.ping(mode=tarantool.Mode.RW)
+        callback_calls = 0
 
         unit = self.pool.pool[f"{self.addrs[0]['host']}:{self.addrs[0]['port']}"]
 
@@ -622,6 +632,7 @@ def expect_instance_unhealthy_and_refresh_alive():
             self.assertTrue(unit.thread.is_alive(),
                             'refresh thread died on a DatabaseError')
             self.assertEqual(unit.state.status, Status.UNHEALTHY)
+            self.assertEqual(callback_calls, 0)
 
         self.retry(func=expect_instance_unhealthy_and_refresh_alive)
 
@@ -636,6 +647,72 @@ def expect_rw_request_succeed():
 
         self.retry(func=expect_rw_request_succeed)
 
+    def test_18_additional_health_check_excludes_and_recovers_instance(self):
+        self.set_cluster_ro([False, False, True, False, False])
+        target_addr = self.addrs[2]
+        target_key = f"{target_addr['host']}:{target_addr['port']}"
+        unhealthy = True
+
+        def additional_health_check(conn):
+            if unhealthy and conn.port == target_addr['port']:
+                return Status.UNHEALTHY
+            return Status.HEALTHY
+
+        self.pool = tarantool.ConnectionPool(
+            addrs=self.addrs,
+            user='test',
+            password='test',
+            refresh_delay=0.2,
+            additional_health_check=additional_health_check)
+
+        self.assertEqual(self.pool.pool[target_key].state.status,
+                         Status.UNHEALTHY)
+        with self.assertRaises(PoolTopologyError):
+            self.pool.ping(mode=tarantool.Mode.RO)
+
+        unhealthy = False
+
+        def expect_instance_recovered():
+            self.assertEqual(self.pool.pool[target_key].state.status,
+                             Status.HEALTHY)
+            self.pool.ping(mode=tarantool.Mode.RO)
+
+        self.retry(func=expect_instance_recovered)
+
+    def test_19_additional_health_check_not_called_before_running(self):
+        warnings.simplefilter('ignore', category=PoolTopologyWarning)
+
+        self.set_cluster_ro([False, True, True, True, True])
+        target_addr = self.addrs[0]
+        callback_calls = 0
+
+        resp = self.servers[0].admin(r"""
+            rawset(_G, 'box_info_backup', box.info)
+            box.info = function()
+                local info = box_info_backup()
+                return {ro = info.ro, status = 'loading'}
+            end
+            return true
+        """)
+        assert_admin_success(resp)
+
+        def additional_health_check(conn):
+            nonlocal callback_calls
+            if conn.port == target_addr['port']:
+                callback_calls += 1
+            return Status.HEALTHY
+
+        self.pool = tarantool.ConnectionPool(
+            addrs=self.addrs,
+            user='test',
+            password='test',
+            additional_health_check=additional_health_check)
+
+        self.assertEqual(callback_calls, 0)
+        self.assertEqual(
+            self.pool.pool[f"{target_addr['host']}:{target_addr['port']}"].state.status,
+            Status.UNHEALTHY)
+
     def tearDown(self):
         if self.pool:
             self.pool.close()