smsb-plus/script/OpenResty/lua/bandwidth_limiter.lua

-- lua/bandwidth_limiter.lua
local _M = {}
local redis = require "resty.redis"

local redis_host = "127.0.0.1"
local redis_port = 6379
local redis_password = "Hycpb@123"
local redis_db = 6

local redis_prefix = "tenant:device_count:"
local max_bandwidth = 1024 * 1024 -- 1MB
local min_bandwidth = 64 * 1024   -- 最小限速 64KB

local function get_redis()
    local red = redis:new()
    red:set_timeout(1000)

    local ok, err = red:connect(redis_host, redis_port)
    if not ok then
        ngx.log(ngx.ERR, "Redis connection error: ", err)
        return nil, err
    end

    -- 认证密码
    if redis_password and redis_password ~= "" then
        local ok, err = red:auth(redis_password)
        if not ok then
            ngx.log(ngx.ERR, "Redis auth failed: ", err)
            return nil, err
        end
    end

    -- 选择数据库
    if redis_db then
        local ok, err = red:select(redis_db)
        if not ok then
            ngx.log(ngx.ERR, "Redis select db failed: ", err)
            return nil, err
        end
    end

    return red
end

function _M.register_device(tenant_id)
    local red, err = get_redis()
    if not red then return min_bandwidth end

    local key = redis_prefix .. tenant_id
    local count, err = red:incr(key)
    if not count then
        ngx.log(ngx.ERR, "Redis INCR failed: ", err)
        return min_bandwidth
    end

    red:expire(key, 60)  -- 设置 60 秒自动过期
    return math.max(math.floor(max_bandwidth / count), min_bandwidth)
end

function _M.unregister_device(tenant_id)
    local red, err = get_redis()
    if not red then return end

    local key = redis_prefix .. tenant_id
    red:decr(key)
end

return _M