top | item 41525690

(no title)

deisteve | 1 year ago

i've written an implementation on postgresql

you can use it like this:

     SELECT * FROM check_rate_limit('client1', 1.0, 10.0, INTERVAL '1 minute');
-- Create a table to store client rate limiting data

CREATE TABLE rate_limiter ( client_id VARCHAR(255) PRIMARY KEY, last_update_time TIMESTAMP WITH TIME ZONE NOT NULL, average_rate DOUBLE PRECISION NOT NULL );

-- Function to check and update rate limit

CREATE OR REPLACE FUNCTION check_rate_limit( client_id VARCHAR(255), cost DOUBLE PRECISION DEFAULT 1.0, limit_value DOUBLE PRECISION DEFAULT 600.0, period INTERVAL DEFAULT INTERVAL '1 hour' ) RETURNS TABLE ( allowed BOOLEAN, next_allowed_time TIMESTAMP WITH TIME ZONE ) AS $$ DECLARE current_time TIMESTAMP WITH TIME ZONE := NOW(); time_interval DOUBLE PRECISION; alpha DOUBLE PRECISION; instantaneous_rate DOUBLE PRECISION; new_rate DOUBLE PRECISION; client_data RECORD; BEGIN

    -- Get client data or use defaults if not exists

    SELECT \* INTO client_data
    FROM rate_limiter
    WHERE rate_limiter.client_id = check_rate_limit.client_id;

    IF NOT FOUND THEN
        client_data := (client_id, current_time - period, 0.0)::rate_limiter;
    END IF;

    -- Calculate interval

    time_interval := EXTRACT(EPOCH FROM (current_time - client_data.last_update_time)) / EXTRACT(EPOCH FROM period);
    time_interval := GREATEST(time_interval, 1.0e-10);

    -- Calculate alpha (exponential smoothing weight)

    alpha := EXP(-time_interval);

    -- Calculate instantaneous rate

    instantaneous_rate := cost / time_interval;

    -- Calculate new average rate

    new_rate := (1 - alpha) * instantaneous_rate + alpha * client_data.average_rate;

    -- Ensure rare requests are counted in full

    new_rate := GREATEST(new_rate, cost);

    -- Check if rate limit is exceeded

    IF new_rate > limit_value THEN
        -- Calculate next allowed time
        next_allowed_time := current_time + (period * LN(new_rate / limit_value));
        allowed := FALSE;
    ELSE

        -- Update client data

        INSERT INTO rate_limiter (client_id, last_update_time, average_rate)
        VALUES (client_id, current_time, new_rate)
        ON CONFLICT (client_id) DO UPDATE
        SET last_update_time = EXCLUDED.last_update_time,
            average_rate = EXCLUDED.average_rate;

        next_allowed_time := current_time;
        allowed := TRUE;
    END IF;

    RETURN NEXT;
END; $$ LANGUAGE plpgsql;

give that a whirl

discuss

order

No comments yet.