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
go.uber.org/zap v1.21.0
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
)

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/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/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-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=

View File

@ -2,14 +2,17 @@ package telegram
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"time"
"go.balki.me/tss/log"
"golang.org/x/time/rate"
)
func main() {
@ -21,8 +24,10 @@ type TelegramSender interface {
}
type telegramSender struct {
client *http.Client
authToken string
client *http.Client
authToken string
rateLimiterPerMin *rate.Limiter
rateLimiterPerSec *rate.Limiter
}
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)
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))
if err != nil {
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 {
//TODO: Rate limit
return &telegramSender{
client: &http.Client{Transport: transport},
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),
}
}