{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Per-job concurrency limits. A payload's 'concurrencyFor' describes how to
-- select a policy and key. Static inspection finds all policies that the migration
-- must initialize. Each selected policy supplies the limit.
module Arbiter.Core.Concurrency.Spec
  ( -- * Core types
    ConcurrencyKey (..)
  , toConcurrencyKey
  , concurrencyKeyText
  , ConcurrencyPolicy (..)
  , concurrencyPolicy

    -- * Selecting a policy per job
  , HasConcurrency (..)
  , ConcurrencyFor
  , noConcurrency
  , concurrencyBy
  , globalConcurrency
  , chooseWhen
  , concurrencyByCase
  , collectPolicies

    -- * Registry reflection
  , RegistryConcurrencyPolicies
  , registryConcurrencyPolicies
  , registryConcurrencyTables
  ) where

import Data.Aeson (FromJSON (..), ToJSON (..))
import Data.Int (Int32)
import Data.Set (Set)
import Data.Text (Text)

import Arbiter.Core.Admission
  ( AdmissionPolicy (..)
  , CollectFor (..)
  , RegistryPolicies (..)
  , prefixedKeyParseJSON
  , prefixedKeyText
  , prefixedKeyToJSON
  , registryPolicies
  , registryPolicyTables
  , selectBy
  , selectNone
  )
import Arbiter.Core.Selector (Selector, chooseWhen, collectPolicies, selectByCase)

-- | A resolved concurrency key with a policy prefix and per-key suffix. The
-- stored form is @prefix:suffix@.
data ConcurrencyKey = ConcurrencyKey
  { ConcurrencyKey -> Text
ckPrefix :: Text
  -- ^ The policy prefix.
  , ConcurrencyKey -> Text
ckSuffix :: Text
  -- ^ The per-key suffix, such as a tenant id.
  }
  deriving stock (ConcurrencyKey -> ConcurrencyKey -> Bool
(ConcurrencyKey -> ConcurrencyKey -> Bool)
-> (ConcurrencyKey -> ConcurrencyKey -> Bool) -> Eq ConcurrencyKey
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ConcurrencyKey -> ConcurrencyKey -> Bool
== :: ConcurrencyKey -> ConcurrencyKey -> Bool
$c/= :: ConcurrencyKey -> ConcurrencyKey -> Bool
/= :: ConcurrencyKey -> ConcurrencyKey -> Bool
Eq, Int -> ConcurrencyKey -> ShowS
[ConcurrencyKey] -> ShowS
ConcurrencyKey -> String
(Int -> ConcurrencyKey -> ShowS)
-> (ConcurrencyKey -> String)
-> ([ConcurrencyKey] -> ShowS)
-> Show ConcurrencyKey
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConcurrencyKey -> ShowS
showsPrec :: Int -> ConcurrencyKey -> ShowS
$cshow :: ConcurrencyKey -> String
show :: ConcurrencyKey -> String
$cshowList :: [ConcurrencyKey] -> ShowS
showList :: [ConcurrencyKey] -> ShowS
Show)

-- | A concurrency key from its prefix and suffix.
toConcurrencyKey :: Text -> Text -> ConcurrencyKey
toConcurrencyKey :: Text -> Text -> ConcurrencyKey
toConcurrencyKey Text
prefix Text
suffix = ConcurrencyKey {ckPrefix :: Text
ckPrefix = Text
prefix, ckSuffix :: Text
ckSuffix = Text
suffix}

instance ToJSON ConcurrencyKey where
  toJSON :: ConcurrencyKey -> Value
toJSON (ConcurrencyKey Text
prefix Text
suffix) = Text -> Text -> Value
prefixedKeyToJSON Text
prefix Text
suffix

instance FromJSON ConcurrencyKey where
  parseJSON :: Value -> Parser ConcurrencyKey
parseJSON = String
-> (Text -> Text -> ConcurrencyKey)
-> Value
-> Parser ConcurrencyKey
forall a. String -> (Text -> Text -> a) -> Value -> Parser a
prefixedKeyParseJSON String
"ConcurrencyKey" Text -> Text -> ConcurrencyKey
toConcurrencyKey

-- | The stored key text, @prefix:suffix@.
concurrencyKeyText :: ConcurrencyKey -> Text
concurrencyKeyText :: ConcurrencyKey -> Text
concurrencyKeyText (ConcurrencyKey Text
prefix Text
suffix) = Text -> Text -> Text
prefixedKeyText Text
prefix Text
suffix

-- | A concurrency policy. At most @cpLimit@ jobs are in flight per key under @cpPrefix@. The
-- default is seeded. An operator override on the policy takes precedence.
data ConcurrencyPolicy = ConcurrencyPolicy
  { ConcurrencyPolicy -> Text
cpPrefix :: Text
  -- ^ The key prefix. It must not contain @:@.
  , ConcurrencyPolicy -> Int32
cpLimit :: Int32
  -- ^ The default in-flight limit per key, in jobs. The migration rejects a
  -- declared limit below 1.
  }
  deriving stock (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
(ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> Eq ConcurrencyPolicy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
== :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
$c/= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
/= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
Eq, Eq ConcurrencyPolicy
Eq ConcurrencyPolicy =>
(ConcurrencyPolicy -> ConcurrencyPolicy -> Ordering)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> Bool)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy)
-> (ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy)
-> Ord ConcurrencyPolicy
ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
ConcurrencyPolicy -> ConcurrencyPolicy -> Ordering
ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: ConcurrencyPolicy -> ConcurrencyPolicy -> Ordering
compare :: ConcurrencyPolicy -> ConcurrencyPolicy -> Ordering
$c< :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
< :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
$c<= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
<= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
$c> :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
> :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
$c>= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
>= :: ConcurrencyPolicy -> ConcurrencyPolicy -> Bool
$cmax :: ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy
max :: ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy
$cmin :: ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy
min :: ConcurrencyPolicy -> ConcurrencyPolicy -> ConcurrencyPolicy
Ord, Int -> ConcurrencyPolicy -> ShowS
[ConcurrencyPolicy] -> ShowS
ConcurrencyPolicy -> String
(Int -> ConcurrencyPolicy -> ShowS)
-> (ConcurrencyPolicy -> String)
-> ([ConcurrencyPolicy] -> ShowS)
-> Show ConcurrencyPolicy
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConcurrencyPolicy -> ShowS
showsPrec :: Int -> ConcurrencyPolicy -> ShowS
$cshow :: ConcurrencyPolicy -> String
show :: ConcurrencyPolicy -> String
$cshowList :: [ConcurrencyPolicy] -> ShowS
showList :: [ConcurrencyPolicy] -> ShowS
Show)

instance AdmissionPolicy ConcurrencyPolicy where
  policyPrefixOf :: ConcurrencyPolicy -> Text
policyPrefixOf = ConcurrencyPolicy -> Text
cpPrefix

-- | A policy named @prefix@ admitting at most @limit@ concurrent jobs per key. The cap is
-- floored at 1. Pause it with @'Arbiter.Concurrency.setConcurrencyLimit' policy {cpLimit = 0}@.
-- The prefix must not contain @:@, the key separator. The migration enforces this.
concurrencyPolicy :: Text -> Int32 -> ConcurrencyPolicy
concurrencyPolicy :: Text -> Int32 -> ConcurrencyPolicy
concurrencyPolicy Text
prefix Int32
limit = ConcurrencyPolicy {cpPrefix :: Text
cpPrefix = Text
prefix, cpLimit :: Int32
cpLimit = Int32 -> Int32 -> Int32
forall a. Ord a => a -> a -> a
max Int32
minConcurrencyLimit Int32
limit}

-- | The lowest per-key limit a policy takes.
minConcurrencyLimit :: Int32
minConcurrencyLimit :: Int32
minConcurrencyLimit = Int32
1

-- | A selective description of the concurrency key for a payload. Evaluation
-- returns the job key. Static inspection returns the reachable policies.
type ConcurrencyFor payload = Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey)

-- | This payload is unbounded.
noConcurrency :: ConcurrencyFor payload
noConcurrency :: forall payload. ConcurrencyFor payload
noConcurrency = Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey)
forall p payload key. Selector p payload (Maybe key)
selectNone

-- | Cap by a fixed policy, keyed by a per-job suffix, such as a tenant id.
concurrencyBy :: ConcurrencyPolicy -> (payload -> Text) -> ConcurrencyFor payload
concurrencyBy :: forall payload.
ConcurrencyPolicy -> (payload -> Text) -> ConcurrencyFor payload
concurrencyBy = (Text -> Text -> ConcurrencyKey)
-> ConcurrencyPolicy
-> (payload -> Text)
-> Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey)
forall p key payload.
AdmissionPolicy p =>
(Text -> Text -> key)
-> p -> (payload -> Text) -> Selector p payload (Maybe key)
selectBy Text -> Text -> ConcurrencyKey
toConcurrencyKey

-- | Cap by a fixed policy under one shared key.
globalConcurrency :: ConcurrencyPolicy -> Text -> ConcurrencyFor payload
globalConcurrency :: forall payload. ConcurrencyPolicy -> Text -> ConcurrencyFor payload
globalConcurrency ConcurrencyPolicy
pol Text
suffix = ConcurrencyPolicy -> (payload -> Text) -> ConcurrencyFor payload
forall payload.
ConcurrencyPolicy -> (payload -> Text) -> ConcurrencyFor payload
concurrencyBy ConcurrencyPolicy
pol (Text -> payload -> Text
forall a b. a -> b -> a
const Text
suffix)

-- | N-way 'chooseWhen'. Maps the job to a finite tag, then each tag to its selector.
-- Policy collection evaluates every tag in @[minBound..maxBound]@. The tag's
-- 'Bounded'\/'Enum' and the selector must be total over @k@.
concurrencyByCase
  :: (Bounded k, Enum k, Eq k) => (payload -> k) -> (k -> ConcurrencyFor payload) -> ConcurrencyFor payload
concurrencyByCase :: forall k payload.
(Bounded k, Enum k, Eq k) =>
(payload -> k)
-> (k -> ConcurrencyFor payload) -> ConcurrencyFor payload
concurrencyByCase = (payload -> k)
-> (k -> Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey))
-> Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey)
forall k payload policy a.
(Bounded k, Enum k, Eq k) =>
(payload -> k)
-> (k -> Selector policy payload a) -> Selector policy payload a
selectByCase

-- | A payload's per-job policy selection. Defaults to unbounded. Only capped
-- payloads need an instance.
class HasConcurrency payload where
  -- | The selector deciding which policy (if any) caps a given job.
  concurrencyFor :: ConcurrencyFor payload
  concurrencyFor = ConcurrencyFor payload
forall payload. ConcurrencyFor payload
noConcurrency

instance {-# OVERLAPPABLE #-} HasConcurrency payload

instance (HasConcurrency payload) => CollectFor payload ConcurrencyPolicy where
  collectFor :: Set ConcurrencyPolicy
collectFor = Selector ConcurrencyPolicy payload (Maybe ConcurrencyKey)
-> Set ConcurrencyPolicy
forall policy payload a.
Ord policy =>
Selector policy payload a -> Set policy
collectPolicies (forall payload. HasConcurrency payload => ConcurrencyFor payload
concurrencyFor @payload)

-- | The constraint that lets the migration collect a registry's declared policies.
type RegistryConcurrencyPolicies registry = RegistryPolicies registry ConcurrencyPolicy

-- | Every distinct policy declared across the registry's payloads.
registryConcurrencyPolicies :: forall registry. (RegistryConcurrencyPolicies registry) => Set ConcurrencyPolicy
registryConcurrencyPolicies :: forall (registry :: JobPayloadRegistry).
RegistryConcurrencyPolicies registry =>
Set ConcurrencyPolicy
registryConcurrencyPolicies = forall (registry :: JobPayloadRegistry) p.
(Ord p, RegistryPolicies registry p) =>
Set p
registryPolicies @registry @ConcurrencyPolicy

-- | Each registry table paired with whether its payload declares any policy.
registryConcurrencyTables :: forall registry. (RegistryConcurrencyPolicies registry) => [(Text, Bool)]
registryConcurrencyTables :: forall (registry :: JobPayloadRegistry).
RegistryConcurrencyPolicies registry =>
[(Text, Bool)]
registryConcurrencyTables = forall (registry :: JobPayloadRegistry) p.
RegistryPolicies registry p =>
[(Text, Bool)]
registryPolicyTables @registry @ConcurrencyPolicy