fix: implement connection tracking in metrics

Right now metrics observation handler does not track database connections but updates a single Gauge based on HasqlPoolObs events. This is problematic because Hasql pool reports various connection events in multiple phases. The connection state machine is not simple and to precisely report the number of connections in various states, it is necessary to track their lifecycles.

This change adds a ConnTrack data structure and logic to track database connections lifecycles. At the moment it supports "connected" and "inUse" connection counts precisely. The "pgrst_db_pool_available" metric is implemented on top of ConnTrack instead of a simple Gauge.

src/PostgREST/Metrics.hs

@@ -5,7 +5,10 @@ Description : Metrics based on the Observation module. See Observation.hs.
 -}
 module PostgREST.Metrics
   ( init
+  , ConnTrack
+  , ConnStats (..)
   , MetricsState (..)
+  , connectionCounts
   , observationMetrics
   , metricsToText
   ) where
@@ -19,13 +22,18 @@ import qualified Prometheus.Metric.GHC as PMG
 
 import PostgREST.Observation
 
-
-import Protolude
+import           Control.Arrow      ((&&&))
+import           Data.Bitraversable (bisequenceA)
+import           Data.Tuple.Extra   (both)
+import           Data.UUID          (UUID)
+import qualified Focus
+import           Protolude
+import qualified StmHamt.SizedHamt  as SH
 
 data MetricsState =
   MetricsState {
     poolTimeouts         :: Counter,
-    poolAvailable        :: Gauge,
+    connTrack            :: ConnTrack,
     poolWaiting          :: Gauge,
     poolMaxSize          :: Gauge,
     schemaCacheLoads     :: Vector Label1 Counter,
@@ -40,7 +48,7 @@ init configDbPoolSize = do
   whenM getRTSStatsEnabled $ void $ register PMG.ghcMetrics
   metricState <- MetricsState <$>
     register (counter (Info "pgrst_db_pool_timeouts_total" "The total number of pool connection timeouts")) <*>
-    register (gauge (Info "pgrst_db_pool_available" "Available connections in the pool")) <*>
+    register (Metric ((identity &&& dbPoolAvailable) <$> connectionTracker)) <*>
     register (gauge (Info "pgrst_db_pool_waiting" "Requests waiting to acquire a pool connection")) <*>
     register (gauge (Info "pgrst_db_pool_max" "Max pool connections")) <*>
     register (vector "status" $ counter (Info "pgrst_schema_cache_loads_total" "The total number of times the schema cache was loaded")) <*>
@@ -50,20 +58,28 @@ init configDbPoolSize = do
     register (counter (Info "pgrst_jwt_cache_evictions_total" "The total number of JWT cache evictions"))
   setGauge (poolMaxSize metricState) (fromIntegral configDbPoolSize)
   pure metricState
+  where
+    dbPoolAvailable = (pure . noLabelsGroup (Info "pgrst_db_pool_available" "Available connections in the pool") GaugeType . calcAvailable <$>) . connectionCounts
+      where
+        calcAvailable = liftA2 (-) connected inUse
+    toSample name labels = Sample name labels . encodeUtf8 . show
+    noLabelsGroup info sampleType = SampleGroup info sampleType . pure . toSample (metricName info) mempty
 
 -- Only some observations are used as metrics
 observationMetrics :: MetricsState -> ObservationHandler
 observationMetrics MetricsState{..} obs = case obs of
   PoolAcqTimeoutObs -> do
     incCounter poolTimeouts
-  (HasqlPoolObs (SQL.ConnectionObservation _ status)) -> case status of
-     SQL.ReadyForUseConnectionStatus _ -> do
-      incGauge poolAvailable
-     SQL.InUseConnectionStatus        -> do
-      decGauge poolAvailable
-     SQL.TerminatedConnectionStatus  _ -> do
-      decGauge poolAvailable
-     SQL.ConnectingConnectionStatus -> pure ()
+  -- Handle pool observations with connection tracking
+  -- this is necessary because it is not possible
+  -- to accurately maintain open/in use conneciton counts
+  -- statelessly based only on pool observation events.
+  -- The reason is that hasql-pool emits TerminatedConnectionStatus
+  -- both for connections successfully established and failed when connecting.
+  -- When receiving TerminatedConnectionStatus we have to find out
+  -- if we can decrement established connection count. To do that we have to track
+  -- established connections.
+  (HasqlPoolObs sqlObs) -> trackConnections connTrack sqlObs
   PoolRequest ->
     incGauge poolWaiting
   PoolRequestFullfilled ->
@@ -81,3 +97,28 @@ observationMetrics MetricsState{..} obs = case obs of
 
 metricsToText :: IO LBS.ByteString
 metricsToText = exportMetricsAsText
+
+data ConnStats = ConnStats {
+    connected :: Int,
+    inUse     :: Int
+} deriving (Eq, Show)
+
+data ConnTrack = ConnTrack { connTrackConnected :: SH.SizedHamt UUID, connTrackInUse :: SH.SizedHamt UUID }
+
+connectionTracker :: IO ConnTrack
+connectionTracker =  ConnTrack <$> SH.newIO <*> SH.newIO
+
+trackConnections :: ConnTrack -> SQL.Observation -> IO ()
+trackConnections ConnTrack{..} (SQL.ConnectionObservation uuid status) = case status of
+  SQL.ReadyForUseConnectionStatus _ -> atomically $
+    SH.insert identity uuid connTrackConnected *>
+    SH.focus Focus.delete identity uuid connTrackInUse
+  SQL.TerminatedConnectionStatus _ -> atomically $
+    SH.focus Focus.delete identity uuid connTrackConnected *>
+    SH.focus Focus.delete identity uuid connTrackInUse
+  SQL.InUseConnectionStatus -> atomically $
+    SH.insert identity uuid connTrackInUse
+  _ -> mempty
+
+connectionCounts :: ConnTrack -> IO ConnStats
+connectionCounts = atomically . fmap (uncurry ConnStats) . bisequenceA . both SH.size . (connTrackConnected &&& connTrackInUse)