基于PostgreSQL存储过程限制特定用户活跃会话数的实现问询

针对你需要限制异步用户最多10个活跃会话、超出则等待的需求，可以通过自定义存储过程+会话检查逻辑实现，核心是利用PostgreSQL的pg_stat_activity监控会话状态，结合原子性锁避免竞态条件，以下是具体实现思路和代码：

核心逻辑要点 #

这类存储过程需要覆盖以下关键逻辑：

• 精准识别目标会话：通过用户名、application_name等标识筛选异步用户的会话
• 实时统计活跃数：排除当前会话，避免统计干扰
• 原子性检查与等待：用数据库 advisory lock 确保统计和等待逻辑的原子性，防止多个会话同时突破阈值
• 超时保护：设置最大等待时长，避免会话无限阻塞
• 安全释放资源：确保锁最终被释放，避免死锁

完整实现代码 #

1. 统计目标用户活跃会话数函数 #

先创建一个辅助函数，用于统计异步用户的当前活跃会话数（排除当前执行统计的会话）：

CREATE OR REPLACE FUNCTION count_active_async_sessions()
RETURNS INTEGER AS $$
BEGIN
  RETURN (
    SELECT COUNT(*)
    FROM pg_stat_activity
    -- 替换为你的异步用户列表，或用application_name筛选：application_name = 'async_application'
    WHERE usename IN ('async_user_1', 'async_user_2')
      AND state = 'active'
      AND pid != pg_backend_pid() -- 排除当前会话，避免统计干扰
  );
END;
$$ LANGUAGE plpgsql STABLE;

2. 会话限制存储过程 #

创建核心的限制逻辑函数，实现活跃数检查、等待与超时控制：

CREATE OR REPLACE FUNCTION enforce_async_session_limit()
RETURNS VOID AS $$
DECLARE
  max_active_sessions INTEGER := 10; -- 异步用户最大活跃会话数
  current_active INTEGER;
  wait_seconds INTEGER := 1; -- 每次等待间隔（秒）
  max_wait_minutes INTEGER := 5; -- 最大等待时长（分钟）
  start_wait_time TIMESTAMP := NOW();
  lock_key BIGINT := 12345; -- 自定义唯一锁键，用于控制并发检查
BEGIN
  -- 获取排他锁，确保统计和等待逻辑的原子性，防止竞态条件
  PERFORM pg_advisory_lock(lock_key);
  
  BEGIN
    LOOP
      -- 实时获取当前活跃会话数
      current_active := count_active_async_sessions();
      
      -- 活跃数未达阈值，退出等待
      IF current_active < max_active_sessions THEN
        EXIT;
      END IF;
      
      -- 检查是否超过最大等待时长，超时则抛出异常
      IF NOW() - start_wait_time > INTERVAL '1 minute' * max_wait_minutes THEN
        RAISE EXCEPTION 'Async session limit exceeded: max wait time reached (5 minutes)';
      END IF;
      
      -- 等待指定时长后重新检查
      PERFORM pg_sleep(wait_seconds);
    END LOOP;
  FINALLY
    -- 无论是否异常，确保释放锁
    PERFORM pg_advisory_unlock(lock_key);
  END;
END;
$$ LANGUAGE plpgsql;

3. 触发逻辑 #

让异步应用在获取连接后、执行业务查询前调用该存储过程：

SELECT enforce_async_session_limit();

如果应用无法修改代码，也可以通过设置用户的session_preload_libraries或使用pg_cron定时清理，但应用主动调用是最可靠的方式。

关键注意事项 #

• 权限配置：确保异步用户拥有SELECT ON pg_stat_activity的权限，或给存储过程添加SECURITY DEFINER（需注意安全风险，限制执行权限）
• 竞态条件避免：必须使用pg_advisory_lock确保统计和等待的原子性，否则多个会话可能同时突破阈值
• 会话状态识别：pg_stat_activity的state字段需注意区分active（正在执行查询）和idle in transaction（事务中但空闲），可根据业务需求调整筛选条件
• 连接池适配：由于等待的会话会占用连接池资源，需确保异步连接池的大小足够容纳等待的会话（比如200个连接池可以覆盖等待需求）

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