Browse Source

feat: add mysql support

master
evanchen333 6 months ago
parent
commit
bce6dbb703
  1. 2
      Makefile
  2. 2
      README.md
  3. 29
      example/main.go
  4. 2
      go.mod
  5. 11
      go.sum
  6. 16
      limiter.go

2
Makefile

@ -2,4 +2,4 @@ run:
go run example/main.go
test:
go test
go test -v -cover=true

2
README.md

@ -1,4 +1,4 @@
# rate-limiter
# Rate Limiter
A go lib for limiting hit rate of a given ip within a given interval.
The counting of the rate is based on the hits within the past **interval**

29
example/main.go

@ -1,17 +1,40 @@
package main
import (
"database/sql"
"net/http"
"os"
"time"
_ "github.com/go-sql-driver/mysql"
limiter "github.com/mutsuki333/rate-limiter"
)
var l *limiter.Limiter
func main() {
l = limiter.Default()
// l.Interval = time.Minute * 10
l.Limit = 5
if len(os.Args) > 1 {
var db *sql.DB
var err error
if os.Args[1] == "mysql" {
db, err = sql.Open("mysql", "root:admin@/limiter?charset=utf8")
} else {
db, err = sql.Open("sqlite3", os.Args[1]+".db")
}
if err != nil {
panic(err)
}
l = &limiter.Limiter{
Interval: time.Minute,
Limit: 5,
Store: db,
}
l.Init()
} else {
l = limiter.Default()
l.Limit = 5
}
http.HandleFunc("/hit", l.Handler)
http.ListenAndServe(":8080", nil)

2
go.mod

@ -3,6 +3,6 @@ module github.com/mutsuki333/rate-limiter
go 1.15
require (
github.com/go-sql-driver/mysql v1.5.0
github.com/mattn/go-sqlite3 v1.14.6
gorm.io/driver/sqlite v1.1.4
)

11
go.sum

@ -1,11 +1,4 @@
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.1 h1:g39TucaRWyV3dwDO++eEc6qf8TVIQ/Da48WmqjZ3i7E=
github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/mattn/go-sqlite3 v1.14.5/go.mod h1:WVKg1VTActs4Qso6iwGbiFih2UIHo0ENGwNd0Lj+XmI=
github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
github.com/mattn/go-sqlite3 v1.14.6 h1:dNPt6NO46WmLVt2DLNpwczCmdV5boIZ6g/tlDrlRUbg=
github.com/mattn/go-sqlite3 v1.14.6/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
gorm.io/driver/sqlite v1.1.4 h1:PDzwYE+sI6De2+mxAneV9Xs11+ZyKV6oxD3wDGkaNvM=
gorm.io/driver/sqlite v1.1.4/go.mod h1:mJCeTFr7+crvS+TRnWc5Z3UvwxUN1BGBLMrf5LA9DYw=
gorm.io/gorm v1.20.7 h1:rMS4CL3pNmYq1V5/X+nHHjh1Dx6dnf27+Cai5zabo+M=
gorm.io/gorm v1.20.7/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw=

16
limiter.go

@ -45,7 +45,15 @@ type Limiter struct {
//Init the db table, and start cleanup goroutine
func (l *Limiter) Init() {
l.Store.Exec(`create table hit (ip string, hit_time datetime);`)
_, err := l.Store.Exec(`drop table if exists hit;`)
if err != nil {
panic(err)
}
_, err = l.Store.Exec(`create table hit (ip varchar(20), hit_time datetime);`)
if err != nil {
panic(err)
}
l.Mux = &sync.Mutex{}
go func() {
for {
time.Sleep(l.Interval)
@ -56,10 +64,10 @@ func (l *Limiter) Init() {
//Rate within interval
func (l *Limiter) Rate(ip string) (rate int, err error) {
l.Mux.Lock()
defer l.Mux.Unlock()
l.Mux.Lock()
defer l.Mux.Unlock()
row := l.Store.QueryRow(
"select COUNT(*) from hit where ip = ? and hit_time > ?",
"select count(*) from hit where ip = ? and hit_time > ?",
ip,
time.Now().Add(-l.Interval),
)

Loading…
Cancel
Save