Skip to content

Instantly share code, notes, and snippets.

@dedalqq
Last active May 24, 2023 06:06
Show Gist options
  • Save dedalqq/78505b481eb6adc998e71169c629e279 to your computer and use it in GitHub Desktop.
Save dedalqq/78505b481eb6adc998e71169c629e279 to your computer and use it in GitHub Desktop.
package main
import (
"fmt"
"log"
"net/http"
"os"
"sync"
"time"
)
// queue реализует циклическую очередь
type queue[T any] struct {
data []T
head, tail int
}
// newQueue возвращает новую очередь на n элементов
func newQueue[T any](n int) *queue[T] {
return &queue[T]{
data: make([]T, n+1),
}
}
// size возвращает количество элементов в очереди
func (q *queue[T]) size() int {
if q.head > q.tail {
return len(q.data) - q.head + q.tail
}
return q.tail - q.head
}
// push пытается добавить элемент если очередь не заполнена и возвращает true если удалось добавить
func (q *queue[T]) push(v T) bool {
if q.size()+1 < len(q.data) {
q.data[q.tail] = v
q.tail = (q.tail + 1) % len(q.data)
return true
}
return false
}
// first позволяет получить первый элемент в очереди не извлекая его
func (q *queue[T]) first() T {
if q.head == q.tail {
panic("")
}
return q.data[q.head]
}
// pop извлекает первый элемент
func (q *queue[T]) pop() T {
if q.head == q.tail {
v := new(T)
return *v
}
value := q.data[q.head]
q.head = (q.head + 1) % len(q.data)
return value
}
type RateLimiter struct {
N int
mx sync.Mutex
q *queue[time.Time]
}
func NewRateLimiter(n int) *RateLimiter {
return &RateLimiter{
N: n,
q: newQueue[time.Time](n),
}
}
func (rl *RateLimiter) check() bool {
rl.mx.Lock()
defer rl.mx.Unlock()
now := time.Now()
if ok := rl.q.push(now); ok {
return true
}
if rl.q.first().Add(time.Second).After(now) {
return false
}
rl.q.pop()
rl.q.push(now)
return true
}
type handlerFunc func(w http.ResponseWriter, r *http.Request)
func (rl *RateLimiter) Middleware(h handlerFunc) handlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if rl.check() {
h(w, r)
return
}
fmt.Println("REQUEST FAILED")
w.WriteHeader(http.StatusTooManyRequests)
}
}
// ################### код ниже не править ######################
func handler(w http.ResponseWriter, r *http.Request) {
fmt.Println("REQUEST COMPLETED")
}
func main() {
rl := NewRateLimiter(10)
http.HandleFunc("/", rl.Middleware(handler))
go func() {
for i := 0; i < 1000; i++ {
time.Sleep(10 * time.Millisecond)
if _, err := http.Get("http://localhost:8080/"); err != nil {
log.Fatal(err)
}
}
os.Exit(0)
}()
err := http.ListenAndServe(":8080", nil)
log.Fatal(err)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment