module Prometheus.Metric.Vector (
    Vector (..)
,   vector
,   withLabel
,   removeLabel
,   clearLabels
,   getVectorWith
) where

import Prometheus.Label
import Prometheus.Metric
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)


type VectorState l m = (Metric m, Map.Map l (m, IO [SampleGroup]))

data Vector l m = MkVector (IORef.IORef (VectorState l m))

instance NFData (Vector l m) where
  rnf :: Vector l m -> ()
rnf (MkVector IORef (VectorState l m)
ioref) = IORef (VectorState l m) -> () -> ()
forall a b. a -> b -> b
seq IORef (VectorState l m)
ioref ()

-- | Creates a new vector of metrics given a label.
vector :: Label l => l -> Metric m -> Metric (Vector l m)
vector :: forall l m. Label l => l -> Metric m -> Metric (Vector l m)
vector l
labels Metric m
gen = IO (Vector l m, IO [SampleGroup]) -> Metric (Vector l m)
forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric (IO (Vector l m, IO [SampleGroup]) -> Metric (Vector l m))
-> IO (Vector l m, IO [SampleGroup]) -> Metric (Vector l m)
forall a b. (a -> b) -> a -> b
$ do
    IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref <- l
-> IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
-> IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
forall l a. Label l => l -> a -> a
checkLabelKeys l
labels (IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
 -> IO (IORef (Metric m, Map l (m, IO [SampleGroup]))))
-> IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
-> IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
forall a b. (a -> b) -> a -> b
$ (Metric m, Map l (m, IO [SampleGroup]))
-> IO (IORef (Metric m, Map l (m, IO [SampleGroup])))
forall a. a -> IO (IORef a)
IORef.newIORef (Metric m
gen, Map l (m, IO [SampleGroup])
forall k a. Map k a
Map.empty)
    (Vector l m, IO [SampleGroup]) -> IO (Vector l m, IO [SampleGroup])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IORef (Metric m, Map l (m, IO [SampleGroup])) -> Vector l m
forall l m. IORef (VectorState l m) -> Vector l m
MkVector IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref, l
-> IORef (Metric m, Map l (m, IO [SampleGroup]))
-> IO [SampleGroup]
forall l m.
Label l =>
l -> IORef (VectorState l m) -> IO [SampleGroup]
collectVector l
labels IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref)

checkLabelKeys :: Label l => l -> a -> a
checkLabelKeys :: forall l a. Label l => l -> a -> a
checkLabelKeys l
keys a
r = (a -> String -> a) -> a -> [String] -> a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl a -> String -> a
forall {a}. a -> String -> a
check a
r ([String] -> a) -> [String] -> a
forall a b. (a -> b) -> a -> b
$ ((Text, Text) -> String) -> [(Text, Text)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> String
T.unpack (Text -> String)
-> ((Text, Text) -> Text) -> (Text, Text) -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Text) -> Text
forall a b. (a, b) -> a
fst) ([(Text, Text)] -> [String]) -> [(Text, Text)] -> [String]
forall a b. (a -> b) -> a -> b
$ l -> l -> [(Text, Text)]
forall l. Label l => l -> l -> [(Text, Text)]
labelPairs l
keys l
keys
    where
        check :: a -> String -> a
check a
_ String
"instance" = String -> a
forall a. HasCallStack => String -> a
error String
"The label 'instance' is reserved."
        check a
_ String
"job"      = String -> a
forall a. HasCallStack => String -> a
error String
"The label 'job' is reserved."
        check a
_ String
"quantile" = String -> a
forall a. HasCallStack => String -> a
error String
"The label 'quantile' is reserved."
        check a
a (Char
k:String
ey)
            | Char -> Bool
validStart Char
k Bool -> Bool -> Bool
&& (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
validRest String
ey = a
a
            | Bool
otherwise = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"The label '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Char
kChar -> String -> String
forall a. a -> [a] -> [a]
:String
ey) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' is not valid."
        check a
_ []         = String -> a
forall a. HasCallStack => String -> a
error String
"Empty labels are not allowed."

        validStart :: Char -> Bool
validStart Char
c =  (Char
'a' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'z')
                     Bool -> Bool -> Bool
|| (Char
'A' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'Z')
                     Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_'

        validRest :: Char -> Bool
validRest Char
c =  (Char
'a' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'z')
                    Bool -> Bool -> Bool
|| (Char
'A' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'Z')
                    Bool -> Bool -> Bool
|| (Char
'0' Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<= Char
'9')
                    Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'_'

-- 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 :: forall l m.
Label l =>
l -> IORef (VectorState l m) -> IO [SampleGroup]
collectVector l
keys IORef (VectorState l m)
ioref = do
    (Metric m
_, Map l (m, IO [SampleGroup])
metricMap) <- IORef (VectorState l m) -> IO (VectorState l m)
forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState l m)
ioref
    [SampleGroup] -> [SampleGroup]
joinSamples ([SampleGroup] -> [SampleGroup])
-> ([[SampleGroup]] -> [SampleGroup])
-> [[SampleGroup]]
-> [SampleGroup]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[SampleGroup]] -> [SampleGroup]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SampleGroup]] -> [SampleGroup])
-> IO [[SampleGroup]] -> IO [SampleGroup]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((l, (m, IO [SampleGroup])) -> IO [SampleGroup])
-> [(l, (m, IO [SampleGroup]))] -> IO [[SampleGroup]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (l, (m, IO [SampleGroup])) -> IO [SampleGroup]
forall {f :: * -> *} {a}.
Functor f =>
(l, (a, f [SampleGroup])) -> f [SampleGroup]
collectInner (Map l (m, IO [SampleGroup]) -> [(l, (m, IO [SampleGroup]))]
forall k a. Map k a -> [(k, a)]
Map.assocs Map l (m, IO [SampleGroup])
metricMap)
    where
        collectInner :: (l, (a, f [SampleGroup])) -> f [SampleGroup]
collectInner (l
labels, (a
_metric, f [SampleGroup]
sampleGroups)) =
            (SampleGroup -> SampleGroup) -> [SampleGroup] -> [SampleGroup]
forall a b. (a -> b) -> [a] -> [b]
map (l -> SampleGroup -> SampleGroup
adjustSamples l
labels) ([SampleGroup] -> [SampleGroup])
-> f [SampleGroup] -> f [SampleGroup]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f [SampleGroup]
sampleGroups

        adjustSamples :: l -> SampleGroup -> SampleGroup
adjustSamples l
labels (SampleGroup Info
info SampleType
ty [Sample]
samples) =
            Info -> SampleType -> [Sample] -> SampleGroup
SampleGroup Info
info SampleType
ty ((Sample -> Sample) -> [Sample] -> [Sample]
forall a b. (a -> b) -> [a] -> [b]
map (l -> Sample -> Sample
prependLabels l
labels) [Sample]
samples)

        prependLabels :: l -> Sample -> Sample
prependLabels l
l (Sample Text
name [(Text, Text)]
labels ByteString
value) =
            Text -> [(Text, Text)] -> ByteString -> Sample
Sample Text
name (l -> l -> [(Text, Text)]
forall l. Label l => l -> l -> [(Text, Text)]
labelPairs l
keys l
l [(Text, Text)] -> [(Text, Text)] -> [(Text, Text)]
forall a. [a] -> [a] -> [a]
++ [(Text, Text)]
labels) ByteString
value

        joinSamples :: [SampleGroup] -> [SampleGroup]
joinSamples []                      = []
        joinSamples s :: [SampleGroup]
s@(SampleGroup Info
i SampleType
t [Sample]
_:[SampleGroup]
_) = [Info -> SampleType -> [Sample] -> SampleGroup
SampleGroup Info
i SampleType
t ([SampleGroup] -> [Sample]
extract [SampleGroup]
s)]

        extract :: [SampleGroup] -> [Sample]
extract [] = []
        extract (SampleGroup Info
_ SampleType
_ [Sample]
s:[SampleGroup]
xs) = [Sample]
s [Sample] -> [Sample] -> [Sample]
forall a. [a] -> [a] -> [a]
++ [SampleGroup] -> [Sample]
extract [SampleGroup]
xs

getVectorWith :: Vector label metric
              -> (metric -> IO a)
              -> IO [(label, a)]
getVectorWith :: forall label metric a.
Vector label metric -> (metric -> IO a) -> IO [(label, a)]
getVectorWith (MkVector IORef (VectorState label metric)
valueTVar) metric -> IO a
f = do
    (Metric metric
_, Map label (metric, IO [SampleGroup])
metricMap) <- IORef (VectorState label metric) -> IO (VectorState label metric)
forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState label metric)
valueTVar
    Map label a -> [(label, a)]
forall k a. Map k a -> [(k, a)]
Map.assocs (Map label a -> [(label, a)])
-> IO (Map label a) -> IO [(label, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map label (metric, IO [SampleGroup])
-> ((metric, IO [SampleGroup]) -> IO a) -> IO (Map label a)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Map label (metric, IO [SampleGroup])
metricMap (metric -> IO a
f (metric -> IO a)
-> ((metric, IO [SampleGroup]) -> metric)
-> (metric, IO [SampleGroup])
-> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (metric, IO [SampleGroup]) -> metric
forall a b. (a, b) -> a
fst)

-- | Given a label, applies an operation to the corresponding metric in the
-- vector.
withLabel :: (Label label, MonadMonitor m)
          => Vector label metric
          -> label
          -> (metric -> IO ())
          -> m ()
withLabel :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
withLabel (MkVector IORef (VectorState label metric)
ioref) label
label metric -> IO ()
f = IO () -> m ()
forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    (Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
_) <- IORef (VectorState label metric) -> IO (VectorState label metric)
forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState label metric)
ioref
    (metric, IO [SampleGroup])
newMetric <- IO (metric, IO [SampleGroup])
gen
    (metric, IO [SampleGroup])
metric <- IORef (VectorState label metric)
-> (VectorState label metric
    -> (VectorState label metric, (metric, IO [SampleGroup])))
-> IO (metric, IO [SampleGroup])
forall a b. IORef a -> (a -> (a, b)) -> IO b
Atomics.atomicModifyIORefCAS IORef (VectorState label metric)
ioref ((VectorState label metric
  -> (VectorState label metric, (metric, IO [SampleGroup])))
 -> IO (metric, IO [SampleGroup]))
-> (VectorState label metric
    -> (VectorState label metric, (metric, IO [SampleGroup])))
-> IO (metric, IO [SampleGroup])
forall a b. (a -> b) -> a -> b
$ \(Metric metric
_, Map label (metric, IO [SampleGroup])
metricMap) ->
        let maybeMetric :: Maybe (metric, IO [SampleGroup])
maybeMetric = label
-> Map label (metric, IO [SampleGroup])
-> Maybe (metric, IO [SampleGroup])
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup label
label Map label (metric, IO [SampleGroup])
metricMap
            updatedMap :: Map label (metric, IO [SampleGroup])
updatedMap  = label
-> (metric, IO [SampleGroup])
-> Map label (metric, IO [SampleGroup])
-> Map label (metric, IO [SampleGroup])
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert label
label (metric, IO [SampleGroup])
newMetric Map label (metric, IO [SampleGroup])
metricMap
        in  case Maybe (metric, IO [SampleGroup])
maybeMetric of
                Maybe (metric, IO [SampleGroup])
Nothing     -> ((IO (metric, IO [SampleGroup]) -> Metric metric
forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
updatedMap), (metric, IO [SampleGroup])
newMetric)
                Just (metric, IO [SampleGroup])
metric -> ((IO (metric, IO [SampleGroup]) -> Metric metric
forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
metricMap), (metric, IO [SampleGroup])
metric)
    metric -> IO ()
f ((metric, IO [SampleGroup]) -> metric
forall a b. (a, b) -> a
fst (metric, IO [SampleGroup])
metric)

-- | Removes a label from a vector.
removeLabel :: (Label label, MonadMonitor m)
            => Vector label metric -> label -> m ()
removeLabel :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> m ()
removeLabel (MkVector IORef (VectorState label metric)
valueTVar) label
label =
    IO () -> m ()
forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef (VectorState label metric)
-> (VectorState label metric -> VectorState label metric) -> IO ()
forall t. IORef t -> (t -> t) -> IO ()
Atomics.atomicModifyIORefCAS_ IORef (VectorState label metric)
valueTVar VectorState label metric -> VectorState label metric
forall {a} {a}. (a, Map label a) -> (a, Map label a)
f
    where f :: (a, Map label a) -> (a, Map label a)
f (a
desc, Map label a
metricMap) = (a
desc, label -> Map label a -> Map label a
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete label
label Map label a
metricMap)

-- | Removes all labels from a vector.
clearLabels :: (Label label, MonadMonitor m)
            => Vector label metric -> m ()
clearLabels :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> m ()
clearLabels (MkVector IORef (VectorState label metric)
valueTVar) =
    IO () -> m ()
forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef (VectorState label metric)
-> (VectorState label metric -> VectorState label metric) -> IO ()
forall t. IORef t -> (t -> t) -> IO ()
Atomics.atomicModifyIORefCAS_ IORef (VectorState label metric)
valueTVar VectorState label metric -> VectorState label metric
forall {a} {b} {k} {a}. (a, b) -> (a, Map k a)
f
    where f :: (a, b) -> (a, Map k a)
f (a
desc, b
_) = (a
desc, Map k a
forall k a. Map k a
Map.empty)