{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}

-- | Backend-parameterized integration tests for per-job rate limiting. Each
-- backend supplies a runner over the shared 'RLReg' registry.
module Arbiter.Test.RateLimit
  ( RLPayload (..)
  , RLReg
  , rateLimitTable
  , setupRateLimitPolicy
  , rateLimitSpec
  ) where

import Arbiter.Core.HighLevel qualified as HL
import Arbiter.Core.Job.DLQ qualified as DLQ
import Arbiter.Core.Job.Schema (jobQueueTable)
import Arbiter.Core.Job.Types
  ( DedupKey (..)
  , JobRead
  , JobStatus (..)
  , JobWrite
  , attempts
  , claimSeq
  , defaultGroupedJob
  , defaultJob
  , jobRateLimitKey
  , payload
  , payloadKeys
  , primaryKey
  , setDedupKey
  )
import Arbiter.Core.MonadArbiter (HasRegistry, getSchema)
import Arbiter.Core.QueueRegistry (Queue)
import Arbiter.Core.RateLimit.Schema
  ( arbiterRateLimitPoliciesTable
  , arbiterRateLimitsTable
  , toPolicyRow
  , upsertPolicyRowSQL
  )
import Arbiter.Core.RateLimit.Spec
  ( HasRateLimit (..)
  , Policy
  , RateLimitFor
  , RateLimitKey (..)
  , chooseWhen
  , collectPolicies
  , limitBy
  , limitByCase
  , noLimit
  , registryRateLimitPolicies
  , runRateLimitFor
  , tokenBucket
  )
import Arbiter.Core.RateLimit.Stats
  ( RateLimitBucketView (..)
  , RateLimitPolicyUpdate (..)
  , RateLimitPolicyView (..)
  )
import Control.Exception (finally)
import Control.Monad (foldM_, void)
import Data.Aeson (FromJSON, ToJSON)
import Data.ByteString (ByteString)
import Data.Foldable (find, traverse_)
import Data.List.NonEmpty qualified as NE
import Data.Maybe (fromMaybe, listToMaybe)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Database.PostgreSQL.Simple (close, connectPostgreSQL)
import GHC.Generics (Generic)
import Hedgehog (Gen, assert, check, evalIO, forAll, property, withTests, (===))
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Hspec
import UnliftIO.Async (mapConcurrently)

import Arbiter.Test.Setup (drainWith, execStatement, execute_)

-- | A payload keyed by tenant, with a per-job token cost.
data RLPayload = RLPayload {RLPayload -> Text
rlTenant :: Text, RLPayload -> Double
rlCost :: Double}
  deriving stock (RLPayload -> RLPayload -> Bool
(RLPayload -> RLPayload -> Bool)
-> (RLPayload -> RLPayload -> Bool) -> Eq RLPayload
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RLPayload -> RLPayload -> Bool
== :: RLPayload -> RLPayload -> Bool
$c/= :: RLPayload -> RLPayload -> Bool
/= :: RLPayload -> RLPayload -> Bool
Eq, (forall x. RLPayload -> Rep RLPayload x)
-> (forall x. Rep RLPayload x -> RLPayload) -> Generic RLPayload
forall x. Rep RLPayload x -> RLPayload
forall x. RLPayload -> Rep RLPayload x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. RLPayload -> Rep RLPayload x
from :: forall x. RLPayload -> Rep RLPayload x
$cto :: forall x. Rep RLPayload x -> RLPayload
to :: forall x. Rep RLPayload x -> RLPayload
Generic, Int -> RLPayload -> ShowS
[RLPayload] -> ShowS
RLPayload -> String
(Int -> RLPayload -> ShowS)
-> (RLPayload -> String)
-> ([RLPayload] -> ShowS)
-> Show RLPayload
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RLPayload -> ShowS
showsPrec :: Int -> RLPayload -> ShowS
$cshow :: RLPayload -> String
show :: RLPayload -> String
$cshowList :: [RLPayload] -> ShowS
showList :: [RLPayload] -> ShowS
Show)
  deriving anyclass (Maybe RLPayload
Value -> Parser [RLPayload]
Value -> Parser RLPayload
(Value -> Parser RLPayload)
-> (Value -> Parser [RLPayload])
-> Maybe RLPayload
-> FromJSON RLPayload
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser RLPayload
parseJSON :: Value -> Parser RLPayload
$cparseJSONList :: Value -> Parser [RLPayload]
parseJSONList :: Value -> Parser [RLPayload]
$comittedField :: Maybe RLPayload
omittedField :: Maybe RLPayload
FromJSON, [RLPayload] -> Value
[RLPayload] -> Encoding
RLPayload -> Bool
RLPayload -> Value
RLPayload -> Encoding
(RLPayload -> Value)
-> (RLPayload -> Encoding)
-> ([RLPayload] -> Value)
-> ([RLPayload] -> Encoding)
-> (RLPayload -> Bool)
-> ToJSON RLPayload
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: RLPayload -> Value
toJSON :: RLPayload -> Value
$ctoEncoding :: RLPayload -> Encoding
toEncoding :: RLPayload -> Encoding
$ctoJSONList :: [RLPayload] -> Value
toJSONList :: [RLPayload] -> Value
$ctoEncodingList :: [RLPayload] -> Encoding
toEncodingList :: [RLPayload] -> Encoding
$comitField :: RLPayload -> Bool
omitField :: RLPayload -> Bool
ToJSON)

-- | A one-queue registry over 'RLPayload'.
type RLReg = '[Queue "arbiter_ratelimit_test" RLPayload]

-- 3 tokens, burst 3, refilling 3 every 2 seconds (1.5 tokens/sec).
rlPolicy :: Policy
rlPolicy :: Policy
rlPolicy = Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
"rl" Double
3 NominalDiffTime
2

instance HasRateLimit RLPayload where
  rateLimitFor :: RateLimitFor RLPayload
rateLimitFor = Policy -> (RLPayload -> Text) -> RateLimitFor RLPayload
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
rlPolicy RLPayload -> Text
rlTenant
  rateLimitCost :: RLPayload -> Double
rateLimitCost = RLPayload -> Double
rlCost

-- | Table name for 'RLReg', shared across backends.
rateLimitTable :: Text
rateLimitTable :: Text
rateLimitTable = Text
"arbiter_ratelimit_test"

-- | Upsert the registry's reflected policy rows into a schema.
setupRateLimitPolicy :: ByteString -> Text -> IO ()
setupRateLimitPolicy :: ByteString -> Text -> IO ()
setupRateLimitPolicy ByteString
connStr Text
schema = do
  conn <- ByteString -> IO Connection
connectPostgreSQL ByteString
connStr
  traverse_ (execute_ conn . upsertPolicyRowSQL schema . toPolicyRow) (Set.toList (registryRateLimitPolicies @RLReg))
  close conn

job :: Text -> JobWrite RLPayload
job :: Text -> JobWrite RLPayload
job Text
tenant = RLPayload -> JobWrite RLPayload
forall payload. payload -> JobWrite payload
defaultJob (Text -> Double -> RLPayload
RLPayload Text
tenant Double
1)

costJob :: Text -> Double -> JobWrite RLPayload
costJob :: Text -> Double -> JobWrite RLPayload
costJob Text
tenant Double
cost = RLPayload -> JobWrite RLPayload
forall payload. payload -> JobWrite payload
defaultJob (Text -> Double -> RLPayload
RLPayload Text
tenant Double
cost)

groupedJob :: Text -> Text -> JobWrite RLPayload
groupedJob :: Text -> Text -> JobWrite RLPayload
groupedJob Text
groupKey Text
tenant = Text -> RLPayload -> JobWrite RLPayload
forall payload. Text -> payload -> JobWrite payload
defaultGroupedJob Text
groupKey (Text -> Double -> RLPayload
RLPayload Text
tenant Double
1)

-- | The rate-limit suite, run against any backend.
rateLimitSpec
  :: forall env m
   . (HasRegistry m RLReg)
  => (forall a. env -> m a -> IO a)
  -> SpecWith env
rateLimitSpec :: forall env (m :: * -> *).
HasRegistry m RLReg =>
(forall a. env -> m a -> IO a) -> SpecWith env
rateLimitSpec forall a. env -> m a -> IO a
runM = do
  let enqueue :: env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [JobWrite RLPayload]
jobs = IO [JobRead RLPayload] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (env -> m [JobRead RLPayload] -> IO [JobRead RLPayload]
forall a. env -> m a -> IO a
runM env
env ([JobWrite RLPayload] -> m [JobRead RLPayload]
forall payload (m :: * -> *).
QueueOperation m payload =>
[JobWrite payload] -> m [JobRead payload]
HL.insertJobsBatch [JobWrite RLPayload]
jobs) :: IO [JobRead RLPayload])
      claim :: env -> IO [JobRead RLPayload]
claim env
env = env -> m [JobRead RLPayload] -> IO [JobRead RLPayload]
forall a. env -> m a -> IO a
runM env
env (Int -> NominalDiffTime -> m [JobRead RLPayload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> NominalDiffTime -> m [JobRead payload]
HL.claimNextVisibleJobs Int
100 NominalDiffTime
60) :: IO [JobRead RLPayload]
      tshow :: Int -> Text
tshow = String -> Text
T.pack (String -> Text) -> (Int -> String) -> Int -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show :: Int -> Text
      -- Delete every bucket, as the reaper's prune does to a full idle bucket.
      deleteBuckets :: env -> IO ()
deleteBuckets env
env = env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
        void $ execStatement ("DELETE FROM " <> arbiterRateLimitsTable schema) []
      -- Clear the queue and keep the buckets.
      deleteJobs :: env -> IO ()
deleteJobs env
env = env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
        void $ execStatement ("DELETE FROM " <> jobQueueTable schema rateLimitTable) []
      -- Fast-forward time by backdating every bucket and job timer. The job UPDATE
      -- fires the groups trigger, which recomputes in_flight_until.
      fastForward :: env -> Int -> IO ()
fastForward env
env Int
secs = env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
        let backdate = Text
" - " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
tshow Int
secs Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" * interval '1 second'"
        void $
          execStatement
            ( "UPDATE "
                <> arbiterRateLimitsTable schema
                <> " SET last_refill = last_refill"
                <> backdate
            )
            []
        void $
          execStatement
            ( "UPDATE "
                <> jobQueueTable schema rateLimitTable
                <> " SET not_visible_until = not_visible_until"
                <> backdate
                <> ", throttled_until = throttled_until"
                <> backdate
            )
            []
      setOverride :: env -> Double -> Double -> IO ()
setOverride env
env Double
maxTokens Double
refillAmount =
        IO ClaimSeq -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ClaimSeq -> IO ()) -> IO ClaimSeq -> IO ()
forall a b. (a -> b) -> a -> b
$
          env -> m ClaimSeq -> IO ClaimSeq
forall a. env -> m a -> IO a
runM
            env
env
            ( Text -> RateLimitPolicyUpdate -> m ClaimSeq
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
Text -> RateLimitPolicyUpdate -> m ClaimSeq
HL.updateRateLimitPolicyOverrides
                Text
"rl"
                (Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> RateLimitPolicyUpdate
RateLimitPolicyUpdate (Maybe Double -> Maybe (Maybe Double)
forall a. a -> Maybe a
Just (Double -> Maybe Double
forall a. a -> Maybe a
Just Double
maxTokens)) (Maybe Double -> Maybe (Maybe Double)
forall a. a -> Maybe a
Just (Double -> Maybe Double
forall a. a -> Maybe a
Just Double
refillAmount)) Maybe (Maybe Double)
forall a. Maybe a
Nothing)
            )
      clearOverride :: env -> IO ()
clearOverride env
env =
        IO ClaimSeq -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ClaimSeq -> IO ()) -> IO ClaimSeq -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m ClaimSeq -> IO ClaimSeq
forall a. env -> m a -> IO a
runM env
env (Text -> RateLimitPolicyUpdate -> m ClaimSeq
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
Text -> RateLimitPolicyUpdate -> m ClaimSeq
HL.updateRateLimitPolicyOverrides Text
"rl" (Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> RateLimitPolicyUpdate
RateLimitPolicyUpdate (Maybe Double -> Maybe (Maybe Double)
forall a. a -> Maybe a
Just Maybe Double
forall a. Maybe a
Nothing) (Maybe Double -> Maybe (Maybe Double)
forall a. a -> Maybe a
Just Maybe Double
forall a. Maybe a
Nothing) Maybe (Maybe Double)
forall a. Maybe a
Nothing))

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"admits up to the bucket size and defers the rest" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
10 (Text -> JobWrite RLPayload
job Text
"burst"))
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"keeps a separate bucket per key" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"iso-a") [JobWrite RLPayload]
-> [JobWrite RLPayload] -> [JobWrite RLPayload]
forall a. Semigroup a => a -> a -> a
<> Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"iso-b"))
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 6

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"skips a throttled key's fresh jobs at claim time" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
10 (Text -> JobWrite RLPayload
job Text
"claimskip"))
    admitted <- env -> IO [JobRead RLPayload]
claim env
env
    length admitted `shouldBe` 3
    -- A drained key's fresh jobs are deferred at claim.
    enqueue env (replicate 3 (job "claimskip"))
    skipped <- claim env
    length skipped `shouldBe` 0

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"moves the claim token when a defer parks a lapsed claim's row" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> JobWrite RLPayload
job Text
"defertoken"]
    [held] <- env -> IO [JobRead RLPayload]
claim env
env
    enqueue env (replicate 2 (job "defertoken"))
    drained <- claim env
    length drained `shouldBe` 2
    void (runM env (HL.setVisibilityTimeout 0 held))
    parked <- claim env
    parked `shouldSatisfy` null
    runM env (HL.setVisibilityTimeoutBatch 120 [held])
      >>= (`shouldBe` [HL.JobReclaimed (primaryKey held) (claimSeq held) (claimSeq held + 1)])

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"refills over time" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"refill"))
    first <- env -> IO [JobRead RLPayload]
claim env
env
    length first `shouldBe` 3
    enqueue env (replicate 3 (job "refill"))
    emptied <- claim env
    length emptied `shouldBe` 0
    fastForward env 3
    refilled <- claim env
    length refilled `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"refills at the policy rate, partially (not all-or-nothing)" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- rl refills 1.5 tokens/sec. Backdating last_refill by 1s accrues about 1.5
    -- tokens, which admits exactly one cost-1 job. The 0.5-token margin absorbs
    -- the real time between backdate and claim.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"refilldet"))
    drained <- env -> IO [JobRead RLPayload]
claim env
env
    length drained `shouldBe` 3
    fastForward env 1
    enqueue env (replicate 3 (job "refilldet"))
    refilled <- claim env
    length refilled `shouldBe` 1

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"caps accrued refill at the bucket max" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Backdating 10s accrues 15 tokens. The bucket max caps it at 3.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"refillcap"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    fastForward env 10
    enqueue env (replicate 5 (job "refillcap"))
    capped <- claim env
    length capped `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"gates grouped heads across groups sharing a key" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> Text -> JobWrite RLPayload
groupedJob (Text
"grp-" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text
tshow Int
index) Text
"g-shared" | Int
index <- [Int
1 .. Int
5]]
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"stalls a throttled group behind its head" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"g-stall"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    enqueue env [groupedJob "stallgrp" "g-stall", groupedJob "stallgrp" "g-stall"]
    firstClaim <- claim env
    length firstClaim `shouldBe` 0
    stalled <- claim env
    length stalled `shouldBe` 0
    fastForward env 3
    resumed <- claim env
    length resumed `shouldBe` 1

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"stalls a mixed-key group on its head" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"mk-head"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    enqueue env [groupedJob "mixed" "mk-head", groupedJob "mixed" "mk-sib"]
    firstClaim <- claim env
    length firstClaim `shouldBe` 0
    stalled <- claim env
    length stalled `shouldBe` 0
    fastForward env 3
    resumed <- claim env
    case resumed of
      [JobRead RLPayload
resumedJob] -> RLPayload -> Text
rlTenant (JobRead RLPayload -> RLPayload
forall payload key q insertedAt adm.
JobRecord payload key q insertedAt adm -> payload
payload JobRead RLPayload
resumedJob) Text -> Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text
"mk-head"
      [JobRead RLPayload]
_ -> HasCallStack => String -> IO ()
String -> IO ()
expectationFailure (String
"expected exactly the head, got " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([JobRead RLPayload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [JobRead RLPayload]
resumed))

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"does not spend a grouped job's attempt when throttled" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- A throttled grouped job is deferred at claim with no attempt charged. It
    -- runs on its first attempt.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"gb"))
    drained <- env -> IO [JobRead RLPayload]
claim env
env
    length drained `shouldBe` 3
    enqueue env [groupedJob "gbgrp" "gb"]
    throttled <- claim env
    length throttled `shouldBe` 0
    fastForward env 3
    resumed <- claim env
    case resumed of
      [JobRead RLPayload
resumedJob] -> do
        RLPayload -> Text
rlTenant (JobRead RLPayload -> RLPayload
forall payload key q insertedAt adm.
JobRecord payload key q insertedAt adm -> payload
payload JobRead RLPayload
resumedJob) Text -> Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text
"gb"
        JobRead RLPayload -> Int32
forall payload q insertedAt adm.
JobRecord payload ClaimSeq q insertedAt adm -> Int32
attempts JobRead RLPayload
resumedJob Int32 -> Int32 -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Int32
1
      [JobRead RLPayload]
_ -> HasCallStack => String -> IO ()
String -> IO ()
expectationFailure (String
"expected exactly the resumed head, got " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([JobRead RLPayload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [JobRead RLPayload]
resumed))

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"preserves group order in batched mode across mixed keys" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> Text -> JobWrite RLPayload
groupedJob Text
"bgroup" Text
"bk-head", Text -> Text -> JobWrite RLPayload
groupedJob Text
"bgroup" Text
"bk-sib"]
    batches <- env
-> m [NonEmpty (JobRead RLPayload)]
-> IO [NonEmpty (JobRead RLPayload)]
forall a. env -> m a -> IO a
runM env
env (Int -> Int -> NominalDiffTime -> m [NonEmpty (JobRead RLPayload)]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> NominalDiffTime -> m [NonEmpty (JobRead payload)]
HL.claimNextVisibleJobsBatched Int
5 Int
100 NominalDiffTime
60) :: IO [NE.NonEmpty (JobRead RLPayload)]
    map (rlTenant . payload) (concatMap NE.toList batches) `shouldBe` ["bk-head", "bk-sib"]

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"stalls a throttled grouped head over a fresh-key sibling in batched mode" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Drain the head's key. A batched claim of a mixed-key group defers the whole
    -- batch on its throttled head.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"bh-head"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    enqueue env [groupedJob "bgrp" "bh-head", groupedJob "bgrp" "bh-sib"]
    stalled <- runM env (HL.claimNextVisibleJobsBatched 5 100 60) :: IO [NE.NonEmpty (JobRead RLPayload)]
    concatMap NE.toList stalled `shouldSatisfy` null
    fastForward env 3
    resumed <- runM env (HL.claimNextVisibleJobsBatched 5 100 60)
    map (rlTenant . payload) (concatMap NE.toList resumed) `shouldBe` ["bh-head", "bh-sib"]

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"re-seeds a missing bucket before admitting" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- A policied key whose bucket was pruned is limited again.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
10 (Text -> JobWrite RLPayload
job Text
"reseed"))
    env -> IO ()
deleteBuckets env
env
    first <- env -> IO [JobRead RLPayload]
claim env
env
    length first `shouldBe` 0
    second <- claim env
    length second `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"keeps a throttled head stalled after a granted batch-mate is acked" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Drain the shared key to one token. A batch grants the head and throttles
    -- its sibling. Acking the granted head keeps the group's in-flight marker.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
2 (Text -> JobWrite RLPayload
job Text
"ov"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    enqueue env [groupedJob "ovgrp" "ov", groupedJob "ovgrp" "ov", groupedJob "ovgrp" "ov-free"]
    batches <- runM env (HL.claimNextVisibleJobsBatched 2 100 60) :: IO [NE.NonEmpty (JobRead RLPayload)]
    let granted = (NonEmpty (JobRead RLPayload) -> [JobRead RLPayload])
-> [NonEmpty (JobRead RLPayload)] -> [JobRead RLPayload]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap NonEmpty (JobRead RLPayload) -> [JobRead RLPayload]
forall a. NonEmpty a -> [a]
NE.toList [NonEmpty (JobRead RLPayload)]
batches
    map (rlTenant . payload) granted `shouldBe` ["ov"]
    _ <- runM env (HL.ackJobsBatch granted)
    -- The fresh-key sibling is never handed out ahead of the throttled head.
    overtaken <- claim env
    map (rlTenant . payload) overtaken `shouldSatisfy` notElem "ov-free"

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"keeps a throttled survivor stalled when a sibling is dedup-moved to another group" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Both grouped "ov" jobs throttle on a drained key. Dedup-replacing one of them
    -- into another group recomputes the old group's in-flight marker from the
    -- throttled survivor.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"ov"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    let mover = Maybe DedupKey -> JobWrite RLPayload -> JobWrite RLPayload
forall payload.
Maybe DedupKey -> JobWrite payload -> JobWrite payload
setDedupKey (DedupKey -> Maybe DedupKey
forall a. a -> Maybe a
Just (Text -> DedupKey
ReplaceDuplicate Text
"mover-key")) (JobWrite RLPayload -> JobWrite RLPayload)
-> JobWrite RLPayload -> JobWrite RLPayload
forall a b. (a -> b) -> a -> b
$ Text -> Text -> JobWrite RLPayload
groupedJob Text
"ovgrp" Text
"ov"
    enqueue env [groupedJob "ovgrp" "ov", mover, groupedJob "ovgrp" "ov-free"]
    batches <- runM env (HL.claimNextVisibleJobsBatched 2 100 60) :: IO [NE.NonEmpty (JobRead RLPayload)]
    length (concatMap NE.toList batches) `shouldBe` 0
    enqueue env [setDedupKey (Just (ReplaceDuplicate "mover-key")) $ groupedJob "othergrp" "ov"]
    overtaken <- claim env
    map (rlTenant . payload) overtaken `shouldSatisfy` notElem "ov-free"

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"defers the over-budget jobs with a future wake" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Five jobs share a key with bucket size 3. The claim admits 3 and parks the
    -- rest with a future wake.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"throttlecb"))
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 3
    again <- claim env
    length again `shouldBe` 0
    fastForward env 3
    woken <- claim env
    length woken `shouldBe` 2

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"reports a throttle-deferred job as throttled status" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"statuskey"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    throttled <- runM env (HL.countJobsFiltered @RLPayload [HL.FilterStatus Throttled])
    throttled `shouldBe` 2

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"spends a job's full cost" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- The bucket holds 3. Of two cost-2 jobs only the first fits.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> Double -> JobWrite RLPayload
costJob Text
"weighted" Double
2, Text -> Double -> JobWrite RLPayload
costJob Text
"weighted" Double
2]
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 1

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"debits the bucket on admission and parks the overflow without charging its attempt" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Freeze refill. The post-claim balance is exact. Admission spends tokens
    -- and denies the overflow.
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (env -> IO ()
clearOverride env
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      env -> Double -> Double -> IO ()
setOverride env
env Double
3 Double
0
      env -> IO ()
deleteBuckets env
env
      env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"debit"))
      kept <- env -> IO [JobRead RLPayload]
claim env
env
      length kept `shouldBe` 3
      remaining <- listToMaybe . map tokens <$> runM env (HL.listRateLimitBuckets "rl" 100 0)
      remaining `shouldBe` Just 0
      -- The top-up refills the frozen bucket and wakes the two parked jobs. They
      -- run on their first attempt.
      runM env (HL.addRateLimitTokens (RateLimitKey "rl" "debit") 3)
      woken <- claim env
      map attempts woken `shouldBe` [1, 1]

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"tops up a bucket manually" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"topup"))
    drained <- env -> IO [JobRead RLPayload]
claim env
env
    length drained `shouldBe` 3
    runM env (HL.addRateLimitTokens (RateLimitKey "rl" "topup") 3)
    enqueue env (replicate 3 (job "topup"))
    toppedUp <- claim env
    length toppedUp `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"wakes a key's deferred jobs on a top-up" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
6 (Text -> JobWrite RLPayload
job Text
"topupwake"))
    admitted <- env -> IO [JobRead RLPayload]
claim env
env
    length admitted `shouldBe` 3
    -- The top-up refills the bucket and wakes the parked jobs.
    runM env (HL.addRateLimitTokens (RateLimitKey "rl" "topupwake") 3)
    woken <- claim env
    length woken `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"seeds an absent bucket at full on a top-up" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- An absent bucket is full. A top-up of 1 leaves it full.
    env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (RateLimitKey -> Double -> m ()
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
RateLimitKey -> Double -> m ()
HL.addRateLimitTokens (Text -> Text -> RateLimitKey
RateLimitKey Text
"rl" Text
"seedfull") Double
1)
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"seedfull"))
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"wakes deferred jobs on reset so a window release is immediate" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
6 (Text -> JobWrite RLPayload
job Text
"windowkey"))
    admitted <- env -> IO [JobRead RLPayload]
claim env
env
    length admitted `shouldBe` 3
    -- Reset wakes the deferred jobs.
    _ <- runM env (HL.resetRateLimitBuckets "rl")
    afterReset <- claim env
    length afterReset `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"clears a stale throttle marker on claim so a wake cannot double-claim it" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> JobWrite RLPayload
job Text
"claimclears"]
    -- A ready job with a stale marker. Claiming clears it.
    env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
      void $
        execStatement
          ( "UPDATE "
              <> jobQueueTable schema rateLimitTable
              <> " SET throttled_until = NOW() + interval '60 second'"
          )
          []
    claimed <- env -> IO [JobRead RLPayload]
claim env
env
    length claimed `shouldBe` 1
    _ <- runM env (HL.resetRateLimitBuckets "rl")
    reclaimed <- claim env
    length reclaimed `shouldBe` 0

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"runs a too-costly job once per window" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Cost 5 exceeds the bucket max 3. The spend clamps to the max. The first job
    -- drains a full bucket and runs. The second is deferred. The pre-fund seeds
    -- the bucket through the top-up path.
    env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (RateLimitKey -> Double -> m ()
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
RateLimitKey -> Double -> m ()
HL.addRateLimitTokens (Text -> Text -> RateLimitKey
RateLimitKey Text
"rl" Text
"cmax") Double
3)
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> Double -> JobWrite RLPayload
costJob Text
"cmax" Double
5, Text -> Double -> JobWrite RLPayload
costJob Text
"cmax" Double
5]
    kept <- env -> IO [JobRead RLPayload]
claim env
env
    length kept `shouldBe` 1

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"preserves the rate-limit key through a DLQ retry" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> JobWrite RLPayload
job Text
"dlqkey"]
    claimed <- env -> IO [JobRead RLPayload]
claim env
env
    case claimed of
      [JobRead RLPayload
claimedJob] -> do
        _ <- env -> m ClaimSeq -> IO ClaimSeq
forall a. env -> m a -> IO a
runM env
env (Text -> JobRead RLPayload -> m ClaimSeq
forall payload (m :: * -> *).
JobOperation m payload =>
Text -> JobRead payload -> m ClaimSeq
HL.moveToDLQ Text
"boom" JobRead RLPayload
claimedJob)
        dlqs <- runM env (HL.listDLQJobs 10 0) :: IO [DLQ.DLQJob RLPayload]
        case find ((== "dlqkey") . rlTenant . payload . DLQ.jobSnapshot) dlqs of
          Just DLQJob RLPayload
dlqJob -> do
            retried <- env
-> m (Maybe (JobRead RLPayload)) -> IO (Maybe (JobRead RLPayload))
forall a. env -> m a -> IO a
runM env
env (ClaimSeq -> m (Maybe (JobRead RLPayload))
forall payload (m :: * -> *).
QueueOperation m payload =>
ClaimSeq -> m (Maybe (JobRead payload))
HL.retryFromDLQ (DLQJob RLPayload -> ClaimSeq
forall payload. DLQJob payload -> ClaimSeq
DLQ.dlqPrimaryKey DLQJob RLPayload
dlqJob)) :: IO (Maybe (JobRead RLPayload))
            (jobRateLimitKey . payloadKeys <$> retried) `shouldBe` Just (Just (RateLimitKey "rl" "dlqkey"))
          Maybe (DLQJob RLPayload)
Nothing -> HasCallStack => String -> IO ()
String -> IO ()
expectationFailure String
"job did not reach the DLQ"
      [JobRead RLPayload]
_ -> HasCallStack => String -> IO ()
String -> IO ()
expectationFailure (String
"expected exactly one claimed job, got " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([JobRead RLPayload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [JobRead RLPayload]
claimed))

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"does not prune a drained bucket" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- A drained bucket is not full. Pruning leaves it in place.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"pd"))
    drained <- env -> IO [JobRead RLPayload]
claim env
env
    length drained `shouldBe` 3
    _ <- runM env (HL.pruneRateLimitBuckets 0)
    enqueue env (replicate 3 (job "pd"))
    again <- claim env
    length again `shouldBe` 0

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"prunes a full bucket without creating a burst" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
3 (Text -> JobWrite RLPayload
job Text
"pf"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    fastForward env 3
    pruned <- runM env (HL.pruneRateLimitBuckets 0)
    pruned `shouldSatisfy` (>= 1)
    enqueue env (replicate 10 (job "pf"))
    kept <- claim env
    length kept `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"never over-admits one key under concurrent claiming" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Many workers claim the same key at once. Exactly the bucket size is admitted.
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
60 (Text -> JobWrite RLPayload
job Text
"conc"))
    results <- (Int -> IO [JobRead RLPayload])
-> [Int] -> IO [[JobRead RLPayload]]
forall (m :: * -> *) (t :: * -> *) a b.
(MonadUnliftIO m, Traversable t) =>
(a -> m b) -> t a -> m (t b)
mapConcurrently (IO [JobRead RLPayload] -> Int -> IO [JobRead RLPayload]
forall a b. a -> b -> a
const (env -> IO [JobRead RLPayload]
claim env
env)) [Int
1 .. Int
10 :: Int]
    length (concat results) `shouldBe` 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"never over-admits across many keys under concurrent claiming" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- A zero refill freezes the buckets. The per-key drain is exact.
    let tenants :: [Text]
tenants = [Int -> Text
tshow Int
index | Int
index <- [Int
1 .. Int
12]]
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (env -> IO ()
clearOverride env
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      env -> Double -> Double -> IO ()
setOverride env
env Double
3 Double
0
      env -> IO ()
deleteBuckets env
env
      env -> [JobWrite RLPayload] -> IO ()
enqueue env
env ([[JobWrite RLPayload]] -> [JobWrite RLPayload]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
8 (Text -> JobWrite RLPayload
job Text
tenant) | Text
tenant <- [Text]
tenants])
      burst <- [[JobRead RLPayload]] -> [JobRead RLPayload]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[JobRead RLPayload]] -> [JobRead RLPayload])
-> IO [[JobRead RLPayload]] -> IO [JobRead RLPayload]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO [JobRead RLPayload])
-> [Int] -> IO [[JobRead RLPayload]]
forall (m :: * -> *) (t :: * -> *) a b.
(MonadUnliftIO m, Traversable t) =>
(a -> m b) -> t a -> m (t b)
mapConcurrently (IO [JobRead RLPayload] -> Int -> IO [JobRead RLPayload]
forall a b. a -> b -> a
const (env -> IO [JobRead RLPayload]
claim env
env)) [Int
1 .. Int
8 :: Int]
      rest <- drainWith (claim env)
      let admitted = [JobRead RLPayload]
burst [JobRead RLPayload] -> [JobRead RLPayload] -> [JobRead RLPayload]
forall a. Semigroup a => a -> a -> a
<> [JobRead RLPayload]
rest
          perTenant Text
tenant = [JobRead RLPayload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((JobRead RLPayload -> Bool)
-> [JobRead RLPayload] -> [JobRead RLPayload]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
tenant) (Text -> Bool)
-> (JobRead RLPayload -> Text) -> JobRead RLPayload -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RLPayload -> Text
rlTenant (RLPayload -> Text)
-> (JobRead RLPayload -> RLPayload) -> JobRead RLPayload -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead RLPayload -> RLPayload
forall payload key q insertedAt adm.
JobRecord payload key q insertedAt adm -> payload
payload) [JobRead RLPayload]
admitted)
      map perTenant tenants `shouldBe` replicate (length tenants) 3

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"reports policy and bucket stats through the management plane" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
2 (Text -> JobWrite RLPayload
job Text
"stat"))
    _ <- env -> IO [JobRead RLPayload]
claim env
env
    policies <- runM env HL.listRateLimitPolicies
    case filter ((== "rl") . prefix) policies of
      [RateLimitPolicyView
policyView] -> do
        RateLimitPolicyView -> Double
defaultMaxTokens RateLimitPolicyView
policyView Double -> Double -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Double
3
        RateLimitPolicyView -> ClaimSeq
bucketCount RateLimitPolicyView
policyView ClaimSeq -> (ClaimSeq -> Bool) -> IO ()
forall a. (HasCallStack, Show a) => a -> (a -> Bool) -> IO ()
`shouldSatisfy` (ClaimSeq -> ClaimSeq -> Bool
forall a. Ord a => a -> a -> Bool
>= ClaimSeq
1)
      [RateLimitPolicyView]
_ -> HasCallStack => String -> IO ()
String -> IO ()
expectationFailure String
"expected exactly the rl policy"
    buckets <- runM env (HL.listRateLimitBuckets "rl" 100 0)
    length buckets `shouldSatisfy` (>= 1)
    map policyPrefix buckets `shouldSatisfy` all (== "rl")

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"applies and clears a policy override through the management plane" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Override max to 0 pauses the prefix. Clearing it restores the default.
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (env -> IO ()
clearOverride env
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      updated <- env -> m ClaimSeq -> IO ClaimSeq
forall a. env -> m a -> IO a
runM env
env (Text -> RateLimitPolicyUpdate -> m ClaimSeq
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
Text -> RateLimitPolicyUpdate -> m ClaimSeq
HL.updateRateLimitPolicyOverrides Text
"rl" (Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> Maybe (Maybe Double)
-> RateLimitPolicyUpdate
RateLimitPolicyUpdate (Maybe Double -> Maybe (Maybe Double)
forall a. a -> Maybe a
Just (Double -> Maybe Double
forall a. a -> Maybe a
Just Double
0)) Maybe (Maybe Double)
forall a. Maybe a
Nothing Maybe (Maybe Double)
forall a. Maybe a
Nothing))
      updated `shouldBe` 1
      enqueue env [job "paused"]
      blocked <- claim env
      length blocked `shouldBe` 0
      clearOverride env
      enqueue env [job "resumed"]
      admitted <- claim env
      length admitted `shouldBe` 1

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"admits a job whose prefix has no policy (fail-open)" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- With no policy row every job runs.
    let restore :: IO ()
restore = env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
          void $ execStatement (upsertPolicyRowSQL schema (toPolicyRow rlPolicy)) []
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally IO ()
restore (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        schema <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
        void $ execStatement ("DELETE FROM " <> arbiterRateLimitPoliciesTable schema <> " WHERE prefix_id = 'rl'") []
      env -> [JobWrite RLPayload] -> IO ()
enqueue env
env (Int -> JobWrite RLPayload -> [JobWrite RLPayload]
forall a. Int -> a -> [a]
replicate Int
5 (Text -> JobWrite RLPayload
job Text
"failopen"))
      admitted <- env -> IO [JobRead RLPayload]
claim env
env
      length admitted `shouldBe` 5

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"gate admission matches the reference token-bucket model over random ops" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- A zero refill override makes token math integral. Each consume enqueues and
    -- claims one cost-job, checked against the pure model.
    let modelTenant :: Text
modelTenant = Text
"model"
        modelKey :: RateLimitKey
modelKey = Text -> Text -> RateLimitKey
RateLimitKey Text
"rl" Text
modelTenant
        consume :: Int -> IO Bool
consume Int
cost = do
          env -> [JobWrite RLPayload] -> IO ()
enqueue env
env [Text -> Double -> JobWrite RLPayload
costJob Text
modelTenant (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cost)]
          claimed <- env -> IO [JobRead RLPayload]
claim env
env
          deleteJobs env
          pure (not (null claimed))
        -- With zero refill the bucket view's token count is the raw stored count.
        readStored :: IO (Maybe Double)
readStored = [Double] -> Maybe Double
forall a. [a] -> Maybe a
listToMaybe ([Double] -> Maybe Double)
-> ([RateLimitBucketView] -> [Double])
-> [RateLimitBucketView]
-> Maybe Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RateLimitBucketView -> Double)
-> [RateLimitBucketView] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map RateLimitBucketView -> Double
tokens ([RateLimitBucketView] -> Maybe Double)
-> IO [RateLimitBucketView] -> IO (Maybe Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> env -> m [RateLimitBucketView] -> IO [RateLimitBucketView]
forall a. env -> m a -> IO a
runM env
env (Text -> Int -> Int -> m [RateLimitBucketView]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Int -> Int -> m [RateLimitBucketView]
HL.listRateLimitBuckets Text
"rl" Int
100 Int
0)
        step :: Int -> Op -> PropertyT IO Int
step Int
balance Op
operation = do
          balance' <- case Op
operation of
            Consume Int
cost -> do
              granted <- IO Bool -> PropertyT IO Bool
forall (m :: * -> *) a.
(MonadTest m, MonadIO m, HasCallStack) =>
IO a -> m a
evalIO (Int -> IO Bool
consume Int
cost)
              let (expected, next) = modelConsume balance cost
              granted === expected
              pure next
            TopUp Int
amt -> do
              IO () -> PropertyT IO ()
forall (m :: * -> *) a.
(MonadTest m, MonadIO m, HasCallStack) =>
IO a -> m a
evalIO (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (RateLimitKey -> Double -> m ()
forall (m :: * -> *).
(MonadArbiter m, RegistryTables (RegistryOf m)) =>
RateLimitKey -> Double -> m ()
HL.addRateLimitTokens RateLimitKey
modelKey (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
amt)))
              Int -> PropertyT IO Int
forall a. a -> PropertyT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Int
modelTopUp Int
balance Int
amt)
            Op
Prune -> do
              IO () -> PropertyT IO ()
forall (m :: * -> *) a.
(MonadTest m, MonadIO m, HasCallStack) =>
IO a -> m a
evalIO (IO ClaimSeq -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (env -> m ClaimSeq -> IO ClaimSeq
forall a. env -> m a -> IO a
runM env
env (NominalDiffTime -> m ClaimSeq
forall (m :: * -> *).
MonadArbiter m =>
NominalDiffTime -> m ClaimSeq
HL.pruneRateLimitBuckets NominalDiffTime
0)))
              Int -> PropertyT IO Int
forall a. a -> PropertyT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
balance
          -- A pruned bucket reads as absent, which equals a full bucket.
          stored <- evalIO readStored
          fromMaybe (fromIntegral modelMax) stored === fromIntegral balance'
          pure balance'
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (env -> IO ()
clearOverride env
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      passed <- Property -> IO Bool
forall (m :: * -> *). MonadIO m => Property -> m Bool
check (Property -> IO Bool) -> Property -> IO Bool
forall a b. (a -> b) -> a -> b
$ TestLimit -> Property -> Property
withTests TestLimit
50 (Property -> Property) -> Property -> Property
forall a b. (a -> b) -> a -> b
$ HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
        ops <- Gen [Op] -> PropertyT IO [Op]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll Gen [Op]
genOps
        evalIO (setOverride env (fromIntegral modelMax) 0 >> deleteBuckets env)
        foldM_ step modelMax ops
      passed `shouldBe` True

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"never over- or under-admits one key under concurrent mixed-cost claims" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
    -- Zero refill, bucket 10. Concurrent claimers drain one key whose jobs carry random costs.
    (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (env -> IO ()
clearOverride env
env) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      passed <- Property -> IO Bool
forall (m :: * -> *). MonadIO m => Property -> m Bool
check (Property -> IO Bool) -> Property -> IO Bool
forall a b. (a -> b) -> a -> b
$ TestLimit -> Property -> Property
withTests TestLimit
20 (Property -> Property) -> Property -> Property
forall a b. (a -> b) -> a -> b
$ HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
        costs <- Gen [Int] -> PropertyT IO [Int]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Range Int -> GenT Identity Int -> Gen [Int]
forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
2 Int
15) (Range Int -> GenT Identity Int
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear (Int
1 :: Int) Int
3)))
        evalIO (setOverride env 10 0 >> deleteBuckets env >> deleteJobs env)
        evalIO (enqueue env [costJob "concmix" (fromIntegral cost) | cost <- costs])
        -- Drain to quiescence.
        let sweep = [[JobRead RLPayload]] -> [JobRead RLPayload]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[JobRead RLPayload]] -> [JobRead RLPayload])
-> IO [[JobRead RLPayload]] -> IO [JobRead RLPayload]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO [JobRead RLPayload])
-> [Int] -> IO [[JobRead RLPayload]]
forall (m :: * -> *) (t :: * -> *) a b.
(MonadUnliftIO m, Traversable t) =>
(a -> m b) -> t a -> m (t b)
mapConcurrently (IO [JobRead RLPayload] -> Int -> IO [JobRead RLPayload]
forall a b. a -> b -> a
const (env -> IO [JobRead RLPayload]
claim env
env)) [Int
1 .. Int
8 :: Int]
        claimed <- evalIO (drainWith sweep)
        let admittedCost = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((JobRead RLPayload -> Double) -> [JobRead RLPayload] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (RLPayload -> Double
rlCost (RLPayload -> Double)
-> (JobRead RLPayload -> RLPayload) -> JobRead RLPayload -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead RLPayload -> RLPayload
forall payload key q insertedAt adm.
JobRecord payload key q insertedAt adm -> payload
payload) [JobRead RLPayload]
claimed) :: Double
            total = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
costs) :: Double
            slack = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int]
costs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) :: Double
        -- Spend never exceeds the cap.
        assert (admittedCost <= 10)
        -- The bucket fills to within one job's cost.
        assert (admittedCost >= min total 10 - slack)
      passed `shouldBe` True

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"collects both branches of a choice selector and runs the chosen one" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
_env -> do
    -- A per-job choice statically yields every policy it could use.
    let policyA :: Policy
policyA = Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
"ca" Double
1 NominalDiffTime
1
        policyB :: Policy
policyB = Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
"cb" Double
2 NominalDiffTime
2
        sel :: RateLimitFor Bool
        sel :: RateLimitFor Bool
sel = (Bool -> Bool)
-> RateLimitFor Bool -> RateLimitFor Bool -> RateLimitFor Bool
forall payload policy a.
(payload -> Bool)
-> Selector policy payload a
-> Selector policy payload a
-> Selector policy payload a
chooseWhen Bool -> Bool
forall a. a -> a
id (Policy -> (Bool -> Text) -> RateLimitFor Bool
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
policyA (Text -> Bool -> Text
forall a b. a -> b -> a
const Text
"x")) (Policy -> (Bool -> Text) -> RateLimitFor Bool
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
policyB (Text -> Bool -> Text
forall a b. a -> b -> a
const Text
"y"))
    Set Policy -> [Policy]
forall a. Set a -> [a]
Set.toList (RateLimitFor Bool -> Set Policy
forall policy payload a.
Ord policy =>
Selector policy payload a -> Set policy
collectPolicies RateLimitFor Bool
sel) [Policy] -> [Policy] -> IO ()
forall a. (HasCallStack, Show a, Eq a) => [a] -> [a] -> IO ()
`shouldMatchList` [Policy
policyA, Policy
policyB]
    (RateLimitKey -> Text
rlkPrefix (RateLimitKey -> Text) -> Maybe RateLimitKey -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> RateLimitFor Bool -> Maybe RateLimitKey
forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor Bool
True RateLimitFor Bool
sel) Maybe Text -> Maybe Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"ca"
    (RateLimitKey -> Text
rlkPrefix (RateLimitKey -> Text) -> Maybe RateLimitKey -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> RateLimitFor Bool -> Maybe RateLimitKey
forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor Bool
False RateLimitFor Bool
sel) Maybe Text -> Maybe Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"cb"

  String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"limitByCase collects every branch and runs the matched one" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
_env -> do
    let policyA :: Policy
policyA = Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
"la" Double
1 NominalDiffTime
1
        policyB :: Policy
policyB = Text -> Double -> NominalDiffTime -> Policy
tokenBucket Text
"lb" Double
2 NominalDiffTime
2
        sel :: RateLimitFor Ordering
        sel :: RateLimitFor Ordering
sel = (Ordering -> Ordering)
-> (Ordering -> RateLimitFor Ordering) -> RateLimitFor Ordering
forall k payload.
(Bounded k, Enum k, Eq k) =>
(payload -> k)
-> (k -> RateLimitFor payload) -> RateLimitFor payload
limitByCase Ordering -> Ordering
forall a. a -> a
id ((Ordering -> RateLimitFor Ordering) -> RateLimitFor Ordering)
-> (Ordering -> RateLimitFor Ordering) -> RateLimitFor Ordering
forall a b. (a -> b) -> a -> b
$ \case
          Ordering
LT -> Policy -> (Ordering -> Text) -> RateLimitFor Ordering
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
policyA (Text -> Ordering -> Text
forall a b. a -> b -> a
const Text
"x")
          Ordering
EQ -> RateLimitFor Ordering
forall payload. RateLimitFor payload
noLimit
          Ordering
GT -> Policy -> (Ordering -> Text) -> RateLimitFor Ordering
forall payload. Policy -> (payload -> Text) -> RateLimitFor payload
limitBy Policy
policyB (Text -> Ordering -> Text
forall a b. a -> b -> a
const Text
"y")
    Set Policy -> [Policy]
forall a. Set a -> [a]
Set.toList (RateLimitFor Ordering -> Set Policy
forall policy payload a.
Ord policy =>
Selector policy payload a -> Set policy
collectPolicies RateLimitFor Ordering
sel) [Policy] -> [Policy] -> IO ()
forall a. (HasCallStack, Show a, Eq a) => [a] -> [a] -> IO ()
`shouldMatchList` [Policy
policyA, Policy
policyB]
    (RateLimitKey -> Text
rlkPrefix (RateLimitKey -> Text) -> Maybe RateLimitKey -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ordering -> RateLimitFor Ordering -> Maybe RateLimitKey
forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor Ordering
LT RateLimitFor Ordering
sel) Maybe Text -> Maybe Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"la"
    Ordering -> RateLimitFor Ordering -> Maybe RateLimitKey
forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor Ordering
EQ RateLimitFor Ordering
sel Maybe RateLimitKey -> Maybe RateLimitKey -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Maybe RateLimitKey
forall a. Maybe a
Nothing
    (RateLimitKey -> Text
rlkPrefix (RateLimitKey -> Text) -> Maybe RateLimitKey -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ordering -> RateLimitFor Ordering -> Maybe RateLimitKey
forall payload.
payload -> RateLimitFor payload -> Maybe RateLimitKey
runRateLimitFor Ordering
GT RateLimitFor Ordering
sel) Maybe Text -> Maybe Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"lb"

-- Pure reference bucket with zero refill. Tokens are integral.

modelMax :: Int
modelMax :: Int
modelMax = Int
10

modelConsume :: Int -> Int -> (Bool, Int)
modelConsume :: Int -> Int -> (Bool, Int)
modelConsume Int
balance Int
cost
  | Int
modelMax Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
balance Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ecost = (Bool
True, Int
balance Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ecost)
  | Bool
otherwise = (Bool
False, Int
balance)
  where
    ecost :: Int
ecost = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
cost Int
modelMax)

modelTopUp :: Int -> Int -> Int
modelTopUp :: Int -> Int -> Int
modelTopUp Int
balance = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
modelMax (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int
balance Int -> Int -> Int
forall a. Num a => a -> a -> a
+)

data Op = Consume Int | TopUp Int | Prune
  deriving stock (Int -> Op -> ShowS
[Op] -> ShowS
Op -> String
(Int -> Op -> ShowS)
-> (Op -> String) -> ([Op] -> ShowS) -> Show Op
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Op -> ShowS
showsPrec :: Int -> Op -> ShowS
$cshow :: Op -> String
show :: Op -> String
$cshowList :: [Op] -> ShowS
showList :: [Op] -> ShowS
Show)

genOps :: Gen [Op]
genOps :: Gen [Op]
genOps =
  Range Int -> GenT Identity Op -> Gen [Op]
forall (m :: * -> *) a. MonadGen m => Range Int -> m a -> m [a]
Gen.list (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
1 Int
30) (GenT Identity Op -> Gen [Op]) -> GenT Identity Op -> Gen [Op]
forall a b. (a -> b) -> a -> b
$
    [GenT Identity Op] -> GenT Identity Op
forall (m :: * -> *) a. (HasCallStack, MonadGen m) => [m a] -> m a
Gen.choice
      [ Int -> Op
Consume (Int -> Op) -> GenT Identity Int -> GenT Identity Op
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Int -> GenT Identity Int
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Int -> Int -> Int -> Range Int
forall a. Integral a => a -> a -> a -> Range a
Range.linearFrom Int
1 (-Int
3) (Int
modelMax Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3))
      , Int -> Op
TopUp (Int -> Op) -> GenT Identity Int -> GenT Identity Op
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Range Int -> GenT Identity Int
forall (m :: * -> *) a. (MonadGen m, Integral a) => Range a -> m a
Gen.integral (Int -> Int -> Range Int
forall a. Integral a => a -> a -> Range a
Range.linear Int
0 (Int
modelMax Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3))
      , Op -> GenT Identity Op
forall a. a -> GenT Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Op
Prune
      ]