perf: optimize user storage retrieval and remove unused rate limiting middleware

This commit is contained in:
krau
2025-04-15 21:04:48 +08:00
parent c8c348a182
commit 609289c16a
6 changed files with 27 additions and 36 deletions
-6
View File
@@ -1,24 +1,18 @@
package bot package bot
import ( import (
"time"
"github.com/celestix/gotgproto/dispatcher" "github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext" "github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice" "github.com/duke-git/lancet/v2/slice"
"github.com/gotd/contrib/middleware/floodwait" "github.com/gotd/contrib/middleware/floodwait"
"github.com/gotd/contrib/middleware/ratelimit"
"github.com/gotd/td/telegram" "github.com/gotd/td/telegram"
"github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/config"
"golang.org/x/time/rate"
) )
func FloodWaitMiddleware() []telegram.Middleware { func FloodWaitMiddleware() []telegram.Middleware {
waiter := floodwait.NewSimpleWaiter().WithMaxRetries(5) waiter := floodwait.NewSimpleWaiter().WithMaxRetries(5)
ratelimiter := ratelimit.New(rate.Every(time.Millisecond*100), 5)
return []telegram.Middleware{ return []telegram.Middleware{
waiter, waiter,
ratelimiter,
} }
} }
+1 -1
View File
@@ -62,7 +62,7 @@ func getSelectStorageMarkup(userChatID int64, fileChatID, fileMessageID int) (*t
} }
buttons = append(buttons, &tg.KeyboardButtonCallback{ buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(), Text: storage.Name(),
Data: []byte(fmt.Sprintf("add %d", cbDataId)), Data: fmt.Appendf(nil, "add %d", cbDataId),
}) })
} }
markup := &tg.ReplyInlineMarkup{} markup := &tg.ReplyInlineMarkup{}
+12 -26
View File
@@ -10,40 +10,26 @@ type userConfig struct {
Blacklist bool `toml:"blacklist" mapstructure:"blacklist" json:"blacklist"` // 黑名单模式, storage names 中的存储将不会被使用, 默认为白名单模式 Blacklist bool `toml:"blacklist" mapstructure:"blacklist" json:"blacklist"` // 黑名单模式, storage names 中的存储将不会被使用, 默认为白名单模式
} }
var userIDs []int64
var storages []string
var userStorages = make(map[int64][]string)
func (c *Config) GetStorageNamesByUserID(userID int64) []string { func (c *Config) GetStorageNamesByUserID(userID int64) []string {
for _, user := range c.Users { us, ok := userStorages[userID]
if user.ID == userID { if ok {
if user.Blacklist { return us
allStorages := make([]string, 0, len(c.Storages))
for _, storage := range c.Storages {
allStorages = append(allStorages, storage.GetName())
}
return slice.Compact(slice.Difference(allStorages, user.Storages))
} else {
return user.Storages
}
}
} }
return nil return nil
} }
func (c *Config) GetUsersID() []int64 { func (c *Config) GetUsersID() []int64 {
var ids []int64 return userIDs
for _, user := range c.Users {
ids = append(ids, user.ID)
}
return ids
} }
func (c *Config) HasStorage(userID int64, storageName string) bool { func (c *Config) HasStorage(userID int64, storageName string) bool {
for _, user := range c.Users { us, ok := userStorages[userID]
if user.ID == userID { if !ok {
if user.Blacklist { return false
return !slice.Contain(user.Storages, storageName)
} else {
return slice.Contain(user.Storages, storageName)
}
}
} }
return false return slice.Contain(us, storageName)
} }
+14
View File
@@ -5,6 +5,7 @@ import (
"os" "os"
"strings" "strings"
"github.com/duke-git/lancet/v2/slice"
"github.com/krau/SaveAny-Bot/config/storage" "github.com/krau/SaveAny-Bot/config/storage"
"github.com/spf13/viper" "github.com/spf13/viper"
) )
@@ -137,6 +138,19 @@ func Init() error {
return fmt.Errorf("workers 和 retry 必须大于 0, 当前值: workers=%d, retry=%d", Cfg.Workers, Cfg.Retry) return fmt.Errorf("workers 和 retry 必须大于 0, 当前值: workers=%d, retry=%d", Cfg.Workers, Cfg.Retry)
} }
for _, storage := range Cfg.Storages {
storages = append(storages, storage.GetName())
}
for _, user := range Cfg.Users {
userIDs = append(userIDs, user.ID)
if user.Blacklist {
userStorages[user.ID] = slice.Compact(slice.Difference(storages, user.Storages))
} else {
userStorages[user.ID] = user.Storages
}
}
return nil return nil
} }
-1
View File
@@ -17,7 +17,6 @@ require (
github.com/spf13/cobra v1.8.1 github.com/spf13/cobra v1.8.1
github.com/spf13/viper v1.19.0 github.com/spf13/viper v1.19.0
golang.org/x/net v0.37.0 golang.org/x/net v0.37.0
golang.org/x/time v0.10.0
) )
require ( require (
-2
View File
@@ -263,8 +263,6 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY= golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY=
golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4= golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4=
golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=