Add rate limiter

This commit is contained in:
Balakrishnan Balasubramanian 2022-05-02 13:24:29 -04:00
parent 8637252785
commit ea7afd951e
3 changed files with 25 additions and 3 deletions

1
go.mod
View File

@ -7,6 +7,7 @@ require (
github.com/stretchr/testify v1.7.1 github.com/stretchr/testify v1.7.1
go.uber.org/zap v1.21.0 go.uber.org/zap v1.21.0
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4
golang.org/x/time v0.0.0-20220411224347-583f2d630306
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
) )

2
go.sum
View File

@ -47,6 +47,8 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/time v0.0.0-20220411224347-583f2d630306 h1:+gHMid33q6pen7kv9xvT+JRinntgeXO2AeZVd0AWD3w=
golang.org/x/time v0.0.0-20220411224347-583f2d630306/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=

View File

@ -2,14 +2,17 @@ package telegram
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"net/url" "net/url"
"time"
"go.balki.me/tss/log" "go.balki.me/tss/log"
"golang.org/x/time/rate"
) )
func main() { func main() {
@ -23,6 +26,8 @@ type TelegramSender interface {
type telegramSender struct { type telegramSender struct {
client *http.Client client *http.Client
authToken string authToken string
rateLimiterPerMin *rate.Limiter
rateLimiterPerSec *rate.Limiter
} }
func (ts *telegramSender) SendLink(link string, channel string, rhash string) error { func (ts *telegramSender) SendLink(link string, channel string, rhash string) error {
@ -42,6 +47,17 @@ func (ts *telegramSender) SendLink(link string, channel string, rhash string) er
} }
apiUrl := fmt.Sprintf("https://api.telegram.org/bot%s/sendMessage", ts.authToken) apiUrl := fmt.Sprintf("https://api.telegram.org/bot%s/sendMessage", ts.authToken)
err = ts.rateLimiterPerMin.Wait(context.TODO())
if err != nil {
return err
}
err = ts.rateLimiterPerSec.Wait(context.TODO())
if err != nil {
return err
}
res, err := ts.client.Post(apiUrl, "application/json", bytes.NewReader(data)) res, err := ts.client.Post(apiUrl, "application/json", bytes.NewReader(data))
if err != nil { if err != nil {
return err return err
@ -60,10 +76,13 @@ func (ts *telegramSender) SendLink(link string, channel string, rhash string) er
} }
func NewTelegramSender(transport http.RoundTripper, authToken string) TelegramSender { func NewTelegramSender(transport http.RoundTripper, authToken string) TelegramSender {
//TODO: Rate limit
return &telegramSender{ return &telegramSender{
client: &http.Client{Transport: transport}, client: &http.Client{Transport: transport},
authToken: authToken, authToken: authToken,
// 20 requests per min with some buffer
rateLimiterPerMin: rate.NewLimiter(rate.Every(65*time.Second), 20),
// 1 msg per sec with some buffer
rateLimiterPerSec: rate.NewLimiter(rate.Every(1050*time.Millisecond), 1),
} }
} }