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

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Per-job token-bucket rate limits. A payload's 'rateLimitFor' describes how
-- to select its policy and key. Static inspection finds all policies that the
-- migration must initialize.
module Arbiter.Core.RateLimit.Spec
  ( -- * Core types
    Durability (..)
  , RateLimitKey (..)
  , toRateLimitKey
  , rateLimitKeyText
  , RateLimitPolicy (..)
  , tokenBucket

    -- * Selecting a policy per job
  , HasRateLimit (..)
  , RateLimitFor
  , noLimit
  , limitBy
  , globalLimit
  , chooseWhen
  , limitByCase
  , collectPolicies

    -- * Registry reflection
  , RegistryRateLimitPolicies
  , registryRateLimitPolicies
  , registryRateLimitTables
  ) where

import Data.Aeson (FromJSON (..), ToJSON (..))
import Data.Set (Set)
import Data.Text (Text)
import Data.Time (NominalDiffTime)

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

-- | Whether the rate-limit bucket table is WAL-logged. Set by the migration config.
data Durability
  = -- | WAL-logged. Buckets survive a crash.
    Durable
  | -- | Not WAL-logged. A crash resets the buckets. The migrations use this by default.
    Unlogged
  deriving stock (Durability -> Durability -> Bool
(Durability -> Durability -> Bool)
-> (Durability -> Durability -> Bool) -> Eq Durability
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Durability -> Durability -> Bool
== :: Durability -> Durability -> Bool
$c/= :: Durability -> Durability -> Bool
/= :: Durability -> Durability -> Bool
Eq, Int -> Durability -> ShowS
[Durability] -> ShowS
Durability -> String
(Int -> Durability -> ShowS)
-> (Durability -> String)
-> ([Durability] -> ShowS)
-> Show Durability
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Durability -> ShowS
showsPrec :: Int -> Durability -> ShowS
$cshow :: Durability -> String
show :: Durability -> String
$cshowList :: [Durability] -> ShowS
showList :: [Durability] -> ShowS
Show)

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

-- | A rate-limit key from its prefix and suffix.
toRateLimitKey :: Text -> Text -> RateLimitKey
toRateLimitKey :: Text -> Text -> RateLimitKey
toRateLimitKey Text
prefix Text
suffix = RateLimitKey {rlkPrefix :: Text
rlkPrefix = Text
prefix, rlkSuffix :: Text
rlkSuffix = Text
suffix}

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

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

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

-- | A token-bucket policy. Burst @policyMax@, refilling continuously at @policyRefill@
-- per @policyInterval@. A @policyRefill@ of 0 is a manually-refilled bucket. Fields are
-- non-negative and the interval is positive. The migration rejects other values.
data RateLimitPolicy = RateLimitPolicy
  { RateLimitPolicy -> Text
policyPrefix :: Text
  -- ^ The key prefix. It must not contain @:@.
  , RateLimitPolicy -> Double
policyMax :: Double
  -- ^ The bucket capacity, in tokens.
  , RateLimitPolicy -> Double
policyRefill :: Double
  -- ^ The tokens added per interval.
  , RateLimitPolicy -> NominalDiffTime
policyInterval :: NominalDiffTime
  -- ^ The refill interval.
  }
  deriving stock (RateLimitPolicy -> RateLimitPolicy -> Bool
(RateLimitPolicy -> RateLimitPolicy -> Bool)
-> (RateLimitPolicy -> RateLimitPolicy -> Bool)
-> Eq RateLimitPolicy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RateLimitPolicy -> RateLimitPolicy -> Bool
== :: RateLimitPolicy -> RateLimitPolicy -> Bool
$c/= :: RateLimitPolicy -> RateLimitPolicy -> Bool
/= :: RateLimitPolicy -> RateLimitPolicy -> Bool
Eq, Eq RateLimitPolicy
Eq RateLimitPolicy =>
(RateLimitPolicy -> RateLimitPolicy -> Ordering)
-> (RateLimitPolicy -> RateLimitPolicy -> Bool)
-> (RateLimitPolicy -> RateLimitPolicy -> Bool)
-> (RateLimitPolicy -> RateLimitPolicy -> Bool)
-> (RateLimitPolicy -> RateLimitPolicy -> Bool)
-> (RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy)
-> (RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy)
-> Ord RateLimitPolicy
RateLimitPolicy -> RateLimitPolicy -> Bool
RateLimitPolicy -> RateLimitPolicy -> Ordering
RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy
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 :: RateLimitPolicy -> RateLimitPolicy -> Ordering
compare :: RateLimitPolicy -> RateLimitPolicy -> Ordering
$c< :: RateLimitPolicy -> RateLimitPolicy -> Bool
< :: RateLimitPolicy -> RateLimitPolicy -> Bool
$c<= :: RateLimitPolicy -> RateLimitPolicy -> Bool
<= :: RateLimitPolicy -> RateLimitPolicy -> Bool
$c> :: RateLimitPolicy -> RateLimitPolicy -> Bool
> :: RateLimitPolicy -> RateLimitPolicy -> Bool
$c>= :: RateLimitPolicy -> RateLimitPolicy -> Bool
>= :: RateLimitPolicy -> RateLimitPolicy -> Bool
$cmax :: RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy
max :: RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy
$cmin :: RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy
min :: RateLimitPolicy -> RateLimitPolicy -> RateLimitPolicy
Ord, Int -> RateLimitPolicy -> ShowS
[RateLimitPolicy] -> ShowS
RateLimitPolicy -> String
(Int -> RateLimitPolicy -> ShowS)
-> (RateLimitPolicy -> String)
-> ([RateLimitPolicy] -> ShowS)
-> Show RateLimitPolicy
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RateLimitPolicy -> ShowS
showsPrec :: Int -> RateLimitPolicy -> ShowS
$cshow :: RateLimitPolicy -> String
show :: RateLimitPolicy -> String
$cshowList :: [RateLimitPolicy] -> ShowS
showList :: [RateLimitPolicy] -> ShowS
Show)

instance AdmissionPolicy RateLimitPolicy where
  policyPrefixOf :: RateLimitPolicy -> Text
policyPrefixOf = RateLimitPolicy -> Text
policyPrefix

-- | @count@ per @period@, with a burst of @count@. The period is floored to a
-- tiny positive value. The prefix must not contain @:@, the key separator. The
-- migration enforces this.
tokenBucket :: Text -> Double -> NominalDiffTime -> RateLimitPolicy
tokenBucket :: Text -> Double -> NominalDiffTime -> RateLimitPolicy
tokenBucket Text
prefix Double
count NominalDiffTime
period =
  RateLimitPolicy
    { policyPrefix :: Text
policyPrefix = Text
prefix
    , policyMax :: Double
policyMax = Double
count
    , policyRefill :: Double
policyRefill = Double
count
    , policyInterval :: NominalDiffTime
policyInterval = NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Ord a => a -> a -> a
max NominalDiffTime
minRefillInterval NominalDiffTime
period
    }

-- | The shortest refill interval a token bucket takes.
minRefillInterval :: NominalDiffTime
minRefillInterval :: NominalDiffTime
minRefillInterval = NominalDiffTime
1e-6

-- | A selective description of the rate-limit key for a payload. Evaluation
-- returns the job key. Static inspection returns the reachable policies.
type RateLimitFor payload = Selector RateLimitPolicy payload (Maybe RateLimitKey)

-- | This payload is unlimited.
noLimit :: RateLimitFor payload
noLimit :: forall payload. RateLimitFor payload
noLimit = Selector RateLimitPolicy payload (Maybe RateLimitKey)
forall p payload key. Selector p payload (Maybe key)
selectNone

-- | Limit by a fixed policy, keyed by a per-job suffix, such as a tenant id.
limitBy :: RateLimitPolicy -> (payload -> Text) -> RateLimitFor payload
limitBy :: forall payload.
RateLimitPolicy -> (payload -> Text) -> RateLimitFor payload
limitBy = (Text -> Text -> RateLimitKey)
-> RateLimitPolicy
-> (payload -> Text)
-> Selector RateLimitPolicy payload (Maybe RateLimitKey)
forall p key payload.
AdmissionPolicy p =>
(Text -> Text -> key)
-> p -> (payload -> Text) -> Selector p payload (Maybe key)
selectBy Text -> Text -> RateLimitKey
toRateLimitKey

-- | Limit by a fixed policy under one shared key (a single global bucket).
globalLimit :: RateLimitPolicy -> Text -> RateLimitFor payload
globalLimit :: forall payload. RateLimitPolicy -> Text -> RateLimitFor payload
globalLimit RateLimitPolicy
pol Text
suffix = RateLimitPolicy -> (payload -> Text) -> RateLimitFor payload
forall payload.
RateLimitPolicy -> (payload -> Text) -> RateLimitFor payload
limitBy RateLimitPolicy
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@.
limitByCase :: (Bounded k, Enum k, Eq k) => (payload -> k) -> (k -> RateLimitFor payload) -> RateLimitFor payload
limitByCase :: forall k payload.
(Bounded k, Enum k, Eq k) =>
(payload -> k)
-> (k -> RateLimitFor payload) -> RateLimitFor payload
limitByCase = (payload -> k)
-> (k -> Selector RateLimitPolicy payload (Maybe RateLimitKey))
-> Selector RateLimitPolicy payload (Maybe RateLimitKey)
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 unlimited. Only limited
-- payloads need an instance.
class HasRateLimit payload where
  -- | The selector deciding which policy (if any) limits a given job.
  rateLimitFor :: RateLimitFor payload
  rateLimitFor = RateLimitFor payload
forall payload. RateLimitFor payload
noLimit

  -- | How many tokens this job spends, clamped to the bucket's range. Defaults to 1.
  rateLimitCost :: payload -> Double
  rateLimitCost payload
_ = Double
defaultRateLimitCost

instance {-# OVERLAPPABLE #-} HasRateLimit payload

-- | The tokens a job spends when its payload sets no cost.
defaultRateLimitCost :: Double
defaultRateLimitCost :: Double
defaultRateLimitCost = Double
1

instance (HasRateLimit payload) => CollectFor payload RateLimitPolicy where
  collectFor :: Set RateLimitPolicy
collectFor = Selector RateLimitPolicy payload (Maybe RateLimitKey)
-> Set RateLimitPolicy
forall policy payload a.
Ord policy =>
Selector policy payload a -> Set policy
collectPolicies (forall payload. HasRateLimit payload => RateLimitFor payload
rateLimitFor @payload)

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

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

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