{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Shared machinery for per-job admission policies (rate limits and concurrency
-- pools). Both kinds pick a policy and a @prefix:suffix@ key per job via a
-- 'Selector', seed every policy a registry references, and store policies in a
-- default\/override table.
module Arbiter.Core.Admission
  ( -- * Keys
    prefixedKeyText
  , prefixedKeyToJSON
  , prefixedKeyParseJSON
  , splitPrefixedSuffix

    -- * Policies and selectors
  , AdmissionPolicy (..)
  , selectNone
  , selectBy

    -- * Registry reflection
  , CollectFor (..)
  , RegistryPolicies (..)
  , registryPolicies
  , registryPolicyTables

    -- * SQL fragments
  , effectivePolicyCol
  , excludedAssignment
  , policyUpsertSQL
  ) where

import Data.Aeson (Value, object, withObject, (.:), (.=))
import Data.Aeson.Types (Parser)
import Data.Proxy (Proxy (..))
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import GHC.TypeLits (KnownSymbol, symbolVal)
import NeatInterpolation (text)

import Arbiter.Core.QueueRegistry (JobPayloadRegistry, SpecName, SpecPayload)
import Arbiter.Core.Selector (Selector, field, usePolicy)

-- Keys -----------------------------------------------------------------------

-- | The stored key text, @prefix:suffix@.
prefixedKeyText :: Text -> Text -> Text
prefixedKeyText :: Text -> Text -> Text
prefixedKeyText Text
prefix Text
suffix = Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
suffix

-- | JSON for a @prefix:suffix@ key.
prefixedKeyToJSON :: Text -> Text -> Value
prefixedKeyToJSON :: Text -> Text -> Value
prefixedKeyToJSON Text
prefix Text
suffix = [Pair] -> Value
object [Key
"prefix" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
prefix, Key
"suffix" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
suffix]

-- | Parse a @prefix:suffix@ key into a constructor.
prefixedKeyParseJSON :: String -> (Text -> Text -> a) -> Value -> Parser a
prefixedKeyParseJSON :: forall a. String -> (Text -> Text -> a) -> Value -> Parser a
prefixedKeyParseJSON String
name Text -> Text -> a
mkKey = String -> (Object -> Parser a) -> Value -> Parser a
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
name ((Object -> Parser a) -> Value -> Parser a)
-> (Object -> Parser a) -> Value -> Parser a
forall a b. (a -> b) -> a -> b
$ \Object
obj -> Text -> Text -> a
mkKey (Text -> Text -> a) -> Parser Text -> Parser (Text -> a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
obj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"prefix" Parser (Text -> a) -> Parser Text -> Parser a
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
obj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"suffix"

-- | Recover the suffix of a stored @prefix:suffix@ key given its prefix.
splitPrefixedSuffix :: Text -> Text -> Text
splitPrefixedSuffix :: Text -> Text -> Text
splitPrefixedSuffix Text
prefix Text
key = Int -> Text -> Text
T.drop (Text -> Int
T.length Text
prefix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Text
key

-- Policies and selectors -----------------------------------------------------

-- | A policy that admits jobs under a named prefix.
class (Ord p) => AdmissionPolicy p where
  policyPrefixOf :: p -> Text

-- | This payload is unrestricted by this policy kind.
selectNone :: Selector p payload (Maybe key)
selectNone :: forall p payload key. Selector p payload (Maybe key)
selectNone = Maybe key -> Select (Prim p payload) (Maybe key)
forall a. a -> Select (Prim p payload) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe key
forall a. Maybe a
Nothing

-- | Restrict by a fixed policy, keyed by a per-job suffix (e.g. a tenant id).
selectBy
  :: (AdmissionPolicy p)
  => (Text -> Text -> key)
  -> p
  -> (payload -> Text)
  -> Selector p payload (Maybe key)
selectBy :: forall p key payload.
AdmissionPolicy p =>
(Text -> Text -> key)
-> p -> (payload -> Text) -> Selector p payload (Maybe key)
selectBy Text -> Text -> key
mkKey p
pol payload -> Text
suffix =
  (\p
policy Text
suffixText -> key -> Maybe key
forall a. a -> Maybe a
Just (Text -> Text -> key
mkKey (p -> Text
forall p. AdmissionPolicy p => p -> Text
policyPrefixOf p
policy) Text
suffixText)) (p -> Text -> Maybe key)
-> Select (Prim p payload) p
-> Select (Prim p payload) (Text -> Maybe key)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p -> Select (Prim p payload) p
forall policy payload. policy -> Selector policy payload policy
usePolicy p
pol Select (Prim p payload) (Text -> Maybe key)
-> Select (Prim p payload) Text
-> Select (Prim p payload) (Maybe key)
forall a b.
Select (Prim p payload) (a -> b)
-> Select (Prim p payload) a -> Select (Prim p payload) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (payload -> Text) -> Select (Prim p payload) Text
forall payload a policy.
(payload -> a) -> Selector policy payload a
field payload -> Text
suffix

-- SQL fragments ---------------------------------------------------------------

-- | The effective (override-or-default) column of the policy at @alias@.
effectivePolicyCol :: Text -> Text -> Text
effectivePolicyCol :: Text -> Text -> Text
effectivePolicyCol Text
alias Text
name = [text|COALESCE(${alias}.override_${name}, ${alias}.default_${name})|]

-- | An @ON CONFLICT DO UPDATE SET@ assignment copying a column from the excluded row.
excludedAssignment :: Text -> Text
excludedAssignment :: Text -> Text
excludedAssignment Text
column = [text|${column} = EXCLUDED.${column}|]

-- | Upsert a policy row's defaults keyed on prefix_id, preserving operator overrides.
policyUpsertSQL :: Text -> Text -> [(Text, Text)] -> Text
policyUpsertSQL :: Text -> Text -> [(Text, Text)] -> Text
policyUpsertSQL Text
policiesTable Text
prefixLit [(Text, Text)]
defaults =
  let names :: Text
names = Text -> [Text] -> Text
T.intercalate Text
", " (((Text, Text) -> Text) -> [(Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Text) -> Text
forall a b. (a, b) -> a
fst [(Text, Text)]
defaults)
      values :: Text
values = Text -> [Text] -> Text
T.intercalate Text
", " (((Text, Text) -> Text) -> [(Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Text) -> Text
forall a b. (a, b) -> b
snd [(Text, Text)]
defaults)
      setClause :: Text
setClause = Text -> [Text] -> Text
T.intercalate Text
", " (((Text, Text) -> Text) -> [(Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Text
excludedAssignment (Text -> Text) -> ((Text, Text) -> Text) -> (Text, Text) -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, Text) -> Text
forall a b. (a, b) -> a
fst) [(Text, Text)]
defaults)
   in [text|INSERT INTO ${policiesTable} (prefix_id, ${names}) VALUES (${prefixLit}, ${values}) ON CONFLICT (prefix_id) DO UPDATE SET ${setClause};|]

-- Registry reflection --------------------------------------------------------

-- | The policies of kind @p@ a single payload's selector can reach. Each feature
-- provides an instance from its 'Arbiter.Core.RateLimit.Spec.HasRateLimit' \/
-- 'Arbiter.Core.Concurrency.Spec.HasConcurrency' selector.
class CollectFor payload p where
  collectFor :: Set p

-- | Each registry table's policies of kind @p@, in registry order.
class RegistryPolicies (registry :: JobPayloadRegistry) p where
  registryTablePolicies :: [(Text, Set p)]

instance RegistryPolicies '[] p where
  registryTablePolicies :: [(Text, Set p)]
registryTablePolicies = []

instance
  (CollectFor (SpecPayload spec) p, KnownSymbol (SpecName spec), RegistryPolicies rest p)
  => RegistryPolicies (spec ': rest) p
  where
  registryTablePolicies :: [(Text, Set p)]
registryTablePolicies =
    (String -> Text
T.pack (Proxy (SpecName spec) -> String
forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
forall (t :: Symbol). Proxy t
Proxy @(SpecName spec))), forall {k} (payload :: k) p. CollectFor payload p => Set p
forall payload p. CollectFor payload p => Set p
collectFor @(SpecPayload spec) @p)
      (Text, Set p) -> [(Text, Set p)] -> [(Text, Set p)]
forall a. a -> [a] -> [a]
: forall (registry :: JobPayloadRegistry) p.
RegistryPolicies registry p =>
[(Text, Set p)]
registryTablePolicies @rest @p

-- | Every policy of kind @p@ declared across a registry's payloads.
registryPolicies :: forall registry p. (Ord p, RegistryPolicies registry p) => Set p
registryPolicies :: forall (registry :: JobPayloadRegistry) p.
(Ord p, RegistryPolicies registry p) =>
Set p
registryPolicies = [Set p] -> Set p
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (((Text, Set p) -> Set p) -> [(Text, Set p)] -> [Set p]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Set p) -> Set p
forall a b. (a, b) -> b
snd (forall (registry :: JobPayloadRegistry) p.
RegistryPolicies registry p =>
[(Text, Set p)]
registryTablePolicies @registry @p))

-- | Each registry table paired with whether its payload declares any policy of kind @p@.
registryPolicyTables :: forall registry p. (RegistryPolicies registry p) => [(Text, Bool)]
registryPolicyTables :: forall (registry :: JobPayloadRegistry) p.
RegistryPolicies registry p =>
[(Text, Bool)]
registryPolicyTables = ((Text, Set p) -> (Text, Bool))
-> [(Text, Set p)] -> [(Text, Bool)]
forall a b. (a -> b) -> [a] -> [b]
map ((Set p -> Bool) -> (Text, Set p) -> (Text, Bool)
forall a b. (a -> b) -> (Text, a) -> (Text, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Bool
not (Bool -> Bool) -> (Set p -> Bool) -> Set p -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set p -> Bool
forall a. Set a -> Bool
Set.null)) (forall (registry :: JobPayloadRegistry) p.
RegistryPolicies registry p =>
[(Text, Set p)]
registryTablePolicies @registry @p)