Go로 PostgreSQL 프록시 만들기 (65) - Per-User / Per-IP Rate Limiting

쿼리 수신
  │
  ▼
[1. Global Rate Limit]  ← 전역 Token Bucket (기존)
  │ pass
  ▼
[2. Per-User Rate Limit] ← 사용자별 Token Bucket
  │ pass
  ▼
[3. Per-IP Rate Limit]   ← IP별 Token Bucket
  │ pass
  ▼
쿼리 처리 계속

Registry
  ├── "app_user"  → RateLimiter(rate=100, burst=50)
  ├── "admin_user" → RateLimiter(rate=200, burst=100)  ← override
  ├── "10.0.0.5"  → RateLimiter(rate=100, burst=50)
  └── (TTL 만료 시 자동 제거)

type entry struct {
    limiter    *RateLimiter
    lastAccess atomic.Int64 // unix nano
}

type RateLimiterRegistry struct {
    mu           sync.RWMutex
    limiters     map[string]*entry
    defaultRate  float64
    defaultBurst int
    overrides    map[string]Override
    cleanupTTL   time.Duration
    stopCleanup  chan struct{}
}

func (r *RateLimiterRegistry) Allow(key string) bool {
    // Fast path: read-lock lookup
    r.mu.RLock()
    e, ok := r.limiters[key]
    r.mu.RUnlock()

    if ok {
        e.touch()
        return e.limiter.Allow()
    }

    // Slow path: create new entry
    r.mu.Lock()
    // Double-check after acquiring write lock
    if e, ok = r.limiters[key]; ok {
        r.mu.Unlock()
        e.touch()
        return e.limiter.Allow()
    }

    rate, burst := r.defaultRate, r.defaultBurst
    if o, found := r.overrides[key]; found {
        rate, burst = o.Rate, o.Burst
    }

    e = &entry{limiter: NewRateLimiter(rate, burst)}
    e.touch()
    r.limiters[key] = e
    r.mu.Unlock()

    return e.limiter.Allow()
}

func (r *RateLimiterRegistry) cleanupLoop() {
    ticker := time.NewTicker(r.cleanupTTL / 2)
    defer ticker.Stop()

    for {
        select {
        case <-r.stopCleanup:
            return
        case <-ticker.C:
            r.evict()
        }
    }
}

func (r *RateLimiterRegistry) evict() {
    cutoff := time.Now().Add(-r.cleanupTTL).UnixNano()
    r.mu.Lock()
    for key, e := range r.limiters {
        if e.lastAccess.Load() < cutoff {
            delete(r.limiters, key)
        }
    }
    r.mu.Unlock()
}

type Server struct {
    rateLimitPtr    atomic.Pointer[resilience.RateLimiter]       // 전역
    perUserRLPtr    atomic.Pointer[resilience.RateLimiterRegistry] // 사용자별
    perIPRLPtr      atomic.Pointer[resilience.RateLimiterRegistry] // IP별
    // ...
}

// Rate limit check — global
if rl := s.getRateLimiter(); rl != nil && !rl.Allow() {
    slog.Warn("rate limited", "scope", "global", "remote", clientConn.RemoteAddr())
    if s.metrics != nil {
        s.metrics.RateLimited.WithLabelValues("global").Inc()
    }
    s.sendError(clientConn, "too many requests")
    _ = protocol.WriteMessage(clientConn, protocol.MsgReadyForQuery, []byte{'I'})
    continue
}

// Rate limit check — per-user
if reg := s.perUserRLPtr.Load(); reg != nil && !reg.Allow(si.User) {
    slog.Warn("rate limited", "scope", "user", "user", si.User, ...)
    // ...
    continue
}

// Rate limit check — per-IP
if reg := s.perIPRLPtr.Load(); reg != nil {
    clientIP := extractIP(clientConn.RemoteAddr().String())
    if !reg.Allow(clientIP) {
        slog.Warn("rate limited", "scope", "ip", "ip", clientIP, ...)
        // ...
        continue
    }
}

// Reload에서 per-user registry 교체
if newCfg.RateLimit.PerUser.Enabled {
    if old := s.perUserRLPtr.Swap(
        resilience.NewRateLimiterRegistry(buildRegistryConfig(...)),
    ); old != nil {
        old.Close() // cleanup 고루틴 종료
    }
} else {
    if old := s.perUserRLPtr.Swap(nil); old != nil {
        old.Close()
    }
}

rate_limit:
  enabled: true
  rate: 1000                        # 전역 초당 허용 쿼리 수
  burst: 100                        # 전역 최대 버스트

  per_user:
    enabled: true
    default_rate: 100               # 사용자별 기본 초당 허용
    default_burst: 50               # 사용자별 기본 버스트
    overrides:
      - key: "batch_worker"         # 배치 워커에 높은 한도
        rate: 500
        burst: 200
      - key: "readonly_dashboard"   # 대시보드에 낮은 한도
        rate: 20
        burst: 10

  per_ip:
    enabled: true
    default_rate: 100
    default_burst: 50
    overrides:
      - key: "10.0.0.100"          # 내부 서비스 IP에 높은 한도
        rate: 500
        burst: 200