Merge pull request #962 from zenixls2/fix/telegram_api_limit

This commit is contained in:
Yo-An Lin 2022-09-21 21:11:53 +08:00 committed by GitHub
commit da6161ddda
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 57 additions and 5 deletions

View File

@ -6,7 +6,9 @@ import (
"strings" "strings"
"time" "time"
"github.com/c9s/bbgo/pkg/util"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"golang.org/x/time/rate"
"gopkg.in/tucnak/telebot.v2" "gopkg.in/tucnak/telebot.v2"
) )
@ -15,6 +17,10 @@ func init() {
_ = Reply(&TelegramReply{}) _ = Reply(&TelegramReply{})
} }
var sendLimiter = rate.NewLimiter(10, 2)
const maxMessageSize int = 3000
type TelegramSessionMap map[int64]*TelegramSession type TelegramSessionMap map[int64]*TelegramSession
type TelegramSession struct { type TelegramSession struct {
@ -62,7 +68,14 @@ type TelegramReply struct {
} }
func (r *TelegramReply) Send(message string) { func (r *TelegramReply) Send(message string) {
checkSendErr(r.bot.Send(r.session.Chat, message)) ctx := context.Background()
splits := util.StringSplitByLength(message, maxMessageSize)
for _, split := range splits {
if err := sendLimiter.Wait(ctx); err != nil {
log.WithError(err).Errorf("telegram send limit exceeded")
}
checkSendErr(r.bot.Send(r.session.Chat, split))
}
} }
func (r *TelegramReply) Message(message string) { func (r *TelegramReply) Message(message string) {
@ -132,7 +145,7 @@ func (tm *Telegram) SetTextMessageResponder(responder Responder) {
tm.textMessageResponder = responder tm.textMessageResponder = responder
} }
func (tm *Telegram) Start(context.Context) { func (tm *Telegram) Start(ctx context.Context) {
tm.Bot.Handle(telebot.OnCallback, func(c *telebot.Callback) { tm.Bot.Handle(telebot.OnCallback, func(c *telebot.Callback) {
log.Infof("[telegram] onCallback: %+v", c) log.Infof("[telegram] onCallback: %+v", c)
}) })
@ -158,7 +171,18 @@ func (tm *Telegram) Start(context.Context) {
if reply.set { if reply.set {
reply.build() reply.build()
if len(reply.message) > 0 || reply.menu != nil { if len(reply.message) > 0 || reply.menu != nil {
checkSendErr(tm.Bot.Send(m.Chat, reply.message, reply.menu)) splits := util.StringSplitByLength(reply.message, maxMessageSize)
for i, split := range splits {
if err := sendLimiter.Wait(ctx); err != nil {
log.WithError(err).Errorf("telegram send limit exceeded")
}
if i == len(splits)-1 {
// only set menu on the last message
checkSendErr(tm.Bot.Send(m.Chat, split, reply.menu))
} else {
checkSendErr(tm.Bot.Send(m.Chat, split))
}
}
} }
} }
}) })

View File

@ -1,6 +1,9 @@
package util package util
import "strings" import (
"strings"
"unicode/utf8"
)
func StringSliceContains(slice []string, needle string) bool { func StringSliceContains(slice []string, needle string) bool {
for _, s := range slice { for _, s := range slice {
@ -27,3 +30,17 @@ func MaskKey(key string) string {
maskKey += key[len(key)-h:] maskKey += key[len(key)-h:]
return maskKey return maskKey
} }
func StringSplitByLength(s string, length int) (result []string) {
var left, right int
for left, right = 0, length; right < len(s); left, right = right, right+length {
for !utf8.RuneStart(s[right]) {
right--
}
result = append(result, s[left:right])
}
if len(s)-left > 0 {
result = append(result, s[left:])
}
return result
}

View File

@ -1,6 +1,10 @@
package util package util
import "testing" import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestMaskKey(t *testing.T) { func TestMaskKey(t *testing.T) {
type args struct { type args struct {
@ -40,3 +44,10 @@ func TestMaskKey(t *testing.T) {
}) })
} }
} }
func TestStringSplitByLength(t *testing.T) {
result := StringSplitByLength("1234567890", 3)
assert.Equal(t, result, []string{"123", "456", "789", "0"})
result = StringSplitByLength("123許456", 4)
assert.Equal(t, result, []string{"123", "許4", "56"})
}