Compare commits

...

5 Commits

Author SHA1 Message Date
fa7a296a66 v1.0.0 beta 7; ratelimt war 2026-03-02 16:49:00 +03:00
7101aba548 v1.0.0 beta 6 2026-03-02 00:08:26 +03:00
2de46a27c8 v1.0.0 beta 5 2026-03-01 23:40:27 +03:00
ae7426c36a 1.0.0 beta 4 2026-03-01 23:08:22 +03:00
61562e8a3b 1.0.0 beta 3 2026-03-01 23:01:06 +03:00
15 changed files with 396 additions and 72 deletions

12
bot.go
View File

@@ -11,9 +11,9 @@ import (
"git.nix13.pw/scuroneko/extypes"
"git.nix13.pw/scuroneko/laniakea/tgapi"
"git.nix13.pw/scuroneko/laniakea/utils"
"git.nix13.pw/scuroneko/slog"
"github.com/alitto/pond/v2"
"golang.org/x/time/rate"
)
type BotOpts struct {
@@ -88,6 +88,7 @@ type Bot[T DbContext] struct {
uploader *tgapi.Uploader
dbContext *T
l10n *L10n
draftProvider *DraftProvider
updateOffsetMu sync.Mutex
updateOffset int
@@ -98,9 +99,9 @@ type Bot[T DbContext] struct {
func NewBot[T any](opts *BotOpts) *Bot[T] {
updateQueue := make(chan *tgapi.Update, 512)
var limiter *rate.Limiter
var limiter *utils.RateLimiter
if opts.RateLimit > 0 {
limiter = rate.NewLimiter(rate.Limit(opts.RateLimit), opts.RateLimit)
limiter = utils.NewRateLimiter()
}
apiOpts := tgapi.NewAPIOpts(opts.Token).SetAPIUrl(opts.APIUrl).UseTestServer(opts.UseTestServer).SetLimiter(limiter)
@@ -122,6 +123,7 @@ func NewBot[T any](opts *BotOpts) *Bot[T] {
runners: make([]Runner[T], 0),
extraLoggers: make([]*slog.Logger, 0),
l10n: &L10n{},
draftProvider: NewRandomDraftProvider(api),
}
bot.extraLoggers = bot.extraLoggers.Push(api.GetLogger()).Push(uploader.GetLogger())
@@ -202,6 +204,10 @@ func (bot *Bot[T]) GetUpdateTypes() []tgapi.UpdateType { return bot.updateTypes
func (bot *Bot[T]) GetLogger() *slog.Logger { return bot.logger }
func (bot *Bot[T]) GetDBContext() *T { return bot.dbContext }
func (bot *Bot[T]) L10n(lang, key string) string { return bot.l10n.Translate(lang, key) }
func (bot *Bot[T]) SetDraftProvider(p *DraftProvider) *Bot[T] {
bot.draftProvider = p
return bot
}
type DbLogger[T DbContext] func(db *T) slog.LoggerWriter

106
drafts.go Normal file
View File

@@ -0,0 +1,106 @@
package laniakea
import (
"math"
"math/rand/v2"
"sync/atomic"
"git.nix13.pw/scuroneko/laniakea/tgapi"
)
type draftIdGenerator interface {
Next() uint64
}
type RandomDraftIdGenerator struct {
draftIdGenerator
}
func (g *RandomDraftIdGenerator) Next() uint64 {
return rand.Uint64N(math.MaxUint64)
}
type LinearDraftIdGenerator struct {
draftIdGenerator
lastId uint64
}
func (g *LinearDraftIdGenerator) Next() uint64 {
return atomic.AddUint64(&g.lastId, 1)
}
type DraftProvider struct {
api *tgapi.API
chatID int64
messageThreadID int
parseMode tgapi.ParseMode
entities []tgapi.MessageEntity
drafts map[uint64]*Draft
generator draftIdGenerator
}
type Draft struct {
api *tgapi.API
chatID int64
messageThreadID int
parseMode tgapi.ParseMode
entities []tgapi.MessageEntity
ID uint64
Message string
}
func NewRandomDraftProvider(api *tgapi.API) *DraftProvider {
return &DraftProvider{
api: api, generator: &RandomDraftIdGenerator{},
drafts: make(map[uint64]*Draft),
}
}
func NewLinearDraftProvider(api *tgapi.API, startValue uint64) *DraftProvider {
return &DraftProvider{
api: api,
generator: &LinearDraftIdGenerator{lastId: startValue},
drafts: make(map[uint64]*Draft),
}
}
func (d *DraftProvider) NewDraft() *Draft {
id := d.generator.Next()
draft := &Draft{d.api, d.chatID, d.messageThreadID, d.parseMode, d.entities, id, ""}
d.drafts[id] = draft
return draft
}
func (d *Draft) Push(newText string) error {
d.Message += newText
params := tgapi.SendMessageDraftP{
ChatID: d.chatID,
DraftID: d.ID,
Text: d.Message,
ParseMode: d.parseMode,
Entities: d.entities,
}
if d.messageThreadID > 0 {
params.MessageThreadID = d.messageThreadID
}
_, err := d.api.SendMessageDraft(params)
return err
}
func (d *Draft) Flush() error {
if d.Message == "" {
return nil
}
params := tgapi.SendMessageP{
ChatID: d.chatID,
ParseMode: d.parseMode,
Entities: d.entities,
Text: d.Message,
}
if d.messageThreadID > 0 {
params.MessageThreadID = d.messageThreadID
}
_, err := d.api.SendMessage(params)
return err
}

View File

@@ -9,7 +9,13 @@ import (
)
func (bot *Bot[T]) handle(u *tgapi.Update) {
ctx := &MsgContext{Update: *u, Api: bot.api, botLogger: bot.logger, errorTemplate: bot.errorTemplate, l10n: bot.l10n}
ctx := &MsgContext{
Update: *u, Api: bot.api,
botLogger: bot.logger,
errorTemplate: bot.errorTemplate,
l10n: bot.l10n,
draftProvider: bot.draftProvider,
}
for _, middleware := range bot.middlewares {
middleware.Execute(ctx, bot.dbContext)
}

View File

@@ -1,6 +1,7 @@
package laniakea
import (
"context"
"fmt"
"git.nix13.pw/scuroneko/laniakea/tgapi"
@@ -24,6 +25,7 @@ type MsgContext struct {
errorTemplate string
botLogger *slog.Logger
l10n *L10n
draftProvider *DraftProvider
}
type AnswerMessage struct {
@@ -77,6 +79,7 @@ func (ctx *MsgContext) editPhotoText(messageId int, text string, kb *InlineKeybo
if kb != nil {
params.ReplyMarkup = kb.Get()
}
msg, _, err := ctx.Api.EditMessageCaption(params)
if err != nil {
ctx.botLogger.Errorln(err)
@@ -105,7 +108,18 @@ func (ctx *MsgContext) answer(text string, keyboard *InlineKeyboard) *AnswerMess
if keyboard != nil {
params.ReplyMarkup = keyboard.Get()
}
if ctx.Msg.MessageThreadID > 0 {
params.MessageThreadID = ctx.Msg.MessageThreadID
}
if ctx.Msg.DirectMessageTopic != nil {
params.DirectMessagesTopicID = ctx.Msg.DirectMessageTopic.TopicID
}
cont := context.Background()
if err := ctx.Api.Limiter.Wait(cont, ctx.Msg.Chat.ID); err != nil {
ctx.botLogger.Errorln(err)
return nil
}
msg, err := ctx.Api.SendMessage(params)
if err != nil {
ctx.botLogger.Errorln(err)
@@ -135,6 +149,10 @@ func (ctx *MsgContext) answerPhoto(photoId, text string, kb *InlineKeyboard) *An
if kb != nil {
params.ReplyMarkup = kb.Get()
}
if ctx.Msg.MessageThreadID > 0 {
params.MessageThreadID = ctx.Msg.MessageThreadID
}
msg, err := ctx.Api.SendPhoto(params)
if err != nil {
ctx.botLogger.Errorln(err)
@@ -162,12 +180,8 @@ func (ctx *MsgContext) delete(messageId int) {
ctx.botLogger.Errorln(err)
}
}
func (m *AnswerMessage) Delete() {
m.ctx.delete(m.MessageID)
}
func (ctx *MsgContext) CallbackDelete() {
ctx.delete(ctx.CallbackMsgId)
}
func (m *AnswerMessage) Delete() { m.ctx.delete(m.MessageID) }
func (ctx *MsgContext) CallbackDelete() { ctx.delete(ctx.CallbackMsgId) }
func (ctx *MsgContext) answerCallbackQuery(url, text string, showAlert bool) {
if len(ctx.CallbackQueryId) == 0 {
@@ -181,23 +195,19 @@ func (ctx *MsgContext) answerCallbackQuery(url, text string, showAlert bool) {
ctx.botLogger.Errorln(err)
}
}
func (ctx *MsgContext) AnswerCbQuery() {
ctx.answerCallbackQuery("", "", false)
}
func (ctx *MsgContext) AnswerCbQueryText(text string) {
ctx.answerCallbackQuery("", text, false)
}
func (ctx *MsgContext) AnswerCbQueryAlert(text string) {
ctx.answerCallbackQuery("", text, true)
}
func (ctx *MsgContext) AnswerCbQueryUrl(u string) {
ctx.answerCallbackQuery(u, "", false)
}
func (ctx *MsgContext) AnswerCbQuery() { ctx.answerCallbackQuery("", "", false) }
func (ctx *MsgContext) AnswerCbQueryText(text string) { ctx.answerCallbackQuery("", text, false) }
func (ctx *MsgContext) AnswerCbQueryAlert(text string) { ctx.answerCallbackQuery("", text, true) }
func (ctx *MsgContext) AnswerCbQueryUrl(u string) { ctx.answerCallbackQuery(u, "", false) }
func (ctx *MsgContext) SendAction(action tgapi.ChatActionType) {
_, err := ctx.Api.SendChatAction(tgapi.SendChatActionP{
params := tgapi.SendChatActionP{
ChatID: ctx.Msg.Chat.ID, Action: action,
})
}
if ctx.Msg.MessageThreadID > 0 {
params.MessageThreadID = ctx.Msg.MessageThreadID
}
_, err := ctx.Api.SendChatAction(params)
if err != nil {
ctx.botLogger.Errorln(err)
}
@@ -213,10 +223,20 @@ func (ctx *MsgContext) error(err error) {
}
ctx.botLogger.Errorln(err)
}
func (ctx *MsgContext) Error(err error) {
ctx.error(err)
func (ctx *MsgContext) Error(err error) { ctx.error(err) }
func (ctx *MsgContext) NewDraft() *Draft {
c := context.Background()
if err := ctx.Api.Limiter.Wait(c, ctx.Msg.Chat.ID); err != nil {
ctx.botLogger.Errorln(err)
return nil
}
draft := ctx.draftProvider.NewDraft()
draft.chatID = ctx.Msg.Chat.ID
draft.messageThreadID = ctx.Msg.MessageThreadID
return draft
}
func (ctx *MsgContext) Translate(key string) string {
if ctx.From == nil {
return key

View File

@@ -12,7 +12,6 @@ import (
"git.nix13.pw/scuroneko/laniakea/utils"
"git.nix13.pw/scuroneko/slog"
"golang.org/x/time/rate"
)
type APIOpts struct {
@@ -21,7 +20,7 @@ type APIOpts struct {
useTestServer bool
apiUrl string
limiter *rate.Limiter
limiter *utils.RateLimiter
dropOverflowLimit bool
}
@@ -46,7 +45,7 @@ func (opts *APIOpts) SetAPIUrl(apiUrl string) *APIOpts {
}
return opts
}
func (opts *APIOpts) SetLimiter(limiter *rate.Limiter) *APIOpts {
func (opts *APIOpts) SetLimiter(limiter *utils.RateLimiter) *APIOpts {
opts.limiter = limiter
return opts
}
@@ -63,7 +62,7 @@ type API struct {
apiUrl string
pool *WorkerPool
limiter *rate.Limiter
Limiter *utils.RateLimiter
dropOverflowLimit bool
}
@@ -88,11 +87,17 @@ func (api *API) CloseApi() error {
}
func (api *API) GetLogger() *slog.Logger { return api.logger }
type ResponseParameters struct {
MigrateToChatID *int64 `json:"migrate_to_chat_id,omitempty"`
RetryAfter *int `json:"retry_after,omitempty"`
}
type ApiResponse[R any] struct {
Ok bool `json:"ok"`
Description string `json:"description,omitempty"`
Result R `json:"result,omitempty"`
ErrorCode int `json:"error_code,omitempty"`
Parameters *ResponseParameters `json:"parameters,omitempty"`
}
type TelegramRequest[R, P any] struct {
method string
@@ -104,13 +109,13 @@ func NewRequest[R, P any](method string, params P) TelegramRequest[R, P] {
}
func (r TelegramRequest[R, P]) doRequest(ctx context.Context, api *API) (R, error) {
var zero R
if api.limiter != nil {
if api.Limiter != nil {
if api.dropOverflowLimit {
if !api.limiter.Allow() {
if !api.Limiter.GlobalAllow() {
return zero, errors.New("rate limited")
}
} else {
if err := api.limiter.Wait(ctx); err != nil {
if err := api.Limiter.GlobalWait(ctx); err != nil {
return zero, err
}
}
@@ -149,10 +154,23 @@ func (r TelegramRequest[R, P]) doRequest(ctx context.Context, api *API) (R, erro
return zero, err
}
api.logger.Debugln("RES", r.method, string(data))
if res.StatusCode != http.StatusOK {
if res.StatusCode != http.StatusOK && res.StatusCode != http.StatusTooManyRequests {
return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(data))
}
return parseBody[R](data)
responseData, err := parseBody[R](data)
if errors.Is(err, ErrRateLimit) {
if responseData.Parameters != nil {
after := 0
if responseData.Parameters.RetryAfter != nil {
after = *responseData.Parameters.RetryAfter
}
api.Limiter.SetGlobalLock(after)
return r.doRequest(ctx, api)
}
return zero, ErrRateLimit
}
return responseData.Result, err
}
func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, error) {
var zero R
@@ -184,15 +202,17 @@ func readBody(body io.ReadCloser) ([]byte, error) {
reader := io.LimitReader(body, 10<<20)
return io.ReadAll(reader)
}
func parseBody[R any](data []byte) (R, error) {
var zero R
func parseBody[R any](data []byte) (ApiResponse[R], error) {
var resp ApiResponse[R]
err := json.Unmarshal(data, &resp)
if err != nil {
return zero, err
return resp, err
}
if !resp.Ok {
return zero, fmt.Errorf("[%d] %s", resp.ErrorCode, resp.Description)
if resp.ErrorCode == 429 {
return resp, ErrRateLimit
}
return resp.Result, nil
return resp, fmt.Errorf("[%d] %s", resp.ErrorCode, resp.Description)
}
return resp, nil
}

View File

@@ -2,7 +2,7 @@ package tgapi
type SendPhotoP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
DirectMessagesTopicID int `json:"direct_messages_topic_id,omitempty"`

View File

@@ -56,6 +56,7 @@ type PromoteChatMember struct {
CanPinMessages bool `json:"can_pin_messages,omitempty"`
CanManageTopics bool `json:"can_manage_topics,omitempty"`
CanManageDirectMessages bool `json:"can_manage_direct_messages,omitempty"`
CanManageTags bool `json:"can_manage_tags,omitempty"`
}
func (api *API) PromoteChatMember(params PromoteChatMember) (bool, error) {
@@ -74,6 +75,17 @@ func (api *API) SetChatAdministratorCustomTitle(params SetChatAdministratorCusto
return req.Do(api)
}
type SetChatMemberTagP struct {
ChatID int `json:"chat_id"`
UserID int `json:"user_id"`
Tag string `json:"tag,omitempty"`
}
func (api *API) SetChatMemberTag(params SetChatMemberTagP) (bool, error) {
req := NewRequest[bool]("setChatMemberTag", params)
return req.Do(api)
}
type BanChatSenderChatP struct {
ChatID int `json:"chat_id"`
SenderChatID int `json:"sender_chat_id"`

View File

@@ -1,7 +1,7 @@
package tgapi
type Chat struct {
ID int `json:"id"`
ID int64 `json:"id"`
Type string `json:"type"`
Title *string `json:"title,omitempty"`
Username *string `json:"username,omitempty"`
@@ -99,6 +99,7 @@ type ChatPermissions struct {
CanSendPolls bool `json:"can_send_polls"`
CanSendOtherMessages bool `json:"can_send_other_messages"`
CanAddWebPagePreview bool `json:"can_add_web_page_preview"`
CatEditTag bool `json:"cat_edit_tag"`
CanChangeInfo bool `json:"can_change_info"`
CanInviteUsers bool `json:"can_invite_users"`
CanPinMessages bool `json:"can_pin_messages"`
@@ -137,6 +138,7 @@ const (
type ChatMember struct {
Status ChatMemberStatusType `json:"status"`
User User `json:"user"`
Tag string `json:"tag,omitempty"`
// Owner
IsAnonymous *bool `json:"is_anonymous"`
@@ -160,6 +162,7 @@ type ChatMember struct {
CanPinMessages *bool `json:"can_pin_messages,omitempty"`
CanManageTopics *bool `json:"can_manage_topics,omitempty"`
CanManageDirectMessages *bool `json:"can_manage_direct_messages,omitempty"`
CanManageTags *bool `json:"can_manage_tags,omitempty"`
// Member
UntilDate *int `json:"until_date,omitempty"`
@@ -175,6 +178,7 @@ type ChatMember struct {
CanSendPolls *bool `json:"can_send_polls,omitempty"`
CanSendOtherMessages *bool `json:"can_send_other_messages,omitempty"`
CanAddWebPagePreview *bool `json:"can_add_web_page_preview,omitempty"`
CanEditTag *bool `json:"can_edit_tag,omitempty"`
}
type ChatBoostSource struct {
@@ -215,6 +219,7 @@ type ChatAdministratorRights struct {
CanPinMessages *bool `json:"can_pin_messages,omitempty"`
CanManageTopics *bool `json:"can_manage_topics,omitempty"`
CanManageDirectMessages *bool `json:"can_manage_direct_messages,omitempty"`
CanManageTags *bool `json:"can_manage_tags,omitempty"`
}
type ChatBoostUpdated struct {

5
tgapi/errors.go Normal file
View File

@@ -0,0 +1,5 @@
package tgapi
import "errors"
var ErrRateLimit = errors.New("rate limit exceeded")

View File

@@ -2,9 +2,9 @@ package tgapi
type SendMessageP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
DirectMessagesTopicID int `json:"direct_messages_topic_id,omitempty"`
DirectMessagesTopicID int64 `json:"direct_messages_topic_id,omitempty"`
Text string `json:"text"`
ParseMode ParseMode `json:"parse_mode,omitempty"`
@@ -266,9 +266,9 @@ func (api *API) SendDice(params SendDiceP) (Message, error) {
}
type SendMessageDraftP struct {
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
DraftID int `json:"draft_id"`
DraftID uint64 `json:"draft_id"`
Text string `json:"text"`
ParseMode ParseMode `json:"parse_mode,omitempty"`
Entities []MessageEntity `json:"entities,omitempty"`
@@ -281,7 +281,7 @@ func (api *API) SendMessageDraft(params SendMessageDraftP) (bool, error) {
type SendChatActionP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
Action ChatActionType `json:"action"`
}
@@ -307,7 +307,7 @@ func (api *API) SetMessageReaction(params SetMessageReactionP) (bool, error) {
type EditMessageTextP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id,omitempty"`
ChatID int64 `json:"chat_id,omitempty"`
MessageID int `json:"message_id,omitempty"`
InlineMessageID string `json:"inline_message_id,omitempty"`
Text string `json:"text"`
@@ -331,7 +331,7 @@ func (api *API) EditMessageText(params EditMessageTextP) (Message, bool, error)
type EditMessageCaptionP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id,omitempty"`
ChatID int64 `json:"chat_id,omitempty"`
MessageID int `json:"message_id,omitempty"`
InlineMessageID string `json:"inline_message_id,omitempty"`
Caption string `json:"caption"`
@@ -495,7 +495,7 @@ func (api *API) DeclineSuggestedPost(params DeclineSuggestedPostP) (bool, error)
}
type DeleteMessageP struct {
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageID int `json:"message_id"`
}

View File

@@ -6,15 +6,22 @@ type MessageReplyMarkup struct {
InlineKeyboard [][]InlineKeyboardButton `json:"inline_keyboard"`
}
type DirectMessageTopic struct {
TopicID int64 `json:"topic_id"`
User *User `json:"user,omitempty"`
}
type Message struct {
MessageID int `json:"message_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
DirectMessageTopic *DirectMessageTopic `json:"direct_message_topic,omitempty"`
BusinessConnectionId string `json:"business_connection_id,omitempty"`
From *User `json:"from,omitempty"`
SenderChat *Chat `json:"sender_chat,omitempty"`
SenderBoostCount int `json:"sender_boost_count,omitempty"`
SenderBusinessBot *User `json:"sender_business_bot,omitempty"`
SenderTag string `json:"sender_tag,omitempty"`
Chat *Chat `json:"chat,omitempty"`
IsTopicMessage bool `json:"is_topic_message,omitempty"`
@@ -74,6 +81,7 @@ const (
MessageEntityTextLink MessageEntityType = "text_link"
MessageEntityTextMention MessageEntityType = "text_mention"
MessageEntityCustomEmoji MessageEntityType = "custom_emoji"
MessageEntityDateTime MessageEntityType = "date_time"
)
type MessageEntity struct {
@@ -85,6 +93,9 @@ type MessageEntity struct {
User *User `json:"user,omitempty"`
Language string `json:"language,omitempty"`
CustomEmojiID string `json:"custom_emoji_id,omitempty"`
UnixTime int `json:"unix_time,omitempty"`
DateTimeFormat string `json:"date_time_format,omitempty"`
}
type ReplyParameters struct {

View File

@@ -66,13 +66,13 @@ func NewUploaderRequest[R, P any](method string, params P, files ...UploaderFile
}
func (r UploaderRequest[R, P]) doRequest(ctx context.Context, up *Uploader) (R, error) {
var zero R
if up.api.limiter != nil {
if up.api.Limiter != nil {
if up.api.dropOverflowLimit {
if !up.api.limiter.Allow() {
if !up.api.Limiter.GlobalAllow() {
return zero, errors.New("rate limited")
}
} else {
if err := up.api.limiter.Wait(ctx); err != nil {
if err := up.api.Limiter.GlobalWait(ctx); err != nil {
return zero, err
}
}
@@ -109,7 +109,11 @@ func (r UploaderRequest[R, P]) doRequest(ctx context.Context, up *Uploader) (R,
return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(body))
}
return parseBody[R](body)
respBody, err := parseBody[R](body)
if err != nil {
return zero, err
}
return respBody.Result, nil
}
func (r UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader) (R, error) {
var zero R

View File

@@ -2,7 +2,7 @@ package tgapi
type UploadPhotoP struct {
BusinessConnectionID string `json:"business_connection_id,omitempty"`
ChatID int `json:"chat_id"`
ChatID int64 `json:"chat_id"`
MessageThreadID int `json:"message_thread_id,omitempty"`
DirectMessagesTopicID int `json:"direct_messages_topic_id,omitempty"`

129
utils/limiter.go Normal file
View File

@@ -0,0 +1,129 @@
package utils
import (
"context"
"sync"
"time"
"golang.org/x/time/rate"
)
type RateLimiter struct {
globalLockUntil time.Time
globalLimiter *rate.Limiter
globalMu sync.RWMutex
chatLocks map[int64]time.Time
chatLimiters map[int64]*rate.Limiter
chatMu sync.Mutex
}
func NewRateLimiter() *RateLimiter {
return &RateLimiter{
// 30 запросов в секунду (burst=30)
globalLimiter: rate.NewLimiter(rate.Limit(30), 30),
chatLimiters: make(map[int64]*rate.Limiter),
}
}
func (rl *RateLimiter) SetGlobalLock(retryAfter int) {
if retryAfter <= 0 {
return
}
rl.globalMu.Lock()
defer rl.globalMu.Unlock()
rl.globalLockUntil = time.Now().Add(time.Duration(retryAfter) * time.Second)
}
func (rl *RateLimiter) SetChatLock(chatID int64, retryAfter int) {
rl.chatMu.Lock()
defer rl.chatMu.Unlock()
rl.chatLocks[chatID] = time.Now().Add(time.Duration(retryAfter) * time.Second)
}
func (rl *RateLimiter) GlobalWait(ctx context.Context) error {
rl.globalMu.RLock()
until := rl.globalLockUntil
rl.globalMu.RUnlock()
if !until.IsZero() {
if time.Now().Before(until) {
// Ждём до окончания блокировки или отмены контекста
select {
case <-time.After(time.Until(until)):
// блокировка снята
case <-ctx.Done():
return ctx.Err()
}
}
}
// Теперь ждём разрешения rate limiter'а
return rl.globalLimiter.Wait(ctx)
}
func (rl *RateLimiter) Wait(ctx context.Context, chatID int64) error {
rl.chatMu.Lock()
until, ok := rl.chatLocks[chatID]
rl.chatMu.Unlock()
if ok && !until.IsZero() {
if time.Now().Before(until) {
select {
case <-time.After(time.Until(until)):
// блокировка снята
case <-ctx.Done():
return ctx.Err()
}
}
}
if err := rl.GlobalWait(ctx); err != nil {
return err
}
rl.chatMu.Lock()
chatLimiter, ok := rl.chatLimiters[chatID]
if !ok {
chatLimiter = rate.NewLimiter(rate.Limit(1), 1)
rl.chatLimiters[chatID] = chatLimiter
}
rl.chatMu.Unlock()
return chatLimiter.Wait(ctx)
}
func (rl *RateLimiter) GlobalAllow() bool {
rl.globalMu.RLock()
until := rl.globalLockUntil
rl.globalMu.RUnlock()
if !until.IsZero() {
if time.Now().Before(until) {
// Ждём до окончания блокировки или отмены контекста
select {
case <-time.After(time.Until(until)):
rl.globalLimiter.Allow()
}
}
}
return rl.globalLimiter.Allow()
}
func (rl *RateLimiter) Allow(chatID int64) bool {
rl.chatMu.Lock()
until, ok := rl.chatLocks[chatID]
rl.chatMu.Unlock()
if ok && !until.IsZero() {
if time.Now().Before(until) {
select {
case <-time.After(time.Until(until)):
}
}
}
if !rl.globalLimiter.Allow() {
return false
}
rl.chatMu.Lock()
chatLimiter, ok := rl.chatLimiters[chatID]
if !ok {
chatLimiter = rate.NewLimiter(rate.Limit(1), 1)
rl.chatLimiters[chatID] = chatLimiter
}
rl.chatMu.Unlock()
return chatLimiter.Allow()
}

View File

@@ -1,9 +1,9 @@
package utils
const (
VersionString = "1.0.0-beta.2"
VersionString = "1.0.0-beta.7"
VersionMajor = 1
VersionMinor = 0
VersionPatch = 0
Beta = 2
Beta = 7
)