ReaderT monad限流实现验证及TestM实例编写求助

首先修正原代码中的拼写错误：Haskell的instance关键字为小写，原代码中大写的Instance会导致编译失败。

修正后原代码 #

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Network.HTTP.Client (Manager, Request, Response, httpLbs, parseRequest_, newManager, defaultManagerSettings)
import Network.HTTP.Types (Status(..), HTTPVersion(HTTP/1.1))
import System.Environment (getEnv)
import Control.Monad.Reader (ReaderT, ask, liftIO)
-- 假设TokenBucket相关定义来自限流库（比如rate-limit）
import qualified RateLimit as RL (TokenBucket, tokenBucketWait, newTokenBucket)

data TestEnv = TestEnv
  { rateLimiter' :: !RL.TokenBucket
  , apiManager :: !Manager
  , apiKey :: !BS.ByteString
  }

type BunnyReaderT m = ReaderT TestEnv m

class MonadIO m => HasBunny m where
  runRequest :: Request -> m (Response BSL.ByteString)
  applyAuth :: Request -> m Request
  fetchAuth :: m BS.ByteString
  -- 默认实现
  applyAuth req = do
    apiKey <- fetchAuth
    return $ req { requestHeaders = ("AccessKey", apiKey) : requestHeaders req }
  fetchAuth = liftIO $ BS.pack <$> getEnv "AccessKey"

instance MonadIO m => HasBunny (BunnyReaderT m) where
  runRequest req = do
    config <- ask
    authReq <- applyAuth req
    let burstSize = 75
        toInvRate r = round (1e6 / r)  -- 转换为微秒间隔
        invRate = toInvRate 75         -- 每秒75次，每次间隔约13333微秒
    liftIO $ RL.tokenBucketWait (rateLimiter' config) burstSize invRate
    liftIO $ httpLbs authReq (apiManager config) 
  fetchAuth = do
    config <- ask
    return $ apiKey config

type TestM = ReaderT TestEnv IO

问题1：验证令牌桶限流是否正确 #

需从参数逻辑和实际测试两方面验证限流效果：

参数逻辑验证 #

• invRate = round(1e6/75)：计算单个令牌的生成间隔（微秒），1秒=1e6微秒，每秒生成75个令牌，对应间隔约13333微秒，参数逻辑正确。
• burstSize=75：令牌桶最大容量为75，允许初始突发75个请求，之后每个请求需等待令牌生成，符合"每秒75次"的长期限流目标。

实际测试验证 #

编写并发测试代码，统计750个请求的总耗时：

import Control.Concurrent.Async (mapConcurrently)
import Data.Time.Clock (getCurrentTime, diffUTCTime)

testRateLimiter :: IO ()
testRateLimiter = do
  -- 初始化线程安全的令牌桶
  rateLimiter <- RL.newTokenBucket
  manager <- newManager defaultManagerSettings
  apiKey <- BS.pack <$> getEnv "AccessKey"
  let env = TestEnv rateLimiter manager apiKey
      -- 生成750个测试请求
      testRequests = replicate 750 (parseRequest_ "http://example.com/api")
  
  start <- getCurrentTime
  -- 并发执行所有请求
  runReaderT (mapConcurrently runRequest testRequests) env
  end <- getCurrentTime
  
  let duration = realToFrac $ diffUTCTime end start
      reqPerSecond = 750 / duration
  putStrLn $ "总耗时: " ++ show duration ++ " 秒"
  putStrLn $ "实际QPS: " ++ show reqPerSecond

• 预期结果：总耗时接近10秒（750/75=10），实际QPS接近75。
• 注意事项：必须确保TokenBucket是线程安全的（内部用MVar或STM维护令牌计数），否则多线程并发时会出现竞态，导致限流失效。

问题2：编写TestM的HasBunny Mock实例 #

TestM的实例需要保留限流逻辑，同时Mock网络请求返回虚拟响应，代码如下：

import Data.Cookie (emptyCookieJar)
import Network.HTTP.Client (ResponseClose(ResponseClose))

instance HasBunny TestM where
  runRequest req = do
    config <- ask
    authReq <- applyAuth req
    -- 保留原有限流逻辑，确保Mock场景下仍遵守限流规则
    let burstSize = 75
        toInvRate r = round (1e6 / r)
        invRate = toInvRate 75
    liftIO $ RL.tokenBucketWait (rateLimiter' config) burstSize invRate
    -- Mock网络请求，返回虚拟响应
    liftIO $ return $ Response
      { responseStatus = Status 200 "OK"
      , responseVersion = HTTP/1.1
      , responseHeaders = []
      , responseBody = BSL.pack "Mock API Response"
      , responseCookieJar = emptyCookieJar
      , responseClose' = ResponseClose (return ())
      }
  -- 复用从环境读取apiKey的逻辑，无需重新实现
  fetchAuth = do
    config <- ask
    return $ apiKey config

Mock场景下的验证 #

使用上述testRateLimiter函数，运行TestM实现的runRequest，同样应该得到接近10秒的总耗时，说明限流规则在Mock场景下依然生效。

内容的提问来源于stack exchange，提问作者Vikram malik