{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

-- | RateLimit SQL templates.
module Arbiter.Core.Sql.RateLimit
  ( defaultThrottleWaitSeconds
  , addRateLimitTokensSQL
  , pruneRateLimitBucketsSQL
  , resetRateLimitBucketsSQL
  , wakeThrottledJobsSQL
  , wakeThrottledJobsForKeySQL
  , refilledExpr
  , listRateLimitPoliciesSQL
  , getRateLimitPolicySQL
  , rateLimitPolicyExistsSQL
  , listRateLimitBucketsSQL
  , updateRateLimitOverridesSQL
  ) where

import Control.Monad (join)
import Data.Int (Int64)
import Data.Maybe (isJust)
import Data.Text (Text)
import Data.Text qualified as T
import NeatInterpolation (text)

import Arbiter.Core.Admission (effectivePolicyCol)
import Arbiter.Core.Codec (rateLimitBucketCodec, rateLimitPolicyViewCodec)
import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueTable)
import Arbiter.Core.RateLimit.Schema (arbiterRateLimitPoliciesTable, arbiterRateLimitsTable)
import Arbiter.Core.RateLimit.Stats (RateLimitBucketView, RateLimitPolicyView)
import Arbiter.Core.Sql.Jobs (throttledPredicateSQL, unionAllOverQueueTables)
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query, rows, sepBy)

-- | Deny-path wait (seconds) when no refill interval yields a real wait.
defaultThrottleWaitSeconds :: Double
defaultThrottleWaitSeconds :: Double
defaultThrottleWaitSeconds = Double
0.5

-- | Add @amount@ to a key's bucket, clamped to @[0, max]@, seeding an absent bucket at
-- full. A no-op without a policy.
addRateLimitTokensSQL :: SchemaName -> Text -> Text -> Double -> Query ()
addRateLimitTokensSQL :: Text -> Text -> Text -> Double -> Query ()
addRateLimitTokensSQL Text
schema Text
key Text
prefix Double
amount =
  let buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
      effRefill :: Text
effRefill = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"refill_amount"
      effInterval :: Text
effInterval = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"interval"
      refilled :: Text
refilled =
        Text -> Text -> Text -> Text -> Text -> Text
refilledExpr
          Text
"(SELECT max_tokens FROM pol)"
          Text
"bucket.tokens + (SELECT amt FROM input)"
          Text
"bucket.last_refill"
          Text
"(SELECT refill_amount FROM pol)"
          Text
"(SELECT refill_interval FROM pol)"
   in [sql|
        WITH input AS (
          SELECT #{key :: CText}::text AS key, #{prefix :: CText}::text AS prefix, #{amount :: CFloat8}::float8 AS amt
        ),
        pol AS (
          SELECT ${effMax} AS max_tokens, ${effRefill} AS refill_amount, ${effInterval} AS refill_interval
          FROM ${policies} policy WHERE policy.prefix_id = (SELECT prefix FROM input)
        )
        INSERT INTO ${buckets} AS bucket (rate_limit_key, policy_prefix, tokens, last_refill)
        SELECT (SELECT key FROM input), (SELECT prefix FROM input), (SELECT max_tokens FROM pol), NOW()
        FROM pol
        ON CONFLICT (rate_limit_key) DO UPDATE
          SET tokens = GREATEST(0, ${refilled}),
              last_refill = NOW()
      |]

-- | Delete idle buckets that have refilled to max.
pruneRateLimitBucketsSQL :: SchemaName -> Double -> Query ()
pruneRateLimitBucketsSQL :: Text -> Double -> Query ()
pruneRateLimitBucketsSQL Text
schema Double
idleSeconds =
  let buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
   in [sql|
        DELETE FROM ${buckets} bucket
        USING ${policies} policy
        WHERE policy.prefix_id = bucket.policy_prefix
          AND bucket.last_refill < NOW() - (#{idleSeconds :: CFloat8}::float8 * interval '1 second')
          AND ${refilledBucketTokens} >= ${effMax}
      |]

-- | Refill a prefix's buckets to full. The fixed-window reset, which also wakes jobs,
-- is the HighLevel resetRateLimitBuckets.
resetRateLimitBucketsSQL :: SchemaName -> Text -> Query ()
resetRateLimitBucketsSQL :: Text -> Text -> Query ()
resetRateLimitBucketsSQL Text
schema Text
prefix =
  let buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
   in [sql|
        UPDATE ${buckets} bucket
        SET tokens = ${effMax}, last_refill = NOW()
        FROM ${policies} policy
        WHERE policy.prefix_id = bucket.policy_prefix AND bucket.policy_prefix = #{prefix :: CText}
      |]

-- | Clear the rate-limit deferral on a prefix's throttled jobs. One statement over
-- all given queue tables.
wakeThrottledJobsSQL :: SchemaName -> [TableName] -> Text -> Query Int64
wakeThrottledJobsSQL :: Text -> [Text] -> Text -> Query Int64
wakeThrottledJobsSQL Text
schema [Text]
tableNames Text
prefix =
  [Query ()] -> Query Int64
wakeThrottledAcrossSQL [Text -> Text -> Text -> Query () -> Query ()
wakeThrottledBody Text
schema Text
tableName Text
prefix Query ()
forall a. Monoid a => a
mempty | Text
tableName <- [Text]
tableNames]

-- | Like 'wakeThrottledJobsSQL' but for one key, after a top-up.
wakeThrottledJobsForKeySQL :: SchemaName -> [TableName] -> Text -> Text -> Query Int64
wakeThrottledJobsForKeySQL :: Text -> [Text] -> Text -> Text -> Query Int64
wakeThrottledJobsForKeySQL Text
schema [Text]
tableNames Text
prefix Text
key =
  [Query ()] -> Query Int64
wakeThrottledAcrossSQL
    [Text -> Text -> Text -> Query () -> Query ()
wakeThrottledBody Text
schema Text
tableName Text
prefix [sql|AND rate_limit_key = #{key :: CText}|] | Text
tableName <- [Text]
tableNames]

-- | One wake UPDATE CTE per queue table, summed.
wakeThrottledAcrossSQL :: [Query ()] -> Query Int64
wakeThrottledAcrossSQL :: [Query ()] -> Query Int64
wakeThrottledAcrossSQL [Query ()]
bodies =
  let named :: [(Text, Query ())]
named = [(String -> Text
T.pack (String
"wake_" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
index), Query ()
body) | (Int
index, Query ()
body) <- [Int] -> [Query ()] -> [(Int, Query ())]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Query ()]
bodies]
      wakeCtes :: Query ()
wakeCtes = Text -> [Query ()] -> Query ()
sepBy Text
", " [[sql|${cteName} AS (${body} RETURNING 1)|] | (Text
cteName, Query ()
body) <- [(Text, Query ())]
named]
      total :: Text
total = Text -> [Text] -> Text
T.intercalate Text
" + " [[text|(SELECT COUNT(*) FROM ${cteName})|] | (Text
cteName, Query ()
_) <- [(Text, Query ())]
named]
   in [sql|WITH ${wakeCtes} SELECT (${total})::int8 AS @{count :: CInt8}|]

-- | Shared wake UPDATE for a prefix, optionally narrowed by @keyFrag@ (e.g. one key).
wakeThrottledBody :: SchemaName -> TableName -> Text -> Query () -> Query ()
wakeThrottledBody :: Text -> Text -> Text -> Query () -> Query ()
wakeThrottledBody Text
schema Text
tableName Text
prefix Query ()
keyFrag =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
   in [sql|
        UPDATE ${tbl}
        SET not_visible_until = NOW(), throttled_until = NULL, updated_at = NOW()
        WHERE rate_limit_prefix = #{prefix :: CText}
          ${keyFrag}
          AND ${throttledPredicateSQL} AND NOT suspended
      |]

-- | Tokens accrued since @lastRefill@ at @refill@ per @interval@ seconds. Arguments
-- are SQL expressions.
accruedTokensExpr :: Text -> Text -> Text -> Text
accruedTokensExpr :: Text -> Text -> Text -> Text
accruedTokensExpr Text
lastRefill Text
refill Text
interval =
  [text|COALESCE(EXTRACT(EPOCH FROM (NOW() - ${lastRefill})) / NULLIF(${interval}, 0) * ${refill}, 0)|]

-- | A bucket's lazily-refilled token count, capped at @maxTokens@. Arguments are SQL
-- expressions.
refilledExpr :: Text -> Text -> Text -> Text -> Text -> Text
refilledExpr :: Text -> Text -> Text -> Text -> Text -> Text
refilledExpr Text
maxTokens Text
tokens Text
lastRefill Text
refill Text
interval =
  let accrued :: Text
accrued = Text -> Text -> Text -> Text
accruedTokensExpr Text
lastRefill Text
refill Text
interval
   in [text|LEAST(${maxTokens}, ${tokens} + ${accrued})|]

-- | The lazily-refilled token count of the bucket at alias @bucket@ under the policy
-- at alias @policy@. Mirrors the gate's accrual.
refilledBucketTokens :: Text
refilledBucketTokens :: Text
refilledBucketTokens =
  Text -> Text -> Text -> Text -> Text -> Text
refilledExpr
    (Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens")
    Text
"bucket.tokens"
    Text
"bucket.last_refill"
    (Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"refill_amount")
    (Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"interval")

-- | List every policy with its default/override params and per-prefix bucket
-- aggregates (count, min and average of lazily-refilled tokens, live throttled
-- count over the given queue tables).
listRateLimitPoliciesSQL :: SchemaName -> [TableName] -> Query RateLimitPolicyView
listRateLimitPoliciesSQL :: Text -> [Text] -> Query RateLimitPolicyView
listRateLimitPoliciesSQL Text
schema [Text]
tableNames = Text -> [Text] -> Maybe Text -> Query RateLimitPolicyView
rateLimitPoliciesSQL Text
schema [Text]
tableNames Maybe Text
forall a. Maybe a
Nothing

-- | Single-prefix variant of 'listRateLimitPoliciesSQL'.
getRateLimitPolicySQL :: SchemaName -> [TableName] -> Text -> Query RateLimitPolicyView
getRateLimitPolicySQL :: Text -> [Text] -> Text -> Query RateLimitPolicyView
getRateLimitPolicySQL Text
schema [Text]
tableNames Text
prefix = Text -> [Text] -> Maybe Text -> Query RateLimitPolicyView
rateLimitPoliciesSQL Text
schema [Text]
tableNames (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
prefix)

-- | Whether a rate-limit policy exists for a prefix.
rateLimitPolicyExistsSQL :: SchemaName -> Text -> Query Bool
rateLimitPolicyExistsSQL :: Text -> Text -> Query Bool
rateLimitPolicyExistsSQL Text
schema Text
prefix =
  let tbl :: Text
tbl = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
   in [sql|SELECT EXISTS (SELECT 1 FROM ${tbl} WHERE prefix_id = #{prefix :: CText}) AS @{result :: CBool}|]

-- | The policy views, every policy or the one a prefix names.
rateLimitPoliciesSQL :: SchemaName -> [TableName] -> Maybe Text -> Query RateLimitPolicyView
rateLimitPoliciesSQL :: Text -> [Text] -> Maybe Text -> Query RateLimitPolicyView
rateLimitPoliciesSQL Text
schema [Text]
tableNames Maybe Text
mPrefix =
  let policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      throttledPerTable :: Text
throttledPerTable = Text -> [Text] -> (Text -> Text -> Text) -> Text
unionAllOverQueueTables Text
schema [Text]
tableNames ((Text -> Text -> Text) -> Text) -> (Text -> Text -> Text) -> Text
forall a b. (a -> b) -> a -> b
$ \Text
_ Text
table ->
        [text|
          SELECT rate_limit_prefix AS prefix, COUNT(*)::int8 AS throttled
          FROM ${table}
          WHERE ${throttledPredicateSQL} AND NOT suspended AND rate_limit_prefix IS NOT NULL
            AND ((SELECT prefix FROM target) IS NULL OR rate_limit_prefix = (SELECT prefix FROM target))
          GROUP BY rate_limit_prefix
        |]
      throttledJoin :: Text
throttledJoin =
        [text|
          LEFT JOIN (
            SELECT prefix, SUM(throttled)::int8 AS throttled
            FROM (${throttledPerTable}) per_table
            GROUP BY prefix
          ) throttled ON throttled.prefix = policy.prefix_id
        |]
      throttledCol, throttledJoinClause :: Text
      (Text
throttledCol, Text
throttledJoinClause) = case [Text]
tableNames of
        [] -> (Text
"0::int8 AS throttled_count", Text
"")
        [Text]
_ -> (Text
"COALESCE(throttled.throttled, 0) AS throttled_count", Text
throttledJoin)
   in RowCodec RateLimitPolicyView
-> Query () -> Query RateLimitPolicyView
forall a. RowCodec a -> Query () -> Query a
rows
        RowCodec RateLimitPolicyView
rateLimitPolicyViewCodec
        [sql|
        WITH target AS (SELECT #{mPrefix :: Maybe CText}::text AS prefix)
        SELECT policy.prefix_id,
               policy.default_max_tokens, policy.default_refill_amount, policy.default_interval,
               policy.override_max_tokens, policy.override_refill_amount, policy.override_interval,
               COALESCE(agg.bucket_count, 0) AS bucket_count,
               ${throttledCol},
               agg.min_tokens, agg.avg_tokens
        FROM ${policies} policy
        LEFT JOIN (
          SELECT policy_prefix, COUNT(*) AS bucket_count,
                 MIN(tokens) AS min_tokens, AVG(tokens) AS avg_tokens
          FROM (
            SELECT bucket.policy_prefix, ${refilledBucketTokens} AS tokens
            FROM ${buckets} bucket JOIN ${policies} policy ON policy.prefix_id = bucket.policy_prefix
            WHERE (SELECT prefix FROM target) IS NULL OR bucket.policy_prefix = (SELECT prefix FROM target)
          ) refilled_bucket
          GROUP BY policy_prefix
        ) agg ON agg.policy_prefix = policy.prefix_id
        ${throttledJoinClause}
        WHERE (SELECT prefix FROM target) IS NULL OR policy.prefix_id = (SELECT prefix FROM target)
        ORDER BY policy.prefix_id
      |]

-- | List a prefix's buckets with effective max and lazily-refilled fill fraction,
-- paginated.
listRateLimitBucketsSQL :: SchemaName -> Text -> Int64 -> Int64 -> Query RateLimitBucketView
listRateLimitBucketsSQL :: Text -> Text -> Int64 -> Int64 -> Query RateLimitBucketView
listRateLimitBucketsSQL Text
schema Text
prefix Int64
limit Int64
offset =
  let policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
   in RowCodec RateLimitBucketView
-> Query () -> Query RateLimitBucketView
forall a. RowCodec a -> Query () -> Query a
rows
        RowCodec RateLimitBucketView
rateLimitBucketCodec
        [sql|
          SELECT rate_limit_key, policy_prefix, tokens, max_tokens,
                 tokens / NULLIF(max_tokens, 0) AS fill_fraction,
                 last_refill
          FROM (
            SELECT bucket.rate_limit_key, bucket.policy_prefix,
                   ${refilledBucketTokens} AS tokens,
                   ${effMax} AS max_tokens,
                   bucket.last_refill
            FROM ${buckets} bucket
            JOIN ${policies} policy ON policy.prefix_id = bucket.policy_prefix
            WHERE bucket.policy_prefix = #{prefix :: CText}
          ) refilled_bucket
          ORDER BY fill_fraction ASC NULLS LAST, rate_limit_key
          LIMIT #{limit :: CInt8} OFFSET #{offset :: CInt8}
        |]

-- | Set or clear a policy's override params. @Nothing@ leaves a field untouched.
-- @Just v@ writes @v@ (a null clears the override back to the default).
updateRateLimitOverridesSQL
  :: SchemaName -> Maybe (Maybe Double) -> Maybe (Maybe Double) -> Maybe (Maybe Double) -> Text -> Query ()
updateRateLimitOverridesSQL :: Text
-> Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> Text
-> Query ()
updateRateLimitOverridesSQL Text
schema Maybe (Maybe Double)
mMax Maybe (Maybe Double)
mRefill Maybe (Maybe Double)
mInterval Text
prefix =
  let policies :: Text
policies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      setMax :: Bool
setMax = Maybe (Maybe Double) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Maybe Double)
mMax
      maxTokens :: Maybe Double
maxTokens = Maybe (Maybe Double) -> Maybe Double
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join Maybe (Maybe Double)
mMax
      setRefill :: Bool
setRefill = Maybe (Maybe Double) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Maybe Double)
mRefill
      refill :: Maybe Double
refill = Maybe (Maybe Double) -> Maybe Double
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join Maybe (Maybe Double)
mRefill
      setInterval :: Bool
setInterval = Maybe (Maybe Double) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Maybe Double)
mInterval
      interval :: Maybe Double
interval = Maybe (Maybe Double) -> Maybe Double
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join Maybe (Maybe Double)
mInterval
   in [sql|
        UPDATE ${policies}
        SET override_max_tokens = CASE WHEN #{setMax :: CBool}::boolean
                                       THEN #{maxTokens :: Maybe CFloat8}::float8
                                       ELSE override_max_tokens END,
            override_refill_amount = CASE WHEN #{setRefill :: CBool}::boolean
                                          THEN #{refill :: Maybe CFloat8}::float8
                                          ELSE override_refill_amount END,
            override_interval = CASE WHEN #{setInterval :: CBool}::boolean
                                     THEN #{interval :: Maybe CFloat8}::float8
                                     ELSE override_interval END
        WHERE prefix_id = #{prefix :: CText}
      |]