REPOSITORY / ScuroNeko/Laniakea
Compare commits
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c0a26024f4
|
||
|
|
786da652e6
|
||
|
|
28ec2b7ca9
|
||
|
|
da122a3be4
|
||
|
|
1bf7499496
|
@@ -1,41 +1,64 @@
|
||||
package laniakea
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"sync"
|
||||
|
||||
"git.nix13.pw/scuroneko/extypes"
|
||||
"git.nix13.pw/scuroneko/laniakea/tgapi"
|
||||
"git.nix13.pw/scuroneko/slog"
|
||||
"github.com/alitto/pond/v2"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
type BotOpts struct {
|
||||
Token string
|
||||
Debug bool
|
||||
ErrorTemplate string
|
||||
Prefixes []string
|
||||
UpdateTypes []string
|
||||
Token string
|
||||
UpdateTypes []string
|
||||
|
||||
Debug bool
|
||||
ErrorTemplate string
|
||||
Prefixes []string
|
||||
|
||||
LoggerBasePath string
|
||||
UseRequestLogger bool
|
||||
WriteToFile bool
|
||||
UseTestServer bool
|
||||
APIUrl string
|
||||
|
||||
UseTestServer bool
|
||||
APIUrl string
|
||||
|
||||
RateLimit int
|
||||
DropRLOverflow bool
|
||||
}
|
||||
|
||||
func NewOpts() *BotOpts { return new(BotOpts) }
|
||||
func LoadOptsFromEnv() *BotOpts {
|
||||
rateLimit := 30
|
||||
if rl := os.Getenv("RATE_LIMIT"); rl != "" {
|
||||
rateLimit, _ = strconv.Atoi(rl)
|
||||
}
|
||||
|
||||
return &BotOpts{
|
||||
Token: os.Getenv("TG_TOKEN"),
|
||||
Debug: os.Getenv("DEBUG") == "true",
|
||||
ErrorTemplate: os.Getenv("ERROR_TEMPLATE"),
|
||||
Prefixes: LoadPrefixesFromEnv(),
|
||||
UpdateTypes: strings.Split(os.Getenv("UPDATE_TYPES"), ";"),
|
||||
Token: os.Getenv("TG_TOKEN"),
|
||||
UpdateTypes: strings.Split(os.Getenv("UPDATE_TYPES"), ";"),
|
||||
|
||||
Debug: os.Getenv("DEBUG") == "true",
|
||||
ErrorTemplate: os.Getenv("ERROR_TEMPLATE"),
|
||||
Prefixes: LoadPrefixesFromEnv(),
|
||||
|
||||
UseRequestLogger: os.Getenv("USE_REQ_LOG") == "true",
|
||||
WriteToFile: os.Getenv("WRITE_TO_FILE") == "true",
|
||||
UseTestServer: os.Getenv("USE_TEST_SERVER") == "true",
|
||||
APIUrl: os.Getenv("API_URL"),
|
||||
|
||||
UseTestServer: os.Getenv("USE_TEST_SERVER") == "true",
|
||||
APIUrl: os.Getenv("API_URL"),
|
||||
|
||||
RateLimit: rateLimit,
|
||||
DropRLOverflow: os.Getenv("DROP_RL_OVERFLOW") == "true",
|
||||
}
|
||||
}
|
||||
func LoadPrefixesFromEnv() []string {
|
||||
@@ -47,6 +70,7 @@ func LoadPrefixesFromEnv() []string {
|
||||
}
|
||||
|
||||
type DbContext interface{}
|
||||
type NoDB struct{ DbContext }
|
||||
type Bot[T DbContext] struct {
|
||||
token string
|
||||
debug bool
|
||||
@@ -54,6 +78,7 @@ type Bot[T DbContext] struct {
|
||||
|
||||
logger *slog.Logger
|
||||
RequestLogger *slog.Logger
|
||||
extraLoggers extypes.Slice[*slog.Logger]
|
||||
|
||||
plugins []Plugin[T]
|
||||
middlewares []Middleware[T]
|
||||
@@ -65,37 +90,41 @@ type Bot[T DbContext] struct {
|
||||
dbContext *T
|
||||
l10n *L10n
|
||||
|
||||
dbWriterRequested extypes.Slice[*slog.Logger]
|
||||
|
||||
updateOffset int
|
||||
updateTypes []tgapi.UpdateType
|
||||
updateQueue *extypes.Queue[*tgapi.Update]
|
||||
updateOffsetMu sync.Mutex
|
||||
updateOffset int
|
||||
updateTypes []tgapi.UpdateType
|
||||
updateQueue chan *tgapi.Update
|
||||
}
|
||||
|
||||
func NewBot[T any](opts *BotOpts) *Bot[T] {
|
||||
updateQueue := extypes.CreateQueue[*tgapi.Update](512)
|
||||
updateQueue := make(chan *tgapi.Update, 512)
|
||||
|
||||
apiOpts := tgapi.NewAPIOpts(opts.Token).SetAPIUrl(opts.APIUrl).UseTestServer(opts.UseTestServer)
|
||||
var limiter *rate.Limiter
|
||||
if opts.RateLimit > 0 {
|
||||
limiter = rate.NewLimiter(rate.Limit(opts.RateLimit), opts.RateLimit)
|
||||
}
|
||||
|
||||
apiOpts := tgapi.NewAPIOpts(opts.Token).SetAPIUrl(opts.APIUrl).UseTestServer(opts.UseTestServer).SetLimiter(limiter)
|
||||
api := tgapi.NewAPI(apiOpts)
|
||||
|
||||
uploader := tgapi.NewUploader(api)
|
||||
|
||||
bot := &Bot[T]{
|
||||
updateOffset: 0,
|
||||
errorTemplate: "%s",
|
||||
updateQueue: updateQueue,
|
||||
api: api,
|
||||
uploader: uploader,
|
||||
debug: opts.Debug,
|
||||
prefixes: opts.Prefixes,
|
||||
token: opts.Token,
|
||||
plugins: make([]Plugin[T], 0),
|
||||
updateTypes: make([]tgapi.UpdateType, 0),
|
||||
runners: make([]Runner[T], 0),
|
||||
dbWriterRequested: make([]*slog.Logger, 0),
|
||||
l10n: &L10n{},
|
||||
updateOffset: 0,
|
||||
errorTemplate: "%s",
|
||||
updateQueue: updateQueue,
|
||||
api: api,
|
||||
uploader: uploader,
|
||||
debug: opts.Debug,
|
||||
prefixes: opts.Prefixes,
|
||||
token: opts.Token,
|
||||
plugins: make([]Plugin[T], 0),
|
||||
updateTypes: make([]tgapi.UpdateType, 0),
|
||||
runners: make([]Runner[T], 0),
|
||||
extraLoggers: make([]*slog.Logger, 0),
|
||||
l10n: &L10n{},
|
||||
}
|
||||
bot.dbWriterRequested = bot.dbWriterRequested.Push(api.GetLogger()).Push(uploader.GetLogger())
|
||||
bot.extraLoggers = bot.extraLoggers.Push(api.GetLogger()).Push(uploader.GetLogger())
|
||||
|
||||
if len(opts.ErrorTemplate) > 0 {
|
||||
bot.errorTemplate = opts.ErrorTemplate
|
||||
@@ -107,8 +136,7 @@ func NewBot[T any](opts *BotOpts) *Bot[T] {
|
||||
|
||||
u, err := api.GetMe()
|
||||
if err != nil {
|
||||
_ = api.CloseApi()
|
||||
_ = uploader.Close()
|
||||
_ = bot.Close()
|
||||
bot.logger.Fatal(err)
|
||||
}
|
||||
bot.logger.Infof("Authorized as %s\n", u.FirstName)
|
||||
@@ -161,21 +189,30 @@ func (bot *Bot[T]) initLoggers(opts *BotOpts) {
|
||||
}
|
||||
}
|
||||
|
||||
func (bot *Bot[T]) GetUpdateOffset() int { return bot.updateOffset }
|
||||
func (bot *Bot[T]) SetUpdateOffset(offset int) { bot.updateOffset = offset }
|
||||
func (bot *Bot[T]) GetUpdateTypes() []tgapi.UpdateType { return bot.updateTypes }
|
||||
func (bot *Bot[T]) GetQueue() *extypes.Queue[*tgapi.Update] { return bot.updateQueue }
|
||||
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]) GetUpdateOffset() int {
|
||||
bot.updateOffsetMu.Lock()
|
||||
defer bot.updateOffsetMu.Unlock()
|
||||
return bot.updateOffset
|
||||
}
|
||||
func (bot *Bot[T]) SetUpdateOffset(offset int) {
|
||||
bot.updateOffsetMu.Lock()
|
||||
defer bot.updateOffsetMu.Unlock()
|
||||
bot.updateOffset = offset
|
||||
}
|
||||
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]) AddDatabaseLogger(writer func(db *T) slog.LoggerWriter) *Bot[T] {
|
||||
type DbLogger[T DbContext] func(db *T) slog.LoggerWriter
|
||||
|
||||
func (bot *Bot[T]) AddDatabaseLoggerWriter(writer DbLogger[T]) *Bot[T] {
|
||||
w := writer(bot.dbContext)
|
||||
bot.logger.AddWriter(w)
|
||||
if bot.RequestLogger != nil {
|
||||
bot.RequestLogger.AddWriter(w)
|
||||
}
|
||||
for _, l := range bot.dbWriterRequested {
|
||||
for _, l := range bot.extraLoggers {
|
||||
l.AddWriter(w)
|
||||
}
|
||||
return bot
|
||||
@@ -209,7 +246,7 @@ func (bot *Bot[T]) Debug(debug bool) *Bot[T] {
|
||||
func (bot *Bot[T]) AddPlugins(plugin ...*Plugin[T]) *Bot[T] {
|
||||
for _, p := range plugin {
|
||||
bot.plugins = append(bot.plugins, *p)
|
||||
bot.logger.Debugln(fmt.Sprintf("plugins with name \"%s\" registered", p.Name))
|
||||
bot.logger.Debugln(fmt.Sprintf("plugins with name \"%s\" registered", p.name))
|
||||
}
|
||||
return bot
|
||||
}
|
||||
@@ -240,7 +277,15 @@ func (bot *Bot[T]) AddL10n(l *L10n) *Bot[T] {
|
||||
return bot
|
||||
}
|
||||
|
||||
func (bot *Bot[T]) Run() {
|
||||
func (bot *Bot[T]) enqueueUpdate(u *tgapi.Update) error {
|
||||
select {
|
||||
case bot.updateQueue <- u:
|
||||
return nil
|
||||
default:
|
||||
return extypes.QueueFullErr
|
||||
}
|
||||
}
|
||||
func (bot *Bot[T]) RunWithContext(ctx context.Context) {
|
||||
if len(bot.prefixes) == 0 {
|
||||
bot.logger.Fatalln("no prefixes defined")
|
||||
return
|
||||
@@ -256,26 +301,36 @@ func (bot *Bot[T]) Run() {
|
||||
bot.logger.Infoln("Bot running. Press CTRL+C to exit.")
|
||||
go func() {
|
||||
for {
|
||||
_, err := bot.Updates()
|
||||
if err != nil {
|
||||
bot.logger.Errorln(err)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
updates, err := bot.Updates()
|
||||
if err != nil {
|
||||
bot.logger.Errorln(err)
|
||||
continue
|
||||
}
|
||||
|
||||
for _, u := range updates {
|
||||
select {
|
||||
case bot.updateQueue <- new(u):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
queue := bot.updateQueue
|
||||
if queue.IsEmpty() {
|
||||
time.Sleep(time.Millisecond * 25)
|
||||
continue
|
||||
}
|
||||
|
||||
u := queue.Dequeue()
|
||||
if u == nil {
|
||||
bot.logger.Errorln("update is nil")
|
||||
continue
|
||||
}
|
||||
|
||||
bot.handle(u)
|
||||
pool := pond.NewPool(16)
|
||||
for update := range bot.updateQueue {
|
||||
update := update
|
||||
log.Println(update)
|
||||
pool.Submit(func() {
|
||||
bot.handle(update)
|
||||
})
|
||||
}
|
||||
}
|
||||
func (bot *Bot[T]) Run() {
|
||||
bot.RunWithContext(context.Background())
|
||||
}
|
||||
|
||||
+5
-1
@@ -27,7 +27,7 @@ func generateBotCommand[T any](cmd Command[T]) tgapi.BotCommand {
|
||||
|
||||
func generateBotCommandForPlugin[T any](pl Plugin[T]) []tgapi.BotCommand {
|
||||
commands := make([]tgapi.BotCommand, 0)
|
||||
for _, cmd := range pl.Commands {
|
||||
for _, cmd := range pl.commands {
|
||||
if cmd.skipAutoCmd {
|
||||
continue
|
||||
}
|
||||
@@ -46,6 +46,10 @@ func (bot *Bot[T]) AutoGenerateCommands() error {
|
||||
|
||||
commands := make([]tgapi.BotCommand, 0)
|
||||
for _, pl := range bot.plugins {
|
||||
if pl.skipAutoCmd {
|
||||
continue
|
||||
}
|
||||
|
||||
commands = append(commands, generateBotCommandForPlugin(pl)...)
|
||||
}
|
||||
if len(commands) > 100 {
|
||||
|
||||
@@ -3,8 +3,10 @@ module git.nix13.pw/scuroneko/laniakea
|
||||
go 1.26
|
||||
|
||||
require (
|
||||
git.nix13.pw/scuroneko/extypes v1.2.0
|
||||
git.nix13.pw/scuroneko/extypes v1.2.1
|
||||
git.nix13.pw/scuroneko/slog v1.0.2
|
||||
github.com/alitto/pond/v2 v2.6.2
|
||||
golang.org/x/time v0.14.0
|
||||
)
|
||||
|
||||
require (
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
git.nix13.pw/scuroneko/extypes v1.2.0 h1:2n2hD6KsMAted+6MGhAyeWyli2Qzc9G2y+pQNB7C1dM=
|
||||
git.nix13.pw/scuroneko/extypes v1.2.0/go.mod h1:uZVs8Yo3RrYAG9dMad6qR6lsYY67t+459D9c65QAYAw=
|
||||
git.nix13.pw/scuroneko/extypes v1.2.1 h1:IYrOjnWKL2EAuJYtYNa+luB1vBe6paE8VY/YD+5/RpQ=
|
||||
git.nix13.pw/scuroneko/extypes v1.2.1/go.mod h1:uZVs8Yo3RrYAG9dMad6qR6lsYY67t+459D9c65QAYAw=
|
||||
git.nix13.pw/scuroneko/slog v1.0.2 h1:vZyUROygxC2d5FJHUQM/30xFEHY1JT/aweDZXA4rm2g=
|
||||
git.nix13.pw/scuroneko/slog v1.0.2/go.mod h1:3Qm2wzkR5KjwOponMfG7TcGSDjmYaFqRAmLvSPTuWJI=
|
||||
github.com/alitto/pond/v2 v2.6.2 h1:Sphe40g0ILeM1pA2c2K+Th0DGU+pt0A/Kprr+WB24Pw=
|
||||
github.com/alitto/pond/v2 v2.6.2/go.mod h1:xkjYEgQ05RSpWdfSd1nM3OVv7TBhLdy7rMp3+2Nq+yE=
|
||||
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
|
||||
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
|
||||
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
||||
@@ -11,3 +13,5 @@ github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
||||
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
|
||||
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
|
||||
|
||||
+2
-2
@@ -46,7 +46,7 @@ func (bot *Bot[T]) handleMessage(update *tgapi.Update, ctx *MsgContext) {
|
||||
text = strings.TrimSpace(text[len(prefix):])
|
||||
|
||||
for _, plugin := range bot.plugins {
|
||||
for cmd := range plugin.Commands {
|
||||
for cmd := range plugin.commands {
|
||||
if !strings.HasPrefix(text, cmd) {
|
||||
continue
|
||||
}
|
||||
@@ -96,7 +96,7 @@ func (bot *Bot[T]) handleCallback(update *tgapi.Update, ctx *MsgContext) {
|
||||
ctx.Args = data.Args
|
||||
|
||||
for _, plugin := range bot.plugins {
|
||||
_, ok := plugin.Payloads[data.Command]
|
||||
_, ok := plugin.payloads[data.Command]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
+5
-21
@@ -2,9 +2,6 @@ package laniakea
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"git.nix13.pw/scuroneko/laniakea/tgapi"
|
||||
)
|
||||
@@ -22,14 +19,8 @@ func (bot *Bot[T]) Updates() ([]tgapi.Update, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, u := range updates {
|
||||
bot.SetUpdateOffset(u.UpdateID + 1)
|
||||
err = bot.GetQueue().Enqueue(&u)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if bot.RequestLogger != nil {
|
||||
if bot.RequestLogger != nil {
|
||||
for _, u := range updates {
|
||||
j, err := json.Marshal(u)
|
||||
if err != nil {
|
||||
bot.GetLogger().Error(err)
|
||||
@@ -37,15 +28,8 @@ func (bot *Bot[T]) Updates() ([]tgapi.Update, error) {
|
||||
bot.RequestLogger.Debugf("UPDATE %s\n", j)
|
||||
}
|
||||
}
|
||||
if len(updates) > 0 {
|
||||
bot.SetUpdateOffset(updates[len(updates)-1].UpdateID + 1)
|
||||
}
|
||||
return updates, err
|
||||
}
|
||||
|
||||
func (bot *Bot[T]) GetFileByLink(link string) ([]byte, error) {
|
||||
u := fmt.Sprintf("https://api.telegram.org/file/bot%s/%s", bot.token, link)
|
||||
res, err := http.Get(u)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
return io.ReadAll(res.Body)
|
||||
}
|
||||
|
||||
+16
-11
@@ -93,37 +93,42 @@ func (c *Command[T]) validateArgs(args []string) error {
|
||||
}
|
||||
|
||||
type Plugin[T DbContext] struct {
|
||||
Name string
|
||||
Commands map[string]Command[T]
|
||||
Payloads map[string]Command[T]
|
||||
Middlewares extypes.Slice[Middleware[T]]
|
||||
name string
|
||||
commands map[string]Command[T]
|
||||
payloads map[string]Command[T]
|
||||
middlewares extypes.Slice[Middleware[T]]
|
||||
skipAutoCmd bool
|
||||
}
|
||||
|
||||
func NewPlugin[T DbContext](name string) *Plugin[T] {
|
||||
return &Plugin[T]{
|
||||
name, map[string]Command[T]{},
|
||||
map[string]Command[T]{}, extypes.Slice[Middleware[T]]{},
|
||||
map[string]Command[T]{}, extypes.Slice[Middleware[T]]{}, false,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Plugin[T]) AddCommand(command *Command[T]) *Plugin[T] {
|
||||
p.Commands[command.command] = *command
|
||||
p.commands[command.command] = *command
|
||||
return p
|
||||
}
|
||||
func (p *Plugin[T]) NewCommand(exec CommandExecutor[T], command string, args ...CommandArg) *Command[T] {
|
||||
return NewCommand(exec, command, args...)
|
||||
}
|
||||
func (p *Plugin[T]) AddPayload(command *Command[T]) *Plugin[T] {
|
||||
p.Payloads[command.command] = *command
|
||||
p.payloads[command.command] = *command
|
||||
return p
|
||||
}
|
||||
func (p *Plugin[T]) AddMiddleware(middleware Middleware[T]) *Plugin[T] {
|
||||
p.Middlewares = p.Middlewares.Push(middleware)
|
||||
p.middlewares = p.middlewares.Push(middleware)
|
||||
return p
|
||||
}
|
||||
func (p *Plugin[T]) SkipCommandAutoGen() *Plugin[T] {
|
||||
p.skipAutoCmd = true
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *Plugin[T]) executeCmd(cmd string, ctx *MsgContext, dbContext *T) {
|
||||
command := p.Commands[cmd]
|
||||
command := p.commands[cmd]
|
||||
if err := command.validateArgs(ctx.Args); err != nil {
|
||||
ctx.error(err)
|
||||
return
|
||||
@@ -131,7 +136,7 @@ func (p *Plugin[T]) executeCmd(cmd string, ctx *MsgContext, dbContext *T) {
|
||||
command.exec(ctx, dbContext)
|
||||
}
|
||||
func (p *Plugin[T]) executePayload(payload string, ctx *MsgContext, dbContext *T) {
|
||||
pl := p.Payloads[payload]
|
||||
pl := p.payloads[payload]
|
||||
if err := pl.validateArgs(ctx.Args); err != nil {
|
||||
ctx.error(err)
|
||||
return
|
||||
@@ -139,7 +144,7 @@ func (p *Plugin[T]) executePayload(payload string, ctx *MsgContext, dbContext *T
|
||||
pl.exec(ctx, dbContext)
|
||||
}
|
||||
func (p *Plugin[T]) executeMiddlewares(ctx *MsgContext, db *T) bool {
|
||||
for _, m := range p.Middlewares {
|
||||
for _, m := range p.middlewares {
|
||||
if !m.Execute(ctx, db) {
|
||||
return false
|
||||
}
|
||||
|
||||
+64
-3
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -11,6 +12,7 @@ import (
|
||||
|
||||
"git.nix13.pw/scuroneko/laniakea/utils"
|
||||
"git.nix13.pw/scuroneko/slog"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
type APIOpts struct {
|
||||
@@ -18,8 +20,13 @@ type APIOpts struct {
|
||||
client *http.Client
|
||||
useTestServer bool
|
||||
apiUrl string
|
||||
|
||||
limiter *rate.Limiter
|
||||
dropOverflowLimit bool
|
||||
}
|
||||
|
||||
var ErrPoolUnexpected = errors.New("unexpected response from pool")
|
||||
|
||||
func NewAPIOpts(token string) *APIOpts {
|
||||
return &APIOpts{token: token, client: nil, useTestServer: false, apiUrl: "https://api.telegram.org"}
|
||||
}
|
||||
@@ -39,6 +46,14 @@ func (opts *APIOpts) SetAPIUrl(apiUrl string) *APIOpts {
|
||||
}
|
||||
return opts
|
||||
}
|
||||
func (opts *APIOpts) SetLimiter(limiter *rate.Limiter) *APIOpts {
|
||||
opts.limiter = limiter
|
||||
return opts
|
||||
}
|
||||
func (opts *APIOpts) SetLimiterDrop(b bool) *APIOpts {
|
||||
opts.dropOverflowLimit = b
|
||||
return opts
|
||||
}
|
||||
|
||||
type API struct {
|
||||
token string
|
||||
@@ -46,6 +61,10 @@ type API struct {
|
||||
logger *slog.Logger
|
||||
useTestServer bool
|
||||
apiUrl string
|
||||
|
||||
pool *WorkerPool
|
||||
limiter *rate.Limiter
|
||||
dropOverflowLimit bool
|
||||
}
|
||||
|
||||
func NewAPI(opts *APIOpts) *API {
|
||||
@@ -55,9 +74,18 @@ func NewAPI(opts *APIOpts) *API {
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: time.Second * 45}
|
||||
}
|
||||
return &API{opts.token, client, l, opts.useTestServer, opts.apiUrl}
|
||||
pool := NewWorkerPool(16, 256)
|
||||
pool.Start(context.Background())
|
||||
return &API{
|
||||
opts.token, client, l,
|
||||
opts.useTestServer, opts.apiUrl,
|
||||
pool, opts.limiter, opts.dropOverflowLimit,
|
||||
}
|
||||
}
|
||||
func (api *API) CloseApi() error {
|
||||
api.pool.Stop()
|
||||
return api.logger.Close()
|
||||
}
|
||||
func (api *API) CloseApi() error { return api.logger.Close() }
|
||||
func (api *API) GetLogger() *slog.Logger { return api.logger }
|
||||
|
||||
type ApiResponse[R any] struct {
|
||||
@@ -74,8 +102,20 @@ type TelegramRequest[R, P any] struct {
|
||||
func NewRequest[R, P any](method string, params P) TelegramRequest[R, P] {
|
||||
return TelegramRequest[R, P]{method: method, params: params}
|
||||
}
|
||||
func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, error) {
|
||||
func (r TelegramRequest[R, P]) doRequest(ctx context.Context, api *API) (R, error) {
|
||||
var zero R
|
||||
if api.limiter != nil {
|
||||
if api.dropOverflowLimit {
|
||||
if !api.limiter.Allow() {
|
||||
return zero, errors.New("rate limited")
|
||||
}
|
||||
} else {
|
||||
if err := api.limiter.Wait(ctx); err != nil {
|
||||
return zero, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
data, err := json.Marshal(r.params)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
@@ -113,7 +153,28 @@ func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R,
|
||||
return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(data))
|
||||
}
|
||||
return parseBody[R](data)
|
||||
}
|
||||
func (r TelegramRequest[R, P]) DoWithContext(ctx context.Context, api *API) (R, error) {
|
||||
var zero R
|
||||
result, err := api.pool.Submit(ctx, func(ctx context.Context) (any, error) {
|
||||
return r.doRequest(ctx, api)
|
||||
})
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return zero, ctx.Err()
|
||||
case res := <-result:
|
||||
if res.Err != nil {
|
||||
return zero, res.Err
|
||||
}
|
||||
if val, ok := res.Value.(R); ok {
|
||||
return val, nil
|
||||
}
|
||||
return zero, ErrPoolUnexpected
|
||||
}
|
||||
}
|
||||
func (r TelegramRequest[R, P]) Do(api *API) (R, error) {
|
||||
return r.DoWithContext(context.Background(), api)
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
package tgapi
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type ParseMode string
|
||||
|
||||
const (
|
||||
@@ -44,3 +50,13 @@ func (api *API) GetFile(params GetFileP) (File, error) {
|
||||
req := NewRequest[File]("getFile", params)
|
||||
return req.Do(api)
|
||||
}
|
||||
|
||||
func (api *API) GetFileByLink(link string) ([]byte, error) {
|
||||
u := fmt.Sprintf("https://api.telegram.org/file/bot%s/%s", api.token, link)
|
||||
res, err := http.Get(u)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
return io.ReadAll(res.Body)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
package tgapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var ErrPoolQueueFull = errors.New("worker pool queue full")
|
||||
|
||||
type RequestEnvelope struct {
|
||||
DoFunc func(context.Context) (any, error) // функция, которая выполнит запрос и вернет any
|
||||
ResultCh chan RequestResult // канал для результата
|
||||
}
|
||||
type RequestResult struct {
|
||||
Value any
|
||||
Err error
|
||||
}
|
||||
|
||||
// WorkerPool управляет воркерами и очередью
|
||||
type WorkerPool struct {
|
||||
taskCh chan RequestEnvelope
|
||||
queueSize int
|
||||
workers int
|
||||
wg sync.WaitGroup
|
||||
quit chan struct{}
|
||||
started bool
|
||||
startedMu sync.Mutex
|
||||
}
|
||||
|
||||
func NewWorkerPool(workers int, queueSize int) *WorkerPool {
|
||||
return &WorkerPool{
|
||||
taskCh: make(chan RequestEnvelope, queueSize),
|
||||
queueSize: queueSize,
|
||||
workers: workers,
|
||||
quit: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start запускает воркеров
|
||||
func (p *WorkerPool) Start(ctx context.Context) {
|
||||
p.startedMu.Lock()
|
||||
defer p.startedMu.Unlock()
|
||||
if p.started {
|
||||
return
|
||||
}
|
||||
p.started = true
|
||||
for i := 0; i < p.workers; i++ {
|
||||
p.wg.Add(1)
|
||||
go p.worker(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop останавливает пул (ждет завершения текущих задач)
|
||||
func (p *WorkerPool) Stop() {
|
||||
close(p.quit)
|
||||
p.wg.Wait()
|
||||
}
|
||||
|
||||
// Submit отправляет задачу в очередь и возвращает канал для результата
|
||||
func (p *WorkerPool) Submit(ctx context.Context, do func(context.Context) (any, error)) (<-chan RequestResult, error) {
|
||||
if len(p.taskCh) >= p.queueSize {
|
||||
return nil, ErrPoolQueueFull
|
||||
}
|
||||
|
||||
resultCh := make(chan RequestResult, 1) // буфер 1, чтобы не блокировать воркера
|
||||
envelope := RequestEnvelope{do, resultCh}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case p.taskCh <- envelope:
|
||||
return resultCh, nil
|
||||
default:
|
||||
return nil, ErrPoolQueueFull
|
||||
}
|
||||
}
|
||||
|
||||
// worker выполняет задачи
|
||||
func (p *WorkerPool) worker(ctx context.Context) {
|
||||
defer p.wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-p.quit:
|
||||
return
|
||||
case envelope := <-p.taskCh:
|
||||
// Выполняем задачу с переданным контекстом (или можно использовать свой)
|
||||
val, err := envelope.DoFunc(ctx)
|
||||
envelope.ResultCh <- RequestResult{Value: val, Err: err}
|
||||
close(envelope.ResultCh)
|
||||
}
|
||||
}
|
||||
}
|
||||
+42
-7
@@ -3,6 +3,7 @@ package tgapi
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
@@ -63,10 +64,21 @@ type UploaderRequest[R, P any] struct {
|
||||
func NewUploaderRequest[R, P any](method string, params P, files ...UploaderFile) UploaderRequest[R, P] {
|
||||
return UploaderRequest[R, P]{method, files, params}
|
||||
}
|
||||
func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader) (R, error) {
|
||||
func (r UploaderRequest[R, P]) doRequest(ctx context.Context, up *Uploader) (R, error) {
|
||||
var zero R
|
||||
if up.api.limiter != nil {
|
||||
if up.api.dropOverflowLimit {
|
||||
if !up.api.limiter.Allow() {
|
||||
return zero, errors.New("rate limited")
|
||||
}
|
||||
} else {
|
||||
if err := up.api.limiter.Wait(ctx); err != nil {
|
||||
return zero, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
buf, contentType, err := prepareMultipart(u.files, u.params)
|
||||
buf, contentType, err := prepareMultipart(r.files, r.params)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
@@ -75,7 +87,7 @@ func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader)
|
||||
if up.api.useTestServer {
|
||||
methodPrefix = "/test"
|
||||
}
|
||||
url := fmt.Sprintf("%s/bot%s%s/%s", up.api.apiUrl, up.api.token, methodPrefix, u.method)
|
||||
url := fmt.Sprintf("%s/bot%s%s/%s", up.api.apiUrl, up.api.token, methodPrefix, r.method)
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", url, buf)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
@@ -84,7 +96,7 @@ func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("User-Agent", fmt.Sprintf("Laniakea/%s", utils.VersionString))
|
||||
|
||||
up.logger.Debugln("UPLOADER REQ", u.method)
|
||||
up.logger.Debugln("UPLOADER REQ", r.method)
|
||||
res, err := up.api.client.Do(req)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
@@ -92,15 +104,38 @@ func (u UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader)
|
||||
defer res.Body.Close()
|
||||
|
||||
body, err := readBody(res.Body)
|
||||
up.logger.Debugln("UPLOADER RES", u.method, string(body))
|
||||
up.logger.Debugln("UPLOADER RES", r.method, string(body))
|
||||
if res.StatusCode != http.StatusOK {
|
||||
return zero, fmt.Errorf("unexpected status code: %d, %s", res.StatusCode, string(body))
|
||||
}
|
||||
|
||||
return parseBody[R](body)
|
||||
}
|
||||
func (u UploaderRequest[R, P]) Do(up *Uploader) (R, error) {
|
||||
return u.DoWithContext(context.Background(), up)
|
||||
func (r UploaderRequest[R, P]) DoWithContext(ctx context.Context, up *Uploader) (R, error) {
|
||||
var zero R
|
||||
|
||||
result, err := up.api.pool.Submit(ctx, func(ctx context.Context) (any, error) {
|
||||
return r.doRequest(ctx, up)
|
||||
})
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return zero, ctx.Err()
|
||||
case res := <-result:
|
||||
if res.Err != nil {
|
||||
return zero, res.Err
|
||||
}
|
||||
if val, ok := res.Value.(R); ok {
|
||||
return val, nil
|
||||
}
|
||||
return zero, ErrPoolUnexpected
|
||||
}
|
||||
}
|
||||
func (r UploaderRequest[R, P]) Do(up *Uploader) (R, error) {
|
||||
return r.DoWithContext(context.Background(), up)
|
||||
}
|
||||
|
||||
func prepareMultipart[P any](files []UploaderFile, params P) (*bytes.Buffer, string, error) {
|
||||
|
||||
+4
-4
@@ -1,9 +1,9 @@
|
||||
package utils
|
||||
|
||||
const (
|
||||
VersionString = "0.8.0-beta.1"
|
||||
VersionMajor = 0
|
||||
VersionMinor = 8
|
||||
VersionString = "1.0.0-beta.2"
|
||||
VersionMajor = 1
|
||||
VersionMinor = 0
|
||||
VersionPatch = 0
|
||||
Beta = 1
|
||||
Beta = 2
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user