Add rate limiter
This commit is contained in:
		
							
								
								
									
										1
									
								
								go.mod
									
									
									
									
									
								
							
							
						
						
									
										1
									
								
								go.mod
									
									
									
									
									
								
							@@ -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
									
									
									
									
									
								
							
							
						
						
									
										2
									
								
								go.sum
									
									
									
									
									
								
							@@ -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=
 | 
			
		||||
 
 | 
			
		||||
@@ -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),
 | 
			
		||||
	}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user