Use stm-containers to avoid contention in Vector

.gitignore

@@ -4,4 +4,6 @@ cabal.sandbox.config
 cabal.config
 .stack-work
 */*.yaml.lock
-.devcontainer
+.devcontainer
+dist-newstyle/
+stack.yaml.lock

prometheus-client/prometheus-client.cabal

@@ -49,6 +49,10 @@ library
     , utf8-string
     , exceptions
     , text
+    , stm-containers
+    , focus
+    , list-t
+    , hashable
     , data-sketches      < 0.4
   ghc-options: -Wall
 
@@ -62,12 +66,17 @@ test-suite doctest
       base               >=4.7 && <5
     , doctest
     , prometheus-client
+    , list-t
+    , hashable
+    , stm-containers
 
 test-suite spec
   type:              exitcode-stdio-1.0
   default-language:  Haskell2010
   hs-source-dirs:    src, tests
   main-is:           Spec.hs
+  build-tool-depends:
+    hspec-discover:hspec-discover
   build-depends:
       QuickCheck
     , atomic-primops
@@ -86,6 +95,10 @@ test-suite spec
     , exceptions
     , text
     , primitive
+    , focus
+    , list-t
+    , hashable
+    , stm-containers
     , data-sketches      < 0.4
   ghc-options: -Wall
 

prometheus-client/src/Prometheus/Label.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
 
 module Prometheus.Label (
     Label (..)
@@ -17,14 +18,23 @@ module Prometheus.Label (
 ) where
 
 import Data.Text
+import Data.Hashable
+
+instance (Hashable a, Hashable b, Hashable c, Hashable d, Hashable e, Hashable f, Hashable g, Hashable h) => Hashable (a, b, c, d, e, f, g, h) where
+    hashWithSalt s (a, b, c, d, e, f, g, h) =
+        s `hashWithSalt` a `hashWithSalt` b `hashWithSalt` c `hashWithSalt` d `hashWithSalt` e `hashWithSalt` f `hashWithSalt` g `hashWithSalt` h
+
+instance (Hashable a, Hashable b, Hashable c, Hashable d, Hashable e, Hashable f, Hashable g, Hashable h, Hashable i) => Hashable (a, b, c, d, e, f, g, h, i) where
+    hashWithSalt s (a, b, c, d, e, f, g, h, i) =
+        s `hashWithSalt` a `hashWithSalt` b `hashWithSalt` c `hashWithSalt` d `hashWithSalt` e `hashWithSalt` f `hashWithSalt` g `hashWithSalt` h `hashWithSalt` i
 
 -- | A list of tuples where the first value is the label and the second is the
 -- value of that label.
 type LabelPairs = [(Text, Text)]
 
 -- | Label describes a class of types that can be used to as the label of
 -- a vector.
-class Ord l => Label l where
+class Hashable l => Label l where
     labelPairs :: l -> l -> LabelPairs
 
 type Label0 = ()

prometheus-client/src/Prometheus/Metric/Summary.hs

@@ -79,7 +79,7 @@ summary info quantiles_ = Metric $ do
 
 instance Observer Summary where
     -- | Adds a new observation to a summary metric.
-    observe s v = doIO $ withMVar (reqSketch s) (`ReqSketch.insert` v)
+    observe s v = doIO $ withMVarMasked (reqSketch s) (\rs -> rs `ReqSketch.insert` v)
 
 -- | Retrieves a list of tuples containing a quantile and its associated value.
 getSummary :: MonadIO m => Summary -> m [(Rational, Double)]

prometheus-client/src/Prometheus/Metric/Vector.hs

@@ -13,25 +13,31 @@ import Prometheus.MonadMonitor
 
 import Control.Applicative ((<$>))
 import Control.DeepSeq
-import qualified Data.Atomics as Atomics
-import qualified Data.IORef as IORef
-import qualified Data.Map.Strict as Map
 import qualified Data.Text as T
 import Data.Traversable (forM)
+import System.IO.Unsafe (unsafeInterleaveIO)
+
+import Control.Concurrent.STM (atomically)
+import Control.Monad.Trans (lift)
+import qualified StmContainers.Map as Map
+import qualified Focus
+import qualified ListT
+import Data.Hashable
 
 
 type VectorState l m = (Metric m, Map.Map l (m, IO [SampleGroup]))
 
-data Vector l m = MkVector (IORef.IORef (VectorState l m))
+data Vector l m = MkVector (VectorState l m)
 
 instance NFData (Vector l m) where
-  rnf (MkVector ioref) = seq ioref ()
+  rnf (MkVector (gen, ioref)) = seq ioref `seq` seq gen ()
 
 -- | Creates a new vector of metrics given a label.
 vector :: Label l => l -> Metric m -> Metric (Vector l m)
 vector labels gen = Metric $ do
-    ioref <- checkLabelKeys labels $ IORef.newIORef (gen, Map.empty)
-    return (MkVector ioref, collectVector labels ioref)
+    m <- Map.newIO
+    vec <- checkLabelKeys labels $ pure $ (gen, m)
+    return (MkVector vec, collectVector labels vec)
 
 checkLabelKeys :: Label l => l -> a -> a
 checkLabelKeys keys r = foldl check r $ map (T.unpack . fst) $ labelPairs keys keys
@@ -56,10 +62,10 @@ checkLabelKeys keys r = foldl check r $ map (T.unpack . fst) $ labelPairs keys k
 -- TODO(will): This currently makes the assumption that all the types and info
 -- for all sample groups returned by a metric's collect method will be the same.
 -- It is not clear that this will always be a valid assumption.
-collectVector :: Label l => l -> IORef.IORef (VectorState l m) -> IO [SampleGroup]
-collectVector keys ioref = do
-    (_, metricMap) <- IORef.readIORef ioref
-    joinSamples <$> concat <$> mapM collectInner (Map.assocs metricMap)
+collectVector :: Label l => l -> VectorState l m -> IO [SampleGroup]
+collectVector keys (_, metricMap) = do
+    assocs <- ListT.toList $ Map.listTNonAtomic metricMap
+    joinSamples <$> concat <$> mapM collectInner assocs
     where
         collectInner (labels, (_metric, sampleGroups)) =
             map (adjustSamples labels) <$> sampleGroups
@@ -79,38 +85,45 @@ collectVector keys ioref = do
 getVectorWith :: Vector label metric
               -> (metric -> IO a)
               -> IO [(label, a)]
-getVectorWith (MkVector valueTVar) f = do
-    (_, metricMap) <- IORef.readIORef valueTVar
-    Map.assocs <$> forM metricMap (f . fst)
+getVectorWith (MkVector (_, metricMap)) f = do
+    ListT.toList $ do
+        (l, (m, _collect)) <- Map.listTNonAtomic metricMap
+        a <- lift $ f m
+        pure (l, a)
+
 
 -- | Given a label, applies an operation to the corresponding metric in the
 -- vector.
-withLabel :: (Label label, MonadMonitor m)
+withLabel :: (Hashable label, Label label, MonadMonitor m)
           => Vector label metric
           -> label
           -> (metric -> IO ())
           -> m ()
-withLabel (MkVector ioref) label f = doIO $ do
-    (Metric gen, _) <- IORef.readIORef ioref
-    newMetric <- gen
-    metric <- Atomics.atomicModifyIORefCAS ioref $ \(_, metricMap) ->
-        let maybeMetric = Map.lookup label metricMap
-            updatedMap  = Map.insert label newMetric metricMap
-        in  case maybeMetric of
-                Nothing     -> ((Metric gen, updatedMap), newMetric)
-                Just metric -> ((Metric gen, metricMap), metric)
+withLabel (MkVector (gen, metricMap)) label f = doIO $ do
+    newMetric <- unsafeInterleaveIO $ construct gen
+    metric <-
+        atomically $
+            Map.focus
+                (Focus.alter (\mmetric ->
+                    case mmetric of
+                        Nothing ->
+                            Just newMetric
+                        Just metric ->
+                            Just metric)
+                *> Focus.lookupWithDefault newMetric)
+                label
+                metricMap
+
     f (fst metric)
 
 -- | Removes a label from a vector.
-removeLabel :: (Label label, MonadMonitor m)
+removeLabel :: (Hashable label, Label label, MonadMonitor m)
             => Vector label metric -> label -> m ()
-removeLabel (MkVector valueTVar) label =
-    doIO $ Atomics.atomicModifyIORefCAS_ valueTVar f
-    where f (desc, metricMap) = (desc, Map.delete label metricMap)
+removeLabel (MkVector (_, metricMap)) label =
+    doIO $ atomically $ Map.delete label metricMap
 
 -- | Removes all labels from a vector.
 clearLabels :: (Label label, MonadMonitor m)
             => Vector label metric -> m ()
-clearLabels (MkVector valueTVar) =
-    doIO $ Atomics.atomicModifyIORefCAS_ valueTVar f
-    where f (desc, _) = (desc, Map.empty)
+clearLabels (MkVector (_, metricMap)) =
+    doIO $ atomically $ Map.reset metricMap

stack.yaml

@@ -5,7 +5,5 @@ packages:
 - prometheus-metrics-ghc
 - prometheus-proc
 - wai-middleware-prometheus
-extra-deps:
-- unix-memory-0.1.2
-- data-sketches-0.3.0.0
-resolver: lts-18.5
+
+resolver: lts-23.25