請使用 go 實現一個 rate limiter
package main
import (
"encoding/json"
"fmt"
"net/http"
"sync"
"time"
)
type RateLimit struct {
MaxRequests int `json:"maxRequests"` // 時間範圍內的最大請求數
Duration time.Duration `json:"duration"` // 時間範圍
}
type SlidingWindow struct {
sync.Mutex
Requests []time.Time // 記錄請求時間的slice
Capacity int // 窗口容量(最大請求數)
Duration time.Duration // 時間範圍
}
var rateLimits = &sync.Map{}
var windows = &sync.Map{}
func setRateLimitHandler(w http.ResponseWriter, r *http.Request) {
var rateLimit RateLimit
err := json.NewDecoder(r.Body).Decode(&rateLimit)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
resource := r.URL.Query().Get("resource")
rateLimits.Store(resource, rateLimit)
window := &SlidingWindow{
Requests: make([]time.Time, 0, rateLimit.MaxRequests),
Capacity: rateLimit.MaxRequests,
Duration: rateLimit.Duration,
}
windows.Store(resource, window)
w.WriteHeader(http.StatusOK)
}
func proxyHandler(w http.ResponseWriter, r *http.Request) {
resource := r.URL.Path[len("/api/resource/"):]
value, ok := rateLimits.Load(resource)
if !ok {
http.Error(w, "Resource not found", http.StatusNotFound)
return
}
rateLimit := value.(RateLimit)
window, _ := windows.LoadOrStore(resource, &SlidingWindow{
Requests: make([]time.Time, 0, rateLimit.MaxRequests),
Capacity: rateLimit.MaxRequests,
Duration: rateLimit.Duration,
})
sw := window.(*SlidingWindow)
sw.Lock()
defer sw.Unlock()
now := time.Now()
sw.cleanupExpiredRequests(now)
if len(sw.Requests) >= sw.Capacity {
fmt.Fprintf(w, "Waiting for request slot...")
sw.waitForSlot(now)
sw.Requests = append(sw.Requests, now)
} else {
sw.Requests = append(sw.Requests, now)
}
// 在這裡你可以轉發請求給後端資源
fmt.Fprintf(w, "Proxying request to resource: %s", resource)
}
func (sw *SlidingWindow) cleanupExpiredRequests(now time.Time) {
for i := 0; i < len(sw.Requests); i++ {
if sw.Requests[i].Add(sw.Duration).Before(now) {
sw.Requests = sw.Requests[i+1:]
return
}
}
}
func (sw *SlidingWindow) waitForSlot(now time.Time) {
for i := 0; i < len(sw.Requests); i++ {
if sw.Requests[i].Add(sw.Duration).After(now) {
time.Sleep(sw.Requests[i].Add(sw.Duration).Sub(now))
return
}
}
}
func main() {
http.HandleFunc("/api/rate-limits", setRateLimitHandler)
http.HandleFunc("/api/resource/", proxyHandler)
http.ListenAndServe(":8080", nil)
}
