Add rate limiter

telegram/telegram.go

@@ -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),
 	}
 }