{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}

-- | 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 (..)
  , rateLimitKeyText
  , Policy (..)
  , tokenBucket

    -- * Selecting a policy per job
  , HasRateLimit (..)
  , RateLimitFor
  , noLimit
  , limitBy
  , globalLimit
  , chooseWhen
  , limitByCase
  , runRateLimitFor
  , 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, runSelector, selectByCase)

-- | Whether the rate-limit bucket table is WAL-logged. Set at migration time.
data Durability = Durable | 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@. The separate prefix supports policy lookup.
data RateLimitKey = RateLimitKey
  { RateLimitKey -> Text
rlkPrefix :: Text
  , RateLimitKey -> Text
rlkSuffix :: Text
  }
  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)

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
RateLimitKey

-- | The bucket key as stored, prefix and suffix joined.
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 @policyRefill@ every
-- @policyInterval@. A @policyRefill@ of 0 is a manually-refilled bucket.
data Policy = Policy
  { Policy -> Text
policyPrefix :: Text
  , Policy -> Double
policyMax :: Double
  , Policy -> Double
policyRefill :: Double
  , Policy -> NominalDiffTime
policyInterval :: NominalDiffTime
  }
  deriving stock (Policy -> Policy -> Bool
(Policy -> Policy -> Bool)
-> (Policy -> Policy -> Bool) -> Eq Policy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Policy -> Policy -> Bool
== :: Policy -> Policy -> Bool
$c/= :: Policy -> Policy -> Bool
/= :: Policy -> Policy -> Bool
Eq, Eq Policy
Eq Policy =>
(Policy -> Policy -> Ordering)
-> (Policy -> Policy -> Bool)
-> (Policy -> Policy -> Bool)
-> (Policy -> Policy -> Bool)
-> (Policy -> Policy -> Bool)
-> (Policy -> Policy -> Policy)
-> (Policy -> Policy -> Policy)
-> Ord Policy
Policy -> Policy -> Bool
Policy -> Policy -> Ordering
Policy -> Policy -> Policy
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 :: Policy -> Policy -> Ordering
compare :: Policy -> Policy -> Ordering
$c< :: Policy -> Policy -> Bool
< :: Policy -> Policy -> Bool
$c<= :: Policy -> Policy -> Bool
<= :: Policy -> Policy -> Bool
$c> :: Policy -> Policy -> Bool
> :: Policy -> Policy -> Bool
$c>= :: Policy -> Policy -> Bool
>= :: Policy -> Policy -> Bool
$cmax :: Policy -> Policy -> Policy
max :: Policy -> Policy -> Policy
$cmin :: Policy -> Policy -> Policy
min :: Policy -> Policy -> Policy
Ord, Int -> Policy -> ShowS
[Policy] -> ShowS
Policy -> String
(Int -> Policy -> ShowS)
-> (Policy -> String) -> ([Policy] -> ShowS) -> Show Policy
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Policy -> ShowS
showsPrec :: Int -> Policy -> ShowS
$cshow :: Policy -> String
show :: Policy -> String
$cshowList :: [Policy] -> ShowS
showList :: [Policy] -> ShowS
Show)

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

-- | "N per period" with burst N (max = refill = n). 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 -> Policy
tokenBucket :: Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
prefix Double
count NominalDiffTime
period = Text -> Double -> Double -> NominalDiffTime -> Policy
Policy Text
prefix Double
count Double
count (NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Ord a => a -> a -> a
max NominalDiffTime
1e-6 NominalDiffTime
period)

-- | 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 Policy payload (Maybe RateLimitKey)

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

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

-- | Limit by a fixed policy under one shared key (a single global bucket).
globalLimit :: Policy -> Text -> RateLimitFor payload
globalLimit :: forall payload. Policy -> Text -> RateLimitFor payload
globalLimit Policy
pol Text
suffix = Policy -> (payload -> Text) -> RateLimitFor payload
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
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 Policy payload (Maybe RateLimitKey))
-> Selector Policy 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

-- | Run a selector against a concrete job to get its key.
runRateLimitFor :: payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor :: forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor = payload
-> Selector Policy payload (Maybe RateLimitKey)
-> Maybe RateLimitKey
forall policy payload a. payload -> Selector policy payload a -> a
runSelector

-- | A payload's per-job key 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. Defaults to 1.
  rateLimitCost :: payload -> Double
  rateLimitCost payload
_ = Double
1

instance {-# OVERLAPPABLE #-} HasRateLimit payload

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

-- | Collect every policy declared across a registry's payloads, by statically
-- inspecting each payload's 'rateLimitFor'. The migration seeds these.
type RegistryRateLimitPolicies registry = RegistryPolicies registry Policy

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

-- | 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 @Policy