From 747d087007dc8cf66e1b5e48b58c5a24860452ad Mon Sep 17 00:00:00 2001
From: xhe <xw897002528@gmail.com>
Date: Mon, 12 Dec 2022 15:53:08 +0800
Subject: [PATCH] server: default rate limit (#148)

---
 go.mod               | 2 ++
 go.sum               | 4 ++++
 pkg/server/server.go | 8 ++++++++
 3 files changed, 14 insertions(+)

diff --git a/go.mod b/go.mod
index 905841e0..e1bddd55 100644
--- a/go.mod
+++ b/go.mod
@@ -19,6 +19,7 @@ require (
 	go.etcd.io/etcd/client/v3 v3.5.6
 	go.etcd.io/etcd/server/v3 v3.5.6
 	go.uber.org/atomic v1.10.0
+	go.uber.org/ratelimit v0.2.0
 	go.uber.org/zap v1.23.0
 	golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0
 	google.golang.org/grpc v1.49.0
@@ -26,6 +27,7 @@ require (
 
 require (
 	github.com/BurntSushi/toml v1.2.0 // indirect
+	github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 // indirect
 	github.com/benbjohnson/clock v1.3.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.1.1 // indirect
diff --git a/go.sum b/go.sum
index 48330cb7..babdb9bb 100644
--- a/go.sum
+++ b/go.sum
@@ -60,6 +60,8 @@ github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRF
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
 github.com/aliyun/alibaba-cloud-sdk-go v1.61.1581 h1:Q/yk4z/cHUVZfgTqtD09qeYBxHwshQAjVRX73qs8UH0=
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 h1:MzBOUgng9orim59UnfUTLRjMpd09C5uEVQ6RPGeCaVI=
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129/go.mod h1:rFgpPQZYZ8vdbc+48xibu8ALc3yeyd64IhHS+PU6Yyg=
 github.com/antihax/optional v0.0.0-20180407024304-ca021399b1a6/go.mod h1:V8iCPQYkqmusNa815XgQio277wI47sdRh1dUOLdyC6Q=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
@@ -723,6 +725,8 @@ go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9i
 go.uber.org/multierr v1.7.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
 go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8=
 go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=
+go.uber.org/ratelimit v0.2.0 h1:UQE2Bgi7p2B85uP5dC2bbRtig0C+OeNRnNEafLjsLPA=
+go.uber.org/ratelimit v0.2.0/go.mod h1:YYBV4e4naJvhpitQrWJu1vCpgB7CboMe0qhltKt6mUg=
 go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
 go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
diff --git a/pkg/server/server.go b/pkg/server/server.go
index df117b23..f88057ef 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -37,9 +37,15 @@ import (
 	"github.com/pingcap/TiProxy/pkg/server/api"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.uber.org/atomic"
+	"go.uber.org/ratelimit"
 	"go.uber.org/zap"
 )
 
+const (
+	// DefAPILimit is the global API limit per second.
+	DefAPILimit = 100
+)
+
 type Server struct {
 	wg waitgroup.WaitGroup
 	// managers
@@ -94,6 +100,7 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error)
 		slogger := lg.Named("gin")
 		gin.SetMode(gin.ReleaseMode)
 		engine := gin.New()
+		limit := ratelimit.New(DefAPILimit)
 		engine.Use(
 			gin.Recovery(),
 			ginzap.Ginzap(slogger, "", true),
@@ -102,6 +109,7 @@ func NewServer(ctx context.Context, sctx *sctx.Context) (srv *Server, err error)
 					c.Abort()
 					c.JSON(http.StatusInternalServerError, "service not ready")
 				}
+				_ = limit.Take()
 			},
 		)