mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-09-07 00:27:36 +08:00
perf: optimize user storage retrieval and remove unused rate limiting middleware
This commit is contained in:
@@ -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
@@ -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
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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 (
|
||||||
|
|||||||
@@ -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=
|
||||||
|
|||||||
Reference in New Issue
Block a user