Skip to content

Commit

Permalink
rate limiter: add (#487)
Browse files Browse the repository at this point in the history
Co-authored-by: demget <30910794+demget@users.noreply.github.com>
  • Loading branch information
setval and demget authored Nov 27, 2023
1 parent 71ac299 commit eb881a6
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 1 deletion.
16 changes: 16 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ go get -u gopkg.in/telebot.v3
- [Editable](#editable)
- [Keyboards](#keyboards)
- [Inline mode](#inline-mode)
- [Rate limiter](#rate-limiter)
* [Contributing](#contributing)
* [Donate](#donate)
* [License](#license)
Expand Down Expand Up @@ -460,6 +461,21 @@ b.Handle(tele.OnQuery, func(c tele.Context) error {
})
```

## Rate limiter
In order not to catch the anti-flood, an implementation has been made that allows you to limit the number of outgoing requests by n times per second.
```go
pref := tele.Settings{
Token: "123456:token",
// PerSeconds limits the number of requests per second executed
// by the client for Raw function. Default per in seconds is -1.
// To enable the queue, set the value greater than zero.
PerSeconds: 30,
// PerBufferSize sets the size of the queue that is waiting for the signal to be sent.
// By default, the buffer is infinite and has a value equal to zero.
PerBufferSize: 0,
}
```

There's not much to talk about really. It also supports some form of authentication
through deep-linking. For that, use fields `SwitchPMText` and `SwitchPMParameter`
of `QueryResponse`.
Expand Down
12 changes: 12 additions & 0 deletions api.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,25 @@ import (
"os"
"strconv"
"strings"
"sync"
"time"
)

// Raw lets you call any method of Bot API manually.
// It also handles API errors, so you only need to unwrap
// result field from json data.
func (b *Bot) Raw(method string, payload interface{}) ([]byte, error) {
if b.ratelimit {
c := sync.NewCond(&sync.Mutex{})
c.L.Lock()
b.raws <- c
c.Wait()
c.L.Unlock()
}
return b.raw(method, payload)
}

func (b *Bot) raw(method string, payload interface{}) ([]byte, error) {
url := b.URL + "/bot" + b.Token + "/" + method

var buf bytes.Buffer
Expand Down
40 changes: 39 additions & 1 deletion bot.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ import (
"os"
"regexp"
"strconv"
"strings"
"time"
"strings"
"sync"

"go.uber.org/ratelimit"
)

// NewBot does try to build a Bot with token `token`, which
Expand All @@ -19,6 +22,9 @@ func NewBot(pref Settings) (*Bot, error) {
if pref.Updates == 0 {
pref.Updates = 100
}
if pref.PerSeconds == 0 {
pref.PerSeconds = -1
}

client := pref.Client
if client == nil {
Expand Down Expand Up @@ -49,8 +55,29 @@ func NewBot(pref Settings) (*Bot, error) {
verbose: pref.Verbose,
parseMode: pref.ParseMode,
client: client,

ratelimit: false,
}

wait := make(chan struct{})
if pref.PerSeconds != -1 {
bot.ratelimit = true
go func(b *Bot) {
rl := ratelimit.New(pref.PerSeconds)
bot.raws = make(chan *sync.Cond, pref.PerBufferSize)
wait <- struct{}{}
for raw := range bot.raws {
rl.Take()
raw.Broadcast()
}
}(bot)
} else {
go func() {
wait <- struct{}{}
}()
}
<-wait

if pref.Offline {
bot.Me = &User{}
} else {
Expand Down Expand Up @@ -82,6 +109,8 @@ type Bot struct {
stop chan chan struct{}
client *http.Client
stopClient chan struct{}
raws chan *sync.Cond
ratelimit bool
}

// Settings represents a utility struct for passing certain
Expand Down Expand Up @@ -119,6 +148,15 @@ type Settings struct {

// Offline allows to create a bot without network for testing purposes.
Offline bool

// PerSeconds limits the number of requests per second executed
// by the client for Raw function. Default per in seconds is -1.
// To enable the queue, set the value greater than zero.
PerSeconds int

// PerBufferSize sets the size of the queue that is waiting for the signal to be sent.
// By default, the buffer is infinite and has a value equal to zero.
PerBufferSize int
}

var defaultOnError = func(err error, c Context) {
Expand Down
5 changes: 5 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ go 1.13

require (
github.com/goccy/go-yaml v1.9.5
github.com/mattn/go-colorable v0.1.12 // indirect
github.com/spf13/cast v1.3.1
go.uber.org/ratelimit v0.2.0 // indirect
golang.org/x/sys v0.0.0-20220114195835-da31bd327af9 // indirect
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
github.com/spf13/viper v1.13.0
github.com/stretchr/testify v1.8.0
)

0 comments on commit eb881a6

Please sign in to comment.