Refactored connection count and SCH metric collection (#4001)

* Refactored connection count and SCH metric collection

* Codestyle fixes

* Fixed correct deprecation description

* Refactored state transitions

* Fixed transition for async

* Revert for backward-compatibility

* Fixed leaking connection states

* Refactor metrics export on cleanup

* Fixed potential deadlock

* Codestyle fixes

* Remvoved double decrement

* Added synchronisation for correct metric recording

---------

Co-authored-by: petyaslavova <petya.slavova@redis.com>

Author: vladvildanov

redis/asyncio/connection.py

@@ -59,6 +59,7 @@
 
 from redis.asyncio.observability.recorder import (
     record_connection_closed,
+    record_connection_count,
     record_connection_create_time,
     record_connection_wait_time,
     record_error_count,
@@ -1332,9 +1333,67 @@ def __repr__(self):
         )
 
     def reset(self):
+        # Record metrics for connections being removed before clearing
+        # (only if attributes exist - they won't during __init__)
+        if hasattr(self, "_available_connections") and hasattr(
+            self, "_in_use_connections"
+        ):
+            idle_count = len(self._available_connections)
+            in_use_count = len(self._in_use_connections)
+            if idle_count > 0 or in_use_count > 0:
+                pool_name = get_pool_name(self)
+                # Note: Using sync version since reset() is sync
+                from redis.observability.recorder import (
+                    record_connection_count as sync_record_connection_count,
+                )
+
+                if idle_count > 0:
+                    sync_record_connection_count(
+                        pool_name=pool_name,
+                        connection_state=ConnectionState.IDLE,
+                        counter=-idle_count,
+                    )
+                if in_use_count > 0:
+                    sync_record_connection_count(
+                        pool_name=pool_name,
+                        connection_state=ConnectionState.USED,
+                        counter=-in_use_count,
+                    )
+
         self._available_connections = []
         self._in_use_connections = weakref.WeakSet()
 
+    def __del__(self) -> None:
+        """Clean up connection pool and record metrics when garbage collected."""
+        try:
+            if not hasattr(self, "_available_connections") or not hasattr(
+                self, "_in_use_connections"
+            ):
+                return
+            idle_count = len(self._available_connections)
+            in_use_count = len(self._in_use_connections)
+            if idle_count > 0 or in_use_count > 0:
+                pool_name = get_pool_name(self)
+                # Note: Using sync version since __del__ is sync
+                from redis.observability.recorder import (
+                    record_connection_count as sync_record_connection_count,
+                )
+
+                if idle_count > 0:
+                    sync_record_connection_count(
+                        pool_name=pool_name,
+                        connection_state=ConnectionState.IDLE,
+                        counter=-idle_count,
+                    )
+                if in_use_count > 0:
+                    sync_record_connection_count(
+                        pool_name=pool_name,
+                        connection_state=ConnectionState.USED,
+                        counter=-in_use_count,
+                    )
+        except Exception:
+            pass
+
     def can_get_connection(self) -> bool:
         """Return True if a connection can be retrieved from the pool."""
         return (
@@ -1361,6 +1420,29 @@ async def get_connection(self, command_name=None, *keys, **options):
             )
             is_created = connections_after > connections_before
 
+        # Record state transition for observability
+        # This ensures counters stay balanced if ensure_connection() fails and release() is called
+        pool_name = get_pool_name(self)
+        if is_created:
+            # New connection created and acquired: just USED +1
+            await record_connection_count(
+                pool_name=pool_name,
+                connection_state=ConnectionState.USED,
+                counter=1,
+            )
+        else:
+            # Existing connection acquired from pool: IDLE -> USED
+            await record_connection_count(
+                pool_name=pool_name,
+                connection_state=ConnectionState.IDLE,
+                counter=-1,
+            )
+            await record_connection_count(
+                pool_name=pool_name,
+                connection_state=ConnectionState.USED,
+                counter=1,
+            )
+
         # We now perform the connection check outside of the lock.
         try:
             await self.ensure_connection(connection)
@@ -1398,6 +1480,8 @@ def get_encoder(self):
 
     def make_connection(self):
         """Create a new connection.  Can be overridden by child classes."""
+        # Note: We don't record IDLE here because async uses a sync make_connection
+        # but async record_connection_count. The recording is handled in get_connection.
         return self.connection_class(**self.connection_kwargs)
 
     async def ensure_connection(self, connection: AbstractConnection):
@@ -1421,6 +1505,7 @@ async def release(self, connection: AbstractConnection):
         # Connections should always be returned to the correct pool,
         # not doing so is an error that will cause an exception here.
         self._in_use_connections.remove(connection)
+
         if connection.should_reconnect():
             await connection.disconnect()
 
@@ -1429,6 +1514,19 @@ async def release(self, connection: AbstractConnection):
             AsyncAfterConnectionReleasedEvent(connection)
         )
 
+        # Record state transition: USED -> IDLE
+        pool_name = get_pool_name(self)
+        await record_connection_count(
+            pool_name=pool_name,
+            connection_state=ConnectionState.USED,
+            counter=-1,
+        )
+        await record_connection_count(
+            pool_name=pool_name,
+            connection_state=ConnectionState.IDLE,
+            counter=1,
+        )
+
     async def disconnect(self, inuse_connections: bool = True):
         """
         Disconnects connections in the pool
@@ -1447,6 +1545,7 @@ async def disconnect(self, inuse_connections: bool = True):
             *(connection.disconnect() for connection in connections),
             return_exceptions=True,
         )
+
         exc = next((r for r in resp if isinstance(r, BaseException)), None)
         if exc:
             raise exc

redis/asyncio/observability/recorder.py

@@ -24,13 +24,14 @@
 from typing import TYPE_CHECKING, List, Optional
 
 from redis.observability.attributes import (
+    ConnectionState,
     GeoFailoverReason,
     PubSubDirection,
 )
 from redis.observability.metrics import CloseReason, RedisMetricsCollector
 from redis.observability.providers import get_observability_instance
 from redis.observability.registry import get_observables_registry_instance
-from redis.utils import str_if_bytes
+from redis.utils import deprecated_function, str_if_bytes
 
 if TYPE_CHECKING:
     from redis.asyncio.connection import ConnectionPool
@@ -168,6 +169,38 @@ async def record_connection_create_time(
         pass
 
 
+async def record_connection_count(
+    pool_name: str,
+    connection_state: ConnectionState,
+    counter: int = 1,
+) -> None:
+    """
+    Record a connection count change for a single state.
+
+    Args:
+        pool_name: Connection pool identifier
+        connection_state: State to update (IDLE or USED)
+        counter: Number to add (positive) or subtract (negative)
+    """
+    collector = _get_or_create_collector()
+    if collector is None:
+        return
+
+    try:
+        collector.record_connection_count(
+            pool_name=pool_name,
+            connection_state=connection_state,
+            counter=counter,
+        )
+    except Exception:
+        pass
+
+
+@deprecated_function(
+    reason="Connection count is now tracked via record_connection_count(). "
+    "This functionality will be removed in the next major version",
+    version="7.4.0",
+)
 async def init_connection_count() -> None:
     """
     Initialize observable gauge for connection count metric.
@@ -194,6 +227,11 @@ def observable_callback(__):
         pass
 
 
+@deprecated_function(
+    reason="Connection count is now tracked via record_connection_count(). "
+    "This functionality will be removed in the next major version",
+    version="7.4.0",
+)
 async def register_pools_connection_count(
     connection_pools: List["ConnectionPool"],
 ) -> None:
@@ -301,7 +339,7 @@ async def record_connection_relaxed_timeout(
     Record a connection timeout relaxation event.
 
     Args:
-        connection_name: Connection identifier
+        connection_name: Connection identifier (pool name)
         maint_notification: Maintenance notification type
         relaxed: True to count up (relaxed), False to count down (unrelaxed)
     """