funstat-mcp-go/internal/ratelimit/limiter.go

package ratelimit

import (
	"context"
	"sync"
	"time"
)

type Limiter struct {
	maxRequests int
	window      time.Duration

	mu         sync.Mutex
	timestamps []time.Time
}

func New(maxRequests int, window time.Duration) *Limiter {
	if maxRequests <= 0 {
		maxRequests = 1
	}
	if window <= 0 {
		window = time.Second
	}

	return &Limiter{
		maxRequests: maxRequests,
		window:      window,
		timestamps:  make([]time.Time, 0, maxRequests),
	}
}

func (l *Limiter) Wait(ctx context.Context) error {
	for {
		l.mu.Lock()
		now := time.Now()

		cutoff := now.Add(-l.window)
		idx := 0
		for ; idx < len(l.timestamps); idx++ {
			if l.timestamps[idx].After(cutoff) {
				break
			}
		}

		if idx > 0 {
			l.timestamps = append([]time.Time(nil), l.timestamps[idx:]...)
		}

		if len(l.timestamps) < l.maxRequests {
			l.timestamps = append(l.timestamps, now)
			l.mu.Unlock()
			return nil
		}

		waitUntil := l.timestamps[0].Add(l.window)
		waitDuration := time.Until(waitUntil)
		l.mu.Unlock()

		if waitDuration <= 0 {
			continue
		}

		timer := time.NewTimer(waitDuration)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}
	}
}