Skip to content

Commit 9d06c64

Browse files
dhsifssdhsifss
authored
fix: ratelimit impl (#56)
* fix: ratelimit impl
1 parent 57bbb1f commit 9d06c64

File tree

7 files changed

+81
-13
lines changed

7 files changed

+81
-13
lines changed

‎backend/cmd/server/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/chaitin/koalaqa/pkg/mq"
1616
"github.com/chaitin/koalaqa/pkg/oss"
1717
"github.com/chaitin/koalaqa/pkg/rag"
18+
"github.com/chaitin/koalaqa/pkg/ratelimit"
1819
"github.com/chaitin/koalaqa/pkg/third_auth"
1920
"github.com/chaitin/koalaqa/pkg/version"
2021
"github.com/chaitin/koalaqa/pkg/webhook"
@@ -48,6 +49,7 @@ func main() {
4849
third_auth.Module,
4950
cron.Module(),
5051
fx.Provide(version.NewInfo),
52+
ratelimit.Module,
5153
)
5254
ctx := context.Background()
5355
if err := app.Start(ctx); err != nil {

‎backend/go.mod‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ module github.com/chaitin/koalaqa
33
go 1.25.0
44

55
require (
6-
github.com/Narasimha1997/ratelimiter v1.1.1
76
github.com/caarlos0/env/v11 v11.3.1
87
github.com/chaitin/ModelKit v1.8.3
98
github.com/chaitin/pandawiki/sdk/rag v0.0.0-20250927130416-bcfc4bde3379
@@ -27,6 +26,7 @@ require (
2726
go.uber.org/fx v1.24.0
2827
golang.org/x/crypto v0.41.0
2928
golang.org/x/oauth2 v0.30.0
29+
golang.org/x/time v0.12.0
3030
gorm.io/driver/postgres v1.6.0
3131
gorm.io/gorm v1.30.1
3232
)
@@ -135,7 +135,6 @@ require (
135135
golang.org/x/sync v0.16.0 // indirect
136136
golang.org/x/sys v0.35.0 // indirect
137137
golang.org/x/text v0.28.0 // indirect
138-
golang.org/x/time v0.12.0 // indirect
139138
golang.org/x/tools v0.36.0 // indirect
140139
google.golang.org/api v0.239.0 // indirect
141140
google.golang.org/genai v1.13.0 // indirect

‎backend/go.sum‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@ github.com/JohannesKaufmann/html-to-markdown/v2 v2.3.3 h1:r3fokGFRDk/8pHmwLwJ8zs
1616
github.com/JohannesKaufmann/html-to-markdown/v2 v2.3.3/go.mod h1:HtsP+1Fchp4dVvaiIsLHAl/yqL3H1YLwqLC9kNwqQEg=
1717
github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc=
1818
github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=
19-
github.com/Narasimha1997/ratelimiter v1.1.1 h1:ndkK0dHqUKdwSElE8Kghz+0gVcGEa9q6/CisEL/h6HU=
20-
github.com/Narasimha1997/ratelimiter v1.1.1/go.mod h1:TCsPmcx5vkQJu64sbTLRcr8xpNNmO22OTnvhfXEWoNw=
2119
github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o=
2220
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
2321
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=

‎backend/pkg/ratelimit/fx.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package ratelimit
2+
3+
import "go.uber.org/fx"
4+
5+
var Module = fx.Options(
6+
fx.Provide(New),
7+
)

‎backend/pkg/ratelimit/ratelimit.go‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package ratelimit
2+
3+
import (
4+
"sync"
5+
"time"
6+
7+
"golang.org/x/time/rate"
8+
)
9+
10+
type Limiter interface {
11+
Allow(key string, period time.Duration, num int) bool
12+
}
13+
14+
type limiter struct {
15+
l *rate.Limiter
16+
lastUse time.Time
17+
}
18+
19+
func newLimiter(period time.Duration, num int) *limiter {
20+
return &limiter{
21+
l: rate.NewLimiter(rate.Every(period), num),
22+
lastUse: time.Now(),
23+
}
24+
}
25+
26+
type multiLimiter struct {
27+
l *sync.Map
28+
}
29+
30+
func (l *multiLimiter) Allow(key string, period time.Duration, num int) bool {
31+
data, _ := l.l.LoadOrStore(key, newLimiter(period, num))
32+
cacheLimiter := data.(*limiter)
33+
cacheLimiter.lastUse = time.Now()
34+
35+
return cacheLimiter.l.Allow()
36+
}
37+
38+
func (l *multiLimiter) clearLoop() {
39+
for {
40+
time.Sleep(time.Hour)
41+
42+
now := time.Now()
43+
44+
l.l.Range(func(key, value any) bool {
45+
li := value.(*limiter)
46+
47+
if li.lastUse.Before(now.Add(-time.Hour)) {
48+
l.l.Delete(key)
49+
}
50+
51+
return true
52+
})
53+
}
54+
}
55+
56+
func New() Limiter {
57+
l := &multiLimiter{
58+
l: &sync.Map{},
59+
}
60+
61+
go l.clearLoop()
62+
return l
63+
}

‎backend/svc/discussion.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,12 @@ import (
99
"path/filepath"
1010
"time"
1111

12-
"github.com/Narasimha1997/ratelimiter"
1312
"github.com/chaitin/koalaqa/model"
1413
"github.com/chaitin/koalaqa/pkg/glog"
1514
"github.com/chaitin/koalaqa/pkg/mq"
1615
"github.com/chaitin/koalaqa/pkg/oss"
1716
"github.com/chaitin/koalaqa/pkg/rag"
17+
"github.com/chaitin/koalaqa/pkg/ratelimit"
1818
"github.com/chaitin/koalaqa/pkg/topic"
1919
"github.com/chaitin/koalaqa/pkg/util"
2020
"github.com/chaitin/koalaqa/pkg/webhook/message"
@@ -36,14 +36,14 @@ type discussionIn struct {
3636
Dataset *repo.Dataset
3737
OC oss.Client
3838
ForumRepo *repo.Forum
39+
Limiter ratelimit.Limiter
3940
}
4041

4142
type Discussion struct {
4243
in discussionIn
4344

4445
logger *glog.Logger
4546
webhookType map[model.DiscussionType]message.Type
46-
limiter *ratelimiter.AttributeBasedLimiter
4747
}
4848

4949
func newDiscussion(in discussionIn) *Discussion {
@@ -55,7 +55,6 @@ func newDiscussion(in discussionIn) *Discussion {
5555
model.DiscussionTypeFeedback: message.TypeNewFeedback,
5656
model.DiscussionTypeBlog: message.TypeNewBlog,
5757
},
58-
limiter: ratelimiter.NewAttributeBasedLimiter(false),
5958
}
6059
}
6160

@@ -91,7 +90,7 @@ func (d *Discussion) limitKey(args ...any) string {
9190
}
9291

9392
func (d *Discussion) allow(args ...any) bool {
94-
return d.limiter.MustShouldAllow(d.limitKey(args...), 1, 3, time.Minute)
93+
return d.in.Limiter.Allow(d.limitKey(args...), time.Second*20, 3)
9594
}
9695

9796
var errRatelimit = errors.New("ratelimit")

‎backend/svc/webhook.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@ import (
66
"sync"
77
"time"
88

9-
"github.com/Narasimha1997/ratelimiter"
109
"github.com/chaitin/koalaqa/model"
1110
"github.com/chaitin/koalaqa/pkg/glog"
11+
"github.com/chaitin/koalaqa/pkg/ratelimit"
1212
"github.com/chaitin/koalaqa/pkg/webhook"
1313
"github.com/chaitin/koalaqa/pkg/webhook/message"
1414
"github.com/chaitin/koalaqa/repo"
@@ -19,7 +19,7 @@ type Webhook struct {
1919
logger *glog.Logger
2020
lock sync.Mutex
2121
webhooks map[uint]webhook.Webhook
22-
limiter *ratelimiter.AttributeBasedLimiter
22+
limiter ratelimit.Limiter
2323

2424
repoWebhook *repo.Webhook
2525
}
@@ -125,7 +125,7 @@ func (w *Webhook) Delete(ctx context.Context, id uint) error {
125125
}
126126

127127
func (w *Webhook) allow(id uint, discID string, msgType message.Type) bool {
128-
return w.limiter.MustShouldAllow(fmt.Sprintf("webhook-%d-%s-%d", id, discID, msgType), 1, 1, time.Minute*5)
128+
return w.limiter.Allow(fmt.Sprintf("webhook-%d-%s-%d", id, discID, msgType), time.Minute*5, 1)
129129
}
130130

131131
func (w *Webhook) Send(ctx context.Context, msg message.Message) error {
@@ -143,12 +143,12 @@ func (w *Webhook) Send(ctx context.Context, msg message.Message) error {
143143
return nil
144144
}
145145

146-
func newWebhook(lc fx.Lifecycle, repoWebhook *repo.Webhook) *Webhook {
146+
func newWebhook(lc fx.Lifecycle, repoWebhook *repo.Webhook, limiter ratelimit.Limiter) *Webhook {
147147
w := &Webhook{
148148
logger: glog.Module("svc", "webhook"),
149149
repoWebhook: repoWebhook,
150150
webhooks: make(map[uint]webhook.Webhook),
151-
limiter: ratelimiter.NewAttributeBasedLimiter(false),
151+
limiter: limiter,
152152
}
153153
lc.Append(fx.Hook{
154154
OnStart: func(ctx context.Context) error {

0 commit comments

Comments
 (0)