Rate Limiting

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"
    "sync"
    "time"
)

// 1. Simple rate limiter
type RateLimiter struct {
    rate       int           // requests per second
    burst      int           // maximum burst
    tokens     int           // current tokens
    lastUpdate time.Time
    mu         sync.Mutex
}

func NewRateLimiter(rate, burst int) *RateLimiter {
    return &RateLimiter{
        rate:       rate,
        burst:      burst,
        tokens:     burst,
        lastUpdate: time.Now(),
    }
}

func (r *RateLimiter) Allow() bool {
    r.mu.Lock()
    defer r.mu.Unlock()

    now := time.Now()
    elapsed := now.Sub(r.lastUpdate)
    r.lastUpdate = now

    // Add new tokens based on elapsed time
    newTokens := float64(elapsed.Nanoseconds()) *
        float64(r.rate) / float64(time.Second)
    r.tokens += int(newTokens)

    // Cap tokens at burst limit
    if r.tokens > r.burst {
        r.tokens = r.burst
    }

    // Check if we can allow the request
    if r.tokens > 0 {
        r.tokens--
        return true
    }

    return false
}

// 2. Per-client rate limiter
type ClientRateLimiter struct {
    limiters map[string]*RateLimiter
    rate     int
    burst    int
    mu       sync.RWMutex
}

func NewClientRateLimiter(rate, burst int) *ClientRateLimiter {
    return &ClientRateLimiter{
        limiters: make(map[string]*RateLimiter),
        rate:     rate,
        burst:    burst,
    }
}

func (c *ClientRateLimiter) Allow(clientID string) bool {
    c.mu.Lock()
    limiter, exists := c.limiters[clientID]
    if !exists {
        limiter = NewRateLimiter(c.rate, c.burst)
        c.limiters[clientID] = limiter
    }
    c.mu.Unlock()

    return limiter.Allow()
}

// 3. Sliding window rate limiter
type SlidingWindowLimiter struct {
    windowSize time.Duration
    limit      int
    requests   map[string][]time.Time
    mu         sync.Mutex
}

func NewSlidingWindowLimiter(windowSize time.Duration,
    limit int) *SlidingWindowLimiter {
    return &SlidingWindowLimiter{
        windowSize: windowSize,
        limit:      limit,
        requests:   make(map[string][]time.Time),
    }
}

func (s *SlidingWindowLimiter) Allow(clientID string) bool {
    s.mu.Lock()
    defer s.mu.Unlock()

    now := time.Now()
    windowStart := now.Add(-s.windowSize)

    // Get requests for this client
    times := s.requests[clientID]
    if times == nil {
        times = make([]time.Time, 0)
    }

    // Remove old requests
    valid := 0
    for _, t := range times {
        if t.After(windowStart) {
            times[valid] = t
            valid++
        }
    }
    times = times[:valid]

    // Check if we can allow the request
    if len(times) >= s.limit {
        return false
    }

    // Add new request
    times = append(times, now)
    s.requests[clientID] = times

    return true
}

// 4. HTTP middleware for rate limiting
func RateLimitMiddleware(limiter *ClientRateLimiter) func(
    http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter,
            r *http.Request) {
            // Get client ID (using IP address as example)
            clientID := r.RemoteAddr

            if !limiter.Allow(clientID) {
                http.Error(w,
                    "Rate limit exceeded",
                    http.StatusTooManyRequests)
                return
            }

            next.ServeHTTP(w, r)
        })
    }
}

// 5. Concurrent task rate limiter
type TaskRateLimiter struct {
    tokens  chan struct{}
    timeout time.Duration
}

func NewTaskRateLimiter(rate int,
    timeout time.Duration) *TaskRateLimiter {
    tokens := make(chan struct{}, rate)
    
    // Fill tokens
    for i := 0; i < rate; i++ {
        tokens <- struct{}{}
    }

    // Refill tokens periodically
    go func() {
        ticker := time.NewTicker(time.Second)
        defer ticker.Stop()

        for range ticker.C {
            for i := len(tokens); i < rate; i++ {
                tokens <- struct{}{}
            }
        }
    }()

    return &TaskRateLimiter{
        tokens:  tokens,
        timeout: timeout,
    }
}

func (t *TaskRateLimiter) Do(ctx context.Context,
    task func()) error {
    select {
    case <-t.tokens:
        task()
        return nil
    case <-time.After(t.timeout):
        return fmt.Errorf("task timed out waiting for rate limit")
    case <-ctx.Done():
        return ctx.Err()
    }
}

func main() {
    // Example 1: Simple rate limiter
    fmt.Println("Simple Rate Limiter Example:")
    limiter := NewRateLimiter(2, 3) // 2 requests/second, burst of 3
    
    for i := 0; i < 5; i++ {
        allowed := limiter.Allow()
        fmt.Printf("Request %d allowed: %v\n", i+1, allowed)
        time.Sleep(300 * time.Millisecond)
    }

    // Example 2: Per-client rate limiter
    fmt.Println("\nPer-Client Rate Limiter Example:")
    clientLimiter := NewClientRateLimiter(2, 3)
    
    clients := []string{"client1", "client2"}
    for _, client := range clients {
        for i := 0; i < 3; i++ {
            allowed := clientLimiter.Allow(client)
            fmt.Printf("%s request %d allowed: %v\n",
                client, i+1, allowed)
        }
    }

    // Example 3: Sliding window rate limiter
    fmt.Println("\nSliding Window Rate Limiter Example:")
    windowLimiter := NewSlidingWindowLimiter(time.Second, 3)
    
    for i := 0; i < 5; i++ {
        allowed := windowLimiter.Allow("client1")
        fmt.Printf("Request %d allowed: %v\n", i+1, allowed)
        time.Sleep(200 * time.Millisecond)
    }

    // Example 4: HTTP server with rate limiting
    fmt.Println("\nHTTP Server with Rate Limiting Example:")
    
    // Create handler
    handler := http.HandlerFunc(func(w http.ResponseWriter,
        r *http.Request) {
        fmt.Fprintf(w, "Hello, World!")
    })

    // Create server
    server := &http.Server{
        Addr:    ":8080",
        Handler: RateLimitMiddleware(clientLimiter)(handler),
    }

    // Start server in goroutine
    go func() {
        if err := server.ListenAndServe(); err != http.ErrServerClosed {
            log.Printf("HTTP server error: %v\n", err)
        }
    }()

    fmt.Println("HTTP server started on :8080")

    // Example 5: Task rate limiter
    fmt.Println("\nTask Rate Limiter Example:")
    taskLimiter := NewTaskRateLimiter(2,
        500*time.Millisecond) // 2 tasks/second
    
    ctx := context.Background()
    for i := 0; i < 5; i++ {
        err := taskLimiter.Do(ctx, func() {
            fmt.Printf("Executing task %d\n", i+1)
        })
        if err != nil {
            fmt.Printf("Task %d error: %v\n", i+1, err)
        }
        time.Sleep(300 * time.Millisecond)
    }

    // Cleanup
    time.Sleep(time.Second)
    server.Shutdown(context.Background())
}