wip: experimental branch with finite-state machine processor (doesn't work yet)
This commit is contained in:
parent
2cd46585a8
commit
b13cd05b45
3
go.mod
3
go.mod
@ -12,6 +12,7 @@ require (
|
|||||||
github.com/maypok86/otter v1.2.1
|
github.com/maypok86/otter v1.2.1
|
||||||
github.com/mymmrac/telego v0.29.2
|
github.com/mymmrac/telego v0.29.2
|
||||||
github.com/nicksnyder/go-i18n/v2 v2.4.0
|
github.com/nicksnyder/go-i18n/v2 v2.4.0
|
||||||
|
github.com/stretchr/testify v1.9.0
|
||||||
go.uber.org/zap v1.27.0
|
go.uber.org/zap v1.27.0
|
||||||
golang.org/x/text v0.14.0
|
golang.org/x/text v0.14.0
|
||||||
gopkg.in/yaml.v3 v3.0.1
|
gopkg.in/yaml.v3 v3.0.1
|
||||||
@ -26,6 +27,7 @@ require (
|
|||||||
github.com/bytedance/sonic v1.11.3 // indirect
|
github.com/bytedance/sonic v1.11.3 // indirect
|
||||||
github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect
|
github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect
|
||||||
github.com/chenzhuoyu/iasm v0.9.1 // indirect
|
github.com/chenzhuoyu/iasm v0.9.1 // indirect
|
||||||
|
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||||
github.com/dolthub/maphash v0.1.0 // indirect
|
github.com/dolthub/maphash v0.1.0 // indirect
|
||||||
github.com/fasthttp/router v1.5.0 // indirect
|
github.com/fasthttp/router v1.5.0 // indirect
|
||||||
github.com/gammazero/deque v0.2.1 // indirect
|
github.com/gammazero/deque v0.2.1 // indirect
|
||||||
@ -42,6 +44,7 @@ require (
|
|||||||
github.com/klauspost/compress v1.17.6 // indirect
|
github.com/klauspost/compress v1.17.6 // indirect
|
||||||
github.com/klauspost/cpuid/v2 v2.2.6 // indirect
|
github.com/klauspost/cpuid/v2 v2.2.6 // indirect
|
||||||
github.com/lib/pq v1.10.9 // indirect
|
github.com/lib/pq v1.10.9 // indirect
|
||||||
|
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||||
github.com/savsgio/gotils v0.0.0-20240303185622-093b76447511 // indirect
|
github.com/savsgio/gotils v0.0.0-20240303185622-093b76447511 // indirect
|
||||||
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
||||||
github.com/valyala/bytebufferpool v1.0.0 // indirect
|
github.com/valyala/bytebufferpool v1.0.0 // indirect
|
||||||
|
@ -4,13 +4,16 @@ import (
|
|||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/config"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/config"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/fsmwizard"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/logger"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/logger"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/types"
|
||||||
"github.com/mymmrac/telego"
|
"github.com/mymmrac/telego"
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
"gorm.io/gorm"
|
"gorm.io/gorm"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
type App struct {
|
type App struct {
|
||||||
@ -20,6 +23,7 @@ type App struct {
|
|||||||
Config *config.Config
|
Config *config.Config
|
||||||
Repositories *db.Repositories
|
Repositories *db.Repositories
|
||||||
Handlers map[handler.Type]handler.Handler
|
Handlers map[handler.Type]handler.Handler
|
||||||
|
updates types.Queue[*telego.Update]
|
||||||
db *gorm.DB
|
db *gorm.DB
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -113,6 +117,8 @@ func (a *App) initTelegram() error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *App) initHandlers() {
|
func (a *App) initHandlers() {
|
||||||
|
a.updates = types.NewQueue[*telego.Update]()
|
||||||
|
fsmwizard.PopulateStates(a)
|
||||||
a.Handlers = map[handler.Type]handler.Handler{
|
a.Handlers = map[handler.Type]handler.Handler{
|
||||||
handler.Noop: handler.NewNoopHandler(a.Logger, a.Config.Debug),
|
handler.Noop: handler.NewNoopHandler(a.Logger, a.Config.Debug),
|
||||||
handler.Message: handler.NewMessageHandler(a),
|
handler.Message: handler.NewMessageHandler(a),
|
||||||
@ -138,6 +144,17 @@ func (a *App) handler(update telego.Update) handler.Handler {
|
|||||||
return a.Handlers[handler.Noop]
|
return a.Handlers[handler.Noop]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (a *App) processUpdates() {
|
||||||
|
for {
|
||||||
|
item := a.updates.Dequeue()
|
||||||
|
if item == nil {
|
||||||
|
time.Sleep(time.Nanosecond * 200)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
a.processUpdate(*item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (a *App) processUpdate(update telego.Update) {
|
func (a *App) processUpdate(update telego.Update) {
|
||||||
defer func() {
|
defer func() {
|
||||||
if r := recover(); r != nil {
|
if r := recover(); r != nil {
|
||||||
@ -155,9 +172,11 @@ func (a *App) longPoll() error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
go a.processUpdates()
|
||||||
defer a.Telegram.StopLongPolling()
|
defer a.Telegram.StopLongPolling()
|
||||||
for update := range updates {
|
for update := range updates {
|
||||||
go a.processUpdate(update)
|
update := update
|
||||||
|
a.updates.Enqueue(&update)
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@ -182,8 +201,10 @@ func (a *App) listenWebhook() error {
|
|||||||
defer func() {
|
defer func() {
|
||||||
_ = a.Telegram.StopWebhook()
|
_ = a.Telegram.StopWebhook()
|
||||||
}()
|
}()
|
||||||
|
go a.processUpdates()
|
||||||
for update := range updates {
|
for update := range updates {
|
||||||
go a.processUpdate(update)
|
update := update
|
||||||
|
a.updates.Enqueue(&update)
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
@ -2,6 +2,7 @@ package handler
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/fsmwizard"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
|
||||||
@ -21,18 +22,25 @@ func NewChatMemberUpdatedHandler(app iface.App) *ChatMemberUpdatedHandler {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (h *ChatMemberUpdatedHandler) Handle(wh telego.Update) error {
|
func (h *ChatMemberUpdatedHandler) Handle(wh telego.Update) error {
|
||||||
cm := wh.MyChatMember
|
return fsmwizard.Get(wh.MyChatMember.From.ID).Handle(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
|
||||||
if !cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == h.App.TGProfile().ID &&
|
if w == nil {
|
||||||
cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == h.App.TGProfile().ID {
|
return &fsmwizard.Wizard{Data: wh}
|
||||||
return h.handleAddToChat(wh.MyChatMember.Chat)
|
}
|
||||||
}
|
w.Data = wh
|
||||||
if cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == h.App.TGProfile().ID &&
|
return w
|
||||||
!cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == h.App.TGProfile().ID {
|
})
|
||||||
return h.handleRemoveFromChat(wh.MyChatMember.Chat)
|
// cm := wh.MyChatMember
|
||||||
}
|
// if !cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == h.App.TGProfile().ID &&
|
||||||
|
// cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == h.App.TGProfile().ID {
|
||||||
h.leaveChat(wh.MyChatMember.Chat.ID)
|
// return h.handleAddToChat(wh.MyChatMember.Chat)
|
||||||
return nil
|
// }
|
||||||
|
// if cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == h.App.TGProfile().ID &&
|
||||||
|
// !cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == h.App.TGProfile().ID {
|
||||||
|
// return h.handleRemoveFromChat(wh.MyChatMember.Chat)
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// h.leaveChat(wh.MyChatMember.Chat.ID)
|
||||||
|
// return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *ChatMemberUpdatedHandler) handleAddToChat(tgChat telego.Chat) error {
|
func (h *ChatMemberUpdatedHandler) handleAddToChat(tgChat telego.Chat) error {
|
||||||
|
151
internal/handler/fsmwizard/add_chat_member_state.go
Normal file
151
internal/handler/fsmwizard/add_chat_member_state.go
Normal file
@ -0,0 +1,151 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
"golang.org/x/text/language"
|
||||||
|
)
|
||||||
|
|
||||||
|
const MaxChatMembers = 32
|
||||||
|
|
||||||
|
const AddChatMemberStateID fsm.StateID = "add_chat_member"
|
||||||
|
|
||||||
|
type AddChatMemberState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewAddChatMemberState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &AddChatMemberState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) Enter(pl *Wizard, mc fsm.MachineControls[Wizard]) error {
|
||||||
|
if pl.Data.MyChatMember == nil {
|
||||||
|
s.Move(mc, WaitingForMemberWebhookStateID, pl)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) Handle(pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
next := WaitingForMemberWebhookStateID
|
||||||
|
defer func() {
|
||||||
|
s.Move(mc, next, pl)
|
||||||
|
}()
|
||||||
|
|
||||||
|
cr := s.App.DB().ForChat()
|
||||||
|
tgChat := pl.Data.MyChatMember.Chat
|
||||||
|
user, chat := s.getUserAndChat(tgChat)
|
||||||
|
if user == nil || chat == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
s.Payload.Loc = s.Localizer(user.Language)
|
||||||
|
if s.handleTooManyMembers(tgChat.ID) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if chat.ID == 0 {
|
||||||
|
chat = &model.Chat{
|
||||||
|
TelegramID: tgChat.ID,
|
||||||
|
UserID: user.ID,
|
||||||
|
}
|
||||||
|
err := cr.Save(chat)
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(tgChat.ID, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
chat.UserID = user.ID
|
||||||
|
err := s.App.DB().ForIntegration().DeleteForChat(chat.ID)
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(tgChat.ID, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
err = cr.Save(chat)
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(tgChat.ID, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(user.ChatID),
|
||||||
|
Text: s.Payload.Loc.Template("bot_was_added", map[string]interface{}{"Name": tgChat.Title}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
pl.User = user
|
||||||
|
pl.UserID = user.TelegramID
|
||||||
|
pl.Chat = chat
|
||||||
|
pl.ChatID = tgChat.ID
|
||||||
|
pl.TGChat = tgChat
|
||||||
|
pl.IsMaster = true
|
||||||
|
next = KeyboardChooserStateID
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) handleTooManyMembers(chatID int64) bool {
|
||||||
|
totalMembers, err := s.App.TG().GetChatMemberCount(&telego.GetChatMemberCountParams{
|
||||||
|
ChatID: tu.ID(chatID),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(chatID, err)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if *totalMembers > MaxChatMembers {
|
||||||
|
_, err = s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(chatID),
|
||||||
|
Text: s.Payload.Loc.Template(
|
||||||
|
"too_many_members_in_the_group", map[string]interface{}{"Limit": MaxChatMembers}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.leaveChat(chatID)
|
||||||
|
s.LogError(err)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) getUserAndChat(tgChat telego.Chat) (*model.User, *model.Chat) {
|
||||||
|
user, err := s.App.DB().ForUser().ByTelegramID(s.Payload.Data.MyChatMember.From.ID)
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(tgChat.ID, err)
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
if user == nil || user.ID == 0 {
|
||||||
|
_, err = s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(tgChat.ID),
|
||||||
|
Text: s.Localizer(language.English.String()).
|
||||||
|
Template("you_should_register_first", map[string]interface{}{"Name": s.App.TGProfile().Username}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.leaveChat(tgChat.ID)
|
||||||
|
s.LogError(err)
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
cr := s.App.DB().ForChat()
|
||||||
|
chat, err := cr.ByTelegramID(tgChat.ID)
|
||||||
|
if err != nil {
|
||||||
|
s.curseAndLeave(tgChat.ID, err)
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return user, chat
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) curseAndLeave(chatID int64, err error) {
|
||||||
|
s.LogError(util.SendInternalError(s.App.TG(), chatID, s.Localizer()))
|
||||||
|
s.leaveChat(chatID)
|
||||||
|
s.LogError(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) leaveChat(chatID int64) {
|
||||||
|
s.LogError(s.App.TG().LeaveChat(&telego.LeaveChatParams{
|
||||||
|
ChatID: tu.ID(chatID),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AddChatMemberState) ID() fsm.StateID {
|
||||||
|
return AddChatMemberStateID
|
||||||
|
}
|
88
internal/handler/fsmwizard/base.go
Normal file
88
internal/handler/fsmwizard/base.go
Normal file
@ -0,0 +1,88 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
"golang.org/x/text/language"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Wizard struct {
|
||||||
|
UserID int64
|
||||||
|
ChatID int64
|
||||||
|
TGChat telego.Chat
|
||||||
|
User *model.User
|
||||||
|
Chat *model.Chat
|
||||||
|
Loc locale.Localizer
|
||||||
|
Data telego.Update
|
||||||
|
KeyClick util.Payload
|
||||||
|
IsMaster bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type State struct {
|
||||||
|
fsm.State[Wizard]
|
||||||
|
App iface.App
|
||||||
|
}
|
||||||
|
|
||||||
|
func newBase(app iface.App) State {
|
||||||
|
return State{App: app}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s State) Localizer(langCode ...string) locale.Localizer {
|
||||||
|
lang := language.English.String()
|
||||||
|
if len(langCode) > 0 {
|
||||||
|
lang = langCode[0]
|
||||||
|
} else {
|
||||||
|
if s.Payload != nil && s.Payload.User != nil {
|
||||||
|
lang = s.Payload.User.Language
|
||||||
|
} else if s.Payload != nil && s.Payload.UserID > 0 {
|
||||||
|
user, _ := s.App.DB().ForUser().ByTelegramID(s.Payload.UserID)
|
||||||
|
if user != nil && user.Language != "" {
|
||||||
|
lang = user.Language
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if s.Payload != nil && s.Payload.Loc != nil && (len(langCode) == 0 || s.hasLocalizerWithLocale(lang)) {
|
||||||
|
return s.Payload.Loc
|
||||||
|
}
|
||||||
|
lang = strings.ToLower(lang)
|
||||||
|
if len(lang) > 2 {
|
||||||
|
lang = lang[:2]
|
||||||
|
}
|
||||||
|
switch lang {
|
||||||
|
case "en", "ru":
|
||||||
|
return s.App.Localizer(lang)
|
||||||
|
default:
|
||||||
|
return s.App.Localizer("en")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s State) hasLocalizerWithLocale(lang string) bool {
|
||||||
|
if s.Payload.Loc == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return s.Payload.Loc.Tag().String() == lang
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s State) LogError(err error) {
|
||||||
|
if err == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if s.App.Conf().Debug {
|
||||||
|
s.App.Log().Errorf("handler error: %s, payload: %#v", err, s.Payload)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.App.Log().Errorf("handler error: %s, user id: %d, chat id: %d", err, s.Payload.UserID, s.Payload.ChatID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s State) Move(mc fsm.MachineControls[Wizard], stateID fsm.StateID, pl *Wizard) {
|
||||||
|
s.LogError(mc.Move(stateID, pl))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s State) MoveForHandle(mc fsm.MachineControls[Wizard], stateID fsm.StateID, pl *Wizard) {
|
||||||
|
s.LogError(mc.MoveForHandle(stateID, pl))
|
||||||
|
}
|
47
internal/handler/fsmwizard/configure_redmine_query_state.go
Normal file
47
internal/handler/fsmwizard/configure_redmine_query_state.go
Normal file
@ -0,0 +1,47 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const ConfigureRedmineQueryStateID fsm.StateID = "configure_redmine_query"
|
||||||
|
|
||||||
|
type ConfigureRedmineQueryState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewConfigureRedmineQueryState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &ConfigureRedmineQueryState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *ConfigureRedmineQueryState) Enter(pl *Wizard, _ fsm.MachineControls[Wizard]) error {
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(pl.User.ChatID),
|
||||||
|
Text: pl.Loc.Message("ask_for_redmine"),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
ReplyMarkup: &telego.InlineKeyboardMarkup{InlineKeyboard: [][]telego.InlineKeyboardButton{
|
||||||
|
{
|
||||||
|
{
|
||||||
|
Text: pl.Loc.Message("yes"),
|
||||||
|
CallbackData: util.NewRedmineQuestionPayload(
|
||||||
|
pl.User.TelegramID, pl.Chat.TelegramID, true).String(),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Text: pl.Loc.Message("no"),
|
||||||
|
CallbackData: util.NewRedmineQuestionPayload(
|
||||||
|
pl.User.TelegramID, pl.Chat.TelegramID, false).String(),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}},
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *ConfigureRedmineQueryState) ID() fsm.StateID {
|
||||||
|
return ConfigureRedmineQueryStateID
|
||||||
|
}
|
3
internal/handler/fsmwizard/doc.go
Normal file
3
internal/handler/fsmwizard/doc.go
Normal file
@ -0,0 +1,3 @@
|
|||||||
|
// Package fsmwizard contains setup wizard implementation written with finite-state machine.
|
||||||
|
// Eventually it will replace wizard package (package name will be changed to wizard then).
|
||||||
|
package fsmwizard
|
21
internal/handler/fsmwizard/error_state.go
Normal file
21
internal/handler/fsmwizard/error_state.go
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ErrorState struct {
|
||||||
|
fsm.ErrorState[Wizard]
|
||||||
|
log *zap.SugaredLogger
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewErrorState(log *zap.SugaredLogger) *ErrorState {
|
||||||
|
return &ErrorState{log: log}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *ErrorState) Handle(err error, cur fsm.StateID, next fsm.StateID, pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
s.log.Errorf("critical wizard error: %s, current state: %s, next state: %s, payload: %#v",
|
||||||
|
err, cur, next, pl)
|
||||||
|
mc.Reset()
|
||||||
|
}
|
33
internal/handler/fsmwizard/help_state.go
Normal file
33
internal/handler/fsmwizard/help_state.go
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const HelpStateID fsm.StateID = "help"
|
||||||
|
|
||||||
|
type HelpState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewHelpState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &HelpState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *HelpState) Enter(pl *Wizard, _ fsm.MachineControls[Wizard]) error {
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(pl.Data.Message.Chat.ID),
|
||||||
|
Text: s.Localizer(pl.Data.Message.From.LanguageCode).
|
||||||
|
Template("help_output", map[string]interface{}{"Name": s.App.TGProfile().Username}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
return fsm.ErrPreventTransition
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *HelpState) ID() fsm.StateID {
|
||||||
|
return HelpStateID
|
||||||
|
}
|
99
internal/handler/fsmwizard/init.go
Normal file
99
internal/handler/fsmwizard/init.go
Normal file
@ -0,0 +1,99 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/types"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
wizards *types.TTLMap[int64, fsm.IMachine[Wizard]]
|
||||||
|
states []fsm.IState[Wizard]
|
||||||
|
errorState *ErrorState
|
||||||
|
appPointer iface.App
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
wizards = types.NewTTLMap[int64, fsm.IMachine[Wizard]](time.Hour * 24 * 7)
|
||||||
|
}
|
||||||
|
|
||||||
|
func Get(userID int64) fsm.IMachine[Wizard] {
|
||||||
|
if machine, ok := wizards.Get(userID); ok && machine != nil {
|
||||||
|
return machine
|
||||||
|
}
|
||||||
|
machine := newWizard()
|
||||||
|
wizards.Set(userID, machine)
|
||||||
|
return machine
|
||||||
|
}
|
||||||
|
|
||||||
|
func newWizard() fsm.IMachine[Wizard] {
|
||||||
|
return fsm.New[Wizard](states[0].ID(), wizardRouter, states, errorState)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PopulateStates will init all state handlers for future use.
|
||||||
|
func PopulateStates(app iface.App) {
|
||||||
|
appPointer = app
|
||||||
|
states = []fsm.IState[Wizard]{
|
||||||
|
NewRegisterState(app),
|
||||||
|
NewWaitingForMemberWebhookState(app),
|
||||||
|
NewAddChatMemberState(app),
|
||||||
|
NewKeyboardChooserState(app),
|
||||||
|
NewConfigureRedmineQueryState(app),
|
||||||
|
NewRemoveChatMemberState(app),
|
||||||
|
NewHelpState(app),
|
||||||
|
NewUnknownCommandState(app),
|
||||||
|
}
|
||||||
|
errorState = NewErrorState(app.Log())
|
||||||
|
}
|
||||||
|
|
||||||
|
func wizardRouter(w *Wizard, mc fsm.MachineControlsWithState[Wizard]) {
|
||||||
|
if w.Data.Message != nil {
|
||||||
|
if w.UserID == 0 {
|
||||||
|
w.UserID = w.Data.Message.From.ID
|
||||||
|
}
|
||||||
|
if w.ChatID == 0 {
|
||||||
|
w.ChatID = w.Data.Message.Chat.ID
|
||||||
|
}
|
||||||
|
if w.User == nil {
|
||||||
|
user, err := appPointer.DB().ForUser().ByTelegramID(w.Data.Message.From.ID)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if user != nil && user.ID > 0 {
|
||||||
|
w.User = user
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if w.Chat == nil {
|
||||||
|
chat, err := appPointer.DB().ForChat().ByTelegramID(w.Data.Message.Chat.ID)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if chat != nil && chat.ID > 0 {
|
||||||
|
w.Chat = chat
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if w.TGChat.ID == 0 {
|
||||||
|
w.TGChat = w.Data.Message.Chat
|
||||||
|
}
|
||||||
|
|
||||||
|
switch {
|
||||||
|
case util.MatchCommand("start", w.Data.Message):
|
||||||
|
_ = mc.Move(RegisterStateID, w)
|
||||||
|
case util.MatchCommand("help", w.Data.Message):
|
||||||
|
_ = mc.Move(HelpStateID, w)
|
||||||
|
case util.HasCommand(w.Data.Message):
|
||||||
|
_ = mc.Move(UnknownCommandStateID, w)
|
||||||
|
default:
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if w.Data.MyChatMember != nil {
|
||||||
|
_ = mc.Move(WaitingForMemberWebhookStateID, w)
|
||||||
|
}
|
||||||
|
}
|
91
internal/handler/fsmwizard/keyboard_chooser_state.go
Normal file
91
internal/handler/fsmwizard/keyboard_chooser_state.go
Normal file
@ -0,0 +1,91 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const KeyboardChooserStateID fsm.StateID = "keyboard_chooser"
|
||||||
|
|
||||||
|
type KeyboardChooserState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewKeyboardChooserState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &KeyboardChooserState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *KeyboardChooserState) Enter(pl *Wizard, _ fsm.MachineControls[Wizard]) error {
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(pl.User.ChatID),
|
||||||
|
Text: s.Localizer().Template("choose_keyboard", map[string]interface{}{"Name": pl.TGChat.Title}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
ReplyMarkup: &telego.InlineKeyboardMarkup{
|
||||||
|
InlineKeyboard: [][]telego.InlineKeyboardButton{
|
||||||
|
{
|
||||||
|
{
|
||||||
|
Text: s.Localizer().Message("standard_vote_keyboard"),
|
||||||
|
CallbackData: util.NewKeyboardChooserPayload(
|
||||||
|
pl.User.TelegramID, pl.TGChat.ID, uint8(model.StandardKeyboard)).String(),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
{
|
||||||
|
Text: s.Localizer().Message("sp_vote_keyboard"),
|
||||||
|
CallbackData: util.NewKeyboardChooserPayload(
|
||||||
|
pl.User.TelegramID, pl.TGChat.ID, uint8(model.StoryPointsKeyboard)).String(),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *KeyboardChooserState) Handle(pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
cr := s.App.DB().ForChat()
|
||||||
|
result := pl.KeyClick.KeyboardChoice()
|
||||||
|
if pl.KeyClick.User != pl.User.TelegramID {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
chat, err := cr.ByTelegramID(pl.KeyClick.Chat)
|
||||||
|
if err != nil {
|
||||||
|
s.LogError(err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if chat == nil || chat.ID == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
chat.KeyboardType = model.KeyboardType(result.Type)
|
||||||
|
if err := cr.Save(chat); err != nil {
|
||||||
|
s.LogError(err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
kbTypeName := "standard_vote_keyboard"
|
||||||
|
if model.KeyboardType(result.Type) == model.StoryPointsKeyboard {
|
||||||
|
kbTypeName = "sp_vote_keyboard"
|
||||||
|
}
|
||||||
|
loc := s.Localizer(pl.User.Language)
|
||||||
|
_, err = s.App.TG().EditMessageText(&telego.EditMessageTextParams{
|
||||||
|
ChatID: tu.ID(pl.User.ChatID),
|
||||||
|
MessageID: pl.Data.CallbackQuery.Message.GetMessageID(),
|
||||||
|
Text: loc.Template("chosen_keyboard", map[string]interface{}{"Name": loc.Message(kbTypeName)}),
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
|
||||||
|
if !s.Payload.IsMaster {
|
||||||
|
mc.Reset()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.Move(mc, ConfigureRedmineQueryStateID, pl)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *KeyboardChooserState) ID() fsm.StateID {
|
||||||
|
return KeyboardChooserStateID
|
||||||
|
}
|
87
internal/handler/fsmwizard/register_state.go
Normal file
87
internal/handler/fsmwizard/register_state.go
Normal file
@ -0,0 +1,87 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db/model"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const RegisterStateID fsm.StateID = "register"
|
||||||
|
|
||||||
|
type RegisterState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRegisterState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &RegisterState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RegisterState) Enter(pl *Wizard, _ fsm.MachineControls[Wizard]) error {
|
||||||
|
if pl.Data.Message != nil && pl.Data.Message.Chat.Type != telego.ChatTypePrivate {
|
||||||
|
return fsm.ErrPreventTransition
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RegisterState) Handle(pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
loc := s.Localizer(pl.Data.Message.From.LanguageCode)
|
||||||
|
userRepo := s.App.DB().ForUser()
|
||||||
|
user, err := userRepo.ByTelegramID(pl.Data.Message.From.ID)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if user != nil && user.ID > 0 {
|
||||||
|
var shouldUpdate bool
|
||||||
|
if user.ChatID != pl.Data.Message.Chat.ID {
|
||||||
|
user.ChatID = pl.Data.Message.Chat.ID
|
||||||
|
shouldUpdate = true
|
||||||
|
}
|
||||||
|
if user.Language != pl.Data.Message.From.LanguageCode {
|
||||||
|
user.Language = pl.Data.Message.From.LanguageCode
|
||||||
|
shouldUpdate = true
|
||||||
|
}
|
||||||
|
if shouldUpdate {
|
||||||
|
s.LogError(userRepo.Save(user))
|
||||||
|
}
|
||||||
|
pl.User = user
|
||||||
|
s.Move(mc, WaitingForMemberWebhookStateID, pl)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
err = userRepo.Save(&model.User{
|
||||||
|
TelegramID: pl.Data.Message.From.ID,
|
||||||
|
ChatID: pl.Data.Message.Chat.ID,
|
||||||
|
Language: pl.Data.Message.From.LanguageCode,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
s.LogError(err)
|
||||||
|
_, err = s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(pl.Data.Message.Chat.ID),
|
||||||
|
Text: loc.Message("internal_error"),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
}
|
||||||
|
pl.User = user
|
||||||
|
s.Move(mc, WaitingForMemberWebhookStateID, pl)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RegisterState) Exit(pl *Wizard) {
|
||||||
|
if pl.Data.Message == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
loc := s.Localizer(pl.Data.Message.From.LanguageCode)
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(pl.Data.Message.Chat.ID),
|
||||||
|
Text: loc.Message("welcome"),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RegisterState) ID() fsm.StateID {
|
||||||
|
return RegisterStateID
|
||||||
|
}
|
56
internal/handler/fsmwizard/remove_chat_member_state.go
Normal file
56
internal/handler/fsmwizard/remove_chat_member_state.go
Normal file
@ -0,0 +1,56 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const RemoveChatMemberStateID fsm.StateID = "remove_chat_member"
|
||||||
|
|
||||||
|
type RemoveChatMemberState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRemoveChatMemberState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &RemoveChatMemberState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RemoveChatMemberState) Enter(pl *Wizard, mc fsm.MachineControls[Wizard]) error {
|
||||||
|
if pl.Data.MyChatMember == nil {
|
||||||
|
s.Move(mc, WaitingForMemberWebhookStateID, pl)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RemoveChatMemberState) Handle(pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
if s.Payload.Chat == nil {
|
||||||
|
chat, err := s.App.DB().ForChat().ByTelegramID(s.Payload.Data.MyChatMember.Chat.ID)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if chat != nil && chat.ID > 0 {
|
||||||
|
s.Payload.Chat = chat
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if s.Payload.Chat != nil && s.Payload.Chat.ID > 0 {
|
||||||
|
user, _ := s.App.DB().ForUser().ByID(s.Payload.Chat.UserID)
|
||||||
|
if user != nil && user.ID > 0 && user.ChatID > 0 {
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(user.ChatID),
|
||||||
|
Text: s.Localizer(user.Language).Template(
|
||||||
|
"bot_was_removed_from_group", map[string]interface{}{"Name": pl.Data.MyChatMember.Chat.Title}),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
}
|
||||||
|
s.LogError(s.App.DB().ForChat().Delete(s.Payload.Chat))
|
||||||
|
}
|
||||||
|
s.Move(mc, WaitingForMemberWebhookStateID, pl)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RemoveChatMemberState) ID() fsm.StateID {
|
||||||
|
return RemoveChatMemberStateID
|
||||||
|
}
|
32
internal/handler/fsmwizard/unknown_command_state.go
Normal file
32
internal/handler/fsmwizard/unknown_command_state.go
Normal file
@ -0,0 +1,32 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const UnknownCommandStateID fsm.StateID = "unknown_command"
|
||||||
|
|
||||||
|
type UnknownCommandState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewUnknownCommandState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &UnknownCommandState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *UnknownCommandState) Enter(pl *Wizard, _ fsm.MachineControls[Wizard]) error {
|
||||||
|
_, err := s.App.TG().SendMessage(&telego.SendMessageParams{
|
||||||
|
ChatID: tu.ID(s.Payload.Data.Message.Chat.ID),
|
||||||
|
Text: s.Localizer(s.Payload.Data.Message.From.LanguageCode).Message("unknown_command"),
|
||||||
|
ParseMode: telego.ModeMarkdown,
|
||||||
|
})
|
||||||
|
s.LogError(err)
|
||||||
|
return fsm.ErrPreventTransition
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *UnknownCommandState) ID() fsm.StateID {
|
||||||
|
return UnknownCommandStateID
|
||||||
|
}
|
@ -0,0 +1,51 @@
|
|||||||
|
package fsmwizard
|
||||||
|
|
||||||
|
import (
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/fsm"
|
||||||
|
"github.com/mymmrac/telego"
|
||||||
|
tu "github.com/mymmrac/telego/telegoutil"
|
||||||
|
)
|
||||||
|
|
||||||
|
const WaitingForMemberWebhookStateID fsm.StateID = "waiting_for_member_webhook"
|
||||||
|
|
||||||
|
type WaitingForMemberWebhookState struct {
|
||||||
|
State
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewWaitingForMemberWebhookState(app iface.App) fsm.IState[Wizard] {
|
||||||
|
return &WaitingForMemberWebhookState{newBase(app)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *WaitingForMemberWebhookState) Enter(pl *Wizard, mc fsm.MachineControls[Wizard]) error {
|
||||||
|
cm := pl.Data.MyChatMember
|
||||||
|
if cm == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == s.App.TGProfile().ID &&
|
||||||
|
cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == s.App.TGProfile().ID {
|
||||||
|
s.MoveForHandle(mc, AddChatMemberStateID, pl)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if cm.OldChatMember.MemberIsMember() && cm.OldChatMember.MemberUser().ID == s.App.TGProfile().ID &&
|
||||||
|
!cm.NewChatMember.MemberIsMember() && cm.NewChatMember.MemberUser().ID == s.App.TGProfile().ID {
|
||||||
|
s.MoveForHandle(mc, RemoveChatMemberStateID, pl)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
s.LogError(s.App.TG().LeaveChat(&telego.LeaveChatParams{
|
||||||
|
ChatID: tu.ID(pl.Data.MyChatMember.Chat.ID),
|
||||||
|
}))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *WaitingForMemberWebhookState) Handle(pl *Wizard, mc fsm.MachineControls[Wizard]) {
|
||||||
|
// By using this we can handle both Move from Register state and Move from waiting for webhook state.
|
||||||
|
// The first one doesn't have data yet and it will be handled by Handle, the second one has data and will be
|
||||||
|
// handled by Enter.
|
||||||
|
_ = s.Enter(pl, mc)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *WaitingForMemberWebhookState) ID() fsm.StateID {
|
||||||
|
return WaitingForMemberWebhookStateID
|
||||||
|
}
|
@ -1,11 +1,8 @@
|
|||||||
package handler
|
package handler
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/group"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/fsmwizard"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/iface"
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/store"
|
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
|
|
||||||
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/wizard"
|
|
||||||
"github.com/mymmrac/telego"
|
"github.com/mymmrac/telego"
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -22,22 +19,30 @@ func (h *MessageHandler) Handle(wh telego.Update) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
if util.MatchCommand("start", wh.Message) {
|
return fsmwizard.Get(wh.Message.From.ID).Handle(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
|
||||||
return wizard.NewRegister(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
if w == nil {
|
||||||
}
|
return &fsmwizard.Wizard{Data: wh}
|
||||||
|
}
|
||||||
|
w.Data = wh
|
||||||
|
return w
|
||||||
|
})
|
||||||
|
|
||||||
if util.MatchCommand("poll", wh.Message) {
|
// if util.MatchCommand("start", wh.Message) {
|
||||||
return group.NewPoll(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
// return wizard.NewRegister(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
||||||
}
|
// }
|
||||||
|
//
|
||||||
setup, found := store.RedmineSetups.Get(wh.Message.Chat.ID)
|
// if util.MatchCommand("poll", wh.Message) {
|
||||||
if found {
|
// return group.NewPoll(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
||||||
return wizard.NewRedmineSetup(h.App, wh.Message.From.ID, wh.Message.Chat.ID, setup).Handle(wh)
|
// }
|
||||||
}
|
//
|
||||||
|
// setup, found := store.RedmineSetups.Get(wh.Message.Chat.ID)
|
||||||
if util.MatchCommand("help", wh.Message) {
|
// if found {
|
||||||
return wizard.NewHelpCommand(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
// return wizard.NewRedmineSetup(h.App, wh.Message.From.ID, wh.Message.Chat.ID, setup).Handle(wh)
|
||||||
}
|
// }
|
||||||
|
//
|
||||||
return wizard.NewUnknownCommand(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
// if util.MatchCommand("help", wh.Message) {
|
||||||
|
// return wizard.NewHelpCommand(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// return wizard.NewUnknownCommand(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
|
||||||
}
|
}
|
||||||
|
@ -1,2 +1,3 @@
|
|||||||
// Package wizard contains setup wizard implementation.
|
// Package wizard contains setup wizard implementation.
|
||||||
|
// Deprecated: will be replaced with fsmwizard in the future.
|
||||||
package wizard
|
package wizard
|
||||||
|
147
pkg/fsm/machine.go
Normal file
147
pkg/fsm/machine.go
Normal file
@ -0,0 +1,147 @@
|
|||||||
|
package fsm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/types"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
// ErrPreventTransition should be returned from Enter if you don't want to perform a state transition and
|
||||||
|
// everything that needed to be done has been done in Enter callback.
|
||||||
|
ErrPreventTransition = errors.New("prevents transition; this is not an error")
|
||||||
|
// ErrStateDoesNotExist will be returned if provided state ID does not exist in this machine.
|
||||||
|
ErrStateDoesNotExist = errors.New("state does not exist")
|
||||||
|
)
|
||||||
|
|
||||||
|
// MachineStateRouter should be provided to IMachine. This function can do two very useful things:
|
||||||
|
// - It can modify Machine's payload.
|
||||||
|
// - It can act as a router by changing Machine's state via provided controls.
|
||||||
|
type MachineStateRouter[T any] func(*T, MachineControlsWithState[T])
|
||||||
|
|
||||||
|
// MachineStateProvider provided to every Handle call. It can be used to set the initial state of the machine or to
|
||||||
|
// update existing machine state.
|
||||||
|
type MachineStateProvider[T any] func(*T) *T
|
||||||
|
|
||||||
|
// IMachine is a Machine contract. The Machine should be able to do the following:
|
||||||
|
// - Move to another state (usually called by the IState itself).
|
||||||
|
// - Handle the state input.
|
||||||
|
// - Reset the machine.
|
||||||
|
type IMachine[T any] interface {
|
||||||
|
MachineControlsWithState[T]
|
||||||
|
// Handle the state input. Handle func will accept the current payload and modify it based on user input.
|
||||||
|
Handle(MachineStateProvider[T]) error
|
||||||
|
// Reset the machine to its initial state.
|
||||||
|
Reset()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Machine is a finite-state machine implementation.
|
||||||
|
type Machine[T any] struct {
|
||||||
|
payload *T
|
||||||
|
stateRouter MachineStateRouter[T]
|
||||||
|
state StateID
|
||||||
|
initialState StateID
|
||||||
|
states *types.Map[StateID, IState[T]]
|
||||||
|
errHandler ErrorState[T]
|
||||||
|
handleNow bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// New machine.
|
||||||
|
func New[T any](initialState StateID, router MachineStateRouter[T], states []IState[T], errHandler ErrorState[T]) IMachine[T] {
|
||||||
|
stateMap := types.NewMap[StateID, IState[T]]()
|
||||||
|
for _, state := range states {
|
||||||
|
stateMap.Set(state.ID(), state)
|
||||||
|
}
|
||||||
|
return &Machine[T]{
|
||||||
|
state: initialState,
|
||||||
|
stateRouter: router,
|
||||||
|
initialState: initialState,
|
||||||
|
states: stateMap,
|
||||||
|
errHandler: errHandler,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Move to another state.
|
||||||
|
// Internal: should never be called outside state callbacks.
|
||||||
|
func (m *Machine[T]) Move(id StateID, payload *T) error {
|
||||||
|
if id == m.state {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
next, err := m.loadState(id, payload)
|
||||||
|
if next == nil || err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
cur, err := m.loadState(m.state, payload)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if cur != nil {
|
||||||
|
cur.Exit(payload)
|
||||||
|
}
|
||||||
|
if err := next.Enter(payload, m); err != nil {
|
||||||
|
if errors.Is(err, ErrPreventTransition) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
m.Reset()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
m.state = id
|
||||||
|
m.payload = payload
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Machine[T]) MoveForHandle(id StateID, payload *T) error {
|
||||||
|
m.handleNow = true
|
||||||
|
return m.Move(id, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle the input.
|
||||||
|
func (m *Machine[T]) Handle(provider MachineStateProvider[T]) error {
|
||||||
|
if provider != nil {
|
||||||
|
m.payload = provider(m.payload)
|
||||||
|
}
|
||||||
|
if m.stateRouter != nil {
|
||||||
|
m.stateRouter(m.payload, m)
|
||||||
|
}
|
||||||
|
st, err := m.loadState(m.state, m.payload)
|
||||||
|
if st == nil || err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
for {
|
||||||
|
st.Handle(m.payload, m)
|
||||||
|
if m.handleNow { // MoveForHandle was called, trying to handle again.
|
||||||
|
m.handleNow = false
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// State of the Machine.
|
||||||
|
func (m *Machine[T]) State() *T {
|
||||||
|
return m.payload
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset the machine.
|
||||||
|
func (m *Machine[T]) Reset() {
|
||||||
|
m.payload = nil
|
||||||
|
m.state = m.initialState
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Machine[T]) loadState(id StateID, payload *T) (IState[T], error) {
|
||||||
|
if id == NilStateID {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
st, ok := m.states.Get(id)
|
||||||
|
if !ok {
|
||||||
|
return nil, m.fatalError(fmt.Errorf("%w: %s", ErrStateDoesNotExist, id), id, payload)
|
||||||
|
}
|
||||||
|
return st, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Machine[T]) fatalError(err error, id StateID, payload *T) error {
|
||||||
|
if m.errHandler != nil {
|
||||||
|
m.errHandler.Handle(err, m.state, id, payload, m)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
85
pkg/fsm/state.go
Normal file
85
pkg/fsm/state.go
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
package fsm
|
||||||
|
|
||||||
|
// StateID is a state identifier. Machine with string state IDs can have *a lot of* states.
|
||||||
|
type StateID string
|
||||||
|
|
||||||
|
// NilStateID is a noop state. Machine won't do anything for this transition.
|
||||||
|
const NilStateID = StateID("")
|
||||||
|
|
||||||
|
// MachineControls is a fragment of IMachine implementation. This one can Move between the states or Reset the machine.
|
||||||
|
// It may fail with an error which should be handled by the IState implementation.
|
||||||
|
type MachineControls[T any] interface {
|
||||||
|
Move(StateID, *T) error
|
||||||
|
// MoveForHandle is the same as Move but it also triggers Handle immediately for the next state.
|
||||||
|
MoveForHandle(StateID, *T) error
|
||||||
|
Reset()
|
||||||
|
}
|
||||||
|
|
||||||
|
// MachineState returns machine state. This state should be immutable.
|
||||||
|
type MachineState[T any] interface {
|
||||||
|
// State returns underlying state. This state SHOULD NOT be modified since there is no locking performed.
|
||||||
|
State() *T
|
||||||
|
}
|
||||||
|
|
||||||
|
// MachineControlsWithState is self-explanatory.
|
||||||
|
type MachineControlsWithState[T any] interface {
|
||||||
|
MachineControls[T]
|
||||||
|
MachineState[T]
|
||||||
|
}
|
||||||
|
|
||||||
|
// IState is a State interface. This contract enforces that any state implementation should have three methods:
|
||||||
|
// Enter, Handle and Exit. The first one is called right after the machine has moved to the state, the second one is
|
||||||
|
// called when handling some state input, the last one is called right before leaving to a next state (which happens
|
||||||
|
// if Handle has called Move on MachineControls.
|
||||||
|
type IState[T any] interface {
|
||||||
|
// ID should return state identifier.
|
||||||
|
ID() StateID
|
||||||
|
// Enter is a state enter callback. Can be used to perform some sort of input query or to move to another
|
||||||
|
// state immediately. Also, if Enter fails the machine won't move to the state and MachineControls's Move will
|
||||||
|
// return an error which was returned from Enter.
|
||||||
|
Enter(*T, MachineControls[T]) error
|
||||||
|
// Handle is called when receiving some sort of input response. This one's signature is nearly identical to Enter,
|
||||||
|
// but it can wait for some sort of input (standard input? webhook) without locking inside the callback
|
||||||
|
// while Enter cannot do that.
|
||||||
|
Handle(*T, MachineControls[T])
|
||||||
|
// Exit is called right before leaving the state. It can be used to modify state's payload (T) or for some other
|
||||||
|
// miscellaneous tasks.
|
||||||
|
// Note: calling Exit doesn't mean that the machine will really transition to the next state.
|
||||||
|
// The next state's Enter callback can return an error which will reset the Machine to default state & payload.
|
||||||
|
Exit(*T)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ErrorState is the Machine's fatal error handler which you should implement yourself.
|
||||||
|
//
|
||||||
|
// Error state is used by Machine in case of a fatal error (for example, for invalid state ID).
|
||||||
|
// You can use ErrorState to make some kind of fatal error response like "Internal error has occurred"
|
||||||
|
// or something similar. Also, ErrorState is special because neither Enter nor Exit exists in it.
|
||||||
|
//
|
||||||
|
// Machine without ErrorState will not do anything in case of fatal errors.
|
||||||
|
type ErrorState[T any] interface {
|
||||||
|
Handle(err error, current StateID, next StateID, payload *T, machine MachineControls[T])
|
||||||
|
}
|
||||||
|
|
||||||
|
// State is the Machine's state. This implementation doesn't do anything and only helps with the
|
||||||
|
// actual IState implementation (you don't need to write empty Enter and Exit callback if you don't use them).
|
||||||
|
type State[T any] struct {
|
||||||
|
Payload *T
|
||||||
|
}
|
||||||
|
|
||||||
|
// ID panics because you need to implement it.
|
||||||
|
func (s *State[T]) ID() StateID {
|
||||||
|
panic("implement ID() StateID method for your state")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enter here will immediately move the Machine to empty state.
|
||||||
|
func (s *State[T]) Enter(payload *T, machine MachineControls[T]) error {
|
||||||
|
return machine.Move(NilStateID, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle here will immediately move the Machine to empty state.
|
||||||
|
func (s *State[T]) Handle(payload *T, machine MachineControls[T]) {
|
||||||
|
_ = machine.Move(NilStateID, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Exit won't do anything.
|
||||||
|
func (s *State[T]) Exit(*T) {}
|
27
pkg/types/map.go
Normal file
27
pkg/types/map.go
Normal file
@ -0,0 +1,27 @@
|
|||||||
|
package types
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Map[K comparable, V any] struct {
|
||||||
|
rw sync.RWMutex
|
||||||
|
m map[K]V
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewMap[K comparable, V any]() *Map[K, V] {
|
||||||
|
return &Map[K, V]{m: make(map[K]V)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Map[K, V]) Set(k K, v V) {
|
||||||
|
defer t.rw.Unlock()
|
||||||
|
t.rw.Lock()
|
||||||
|
t.m[k] = v
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *Map[K, V]) Get(k K) (val V, ok bool) {
|
||||||
|
defer t.rw.RUnlock()
|
||||||
|
t.rw.RLock()
|
||||||
|
val, ok = t.m[k]
|
||||||
|
return
|
||||||
|
}
|
62
pkg/types/queue.go
Normal file
62
pkg/types/queue.go
Normal file
@ -0,0 +1,62 @@
|
|||||||
|
package types
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync/atomic"
|
||||||
|
"unsafe"
|
||||||
|
)
|
||||||
|
|
||||||
|
type node[T any] struct {
|
||||||
|
val atomic.Pointer[T]
|
||||||
|
next atomic.Pointer[node[T]]
|
||||||
|
}
|
||||||
|
|
||||||
|
type Queue[T any] interface {
|
||||||
|
Enqueue(T)
|
||||||
|
Dequeue() T
|
||||||
|
Len() uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
type queue[T any] struct {
|
||||||
|
head, tail *node[T]
|
||||||
|
length atomic.Uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewQueue[T any]() Queue[T] {
|
||||||
|
n := node[T]{}
|
||||||
|
return &queue[T]{head: &n, tail: &n}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *queue[T]) Enqueue(v T) {
|
||||||
|
n := node[T]{val: atomic.Pointer[T]{}}
|
||||||
|
n.val.Store(&v)
|
||||||
|
for {
|
||||||
|
if q.tail.next.CompareAndSwap(nil, &n) {
|
||||||
|
atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&q.tail)), unsafe.Pointer(&n))
|
||||||
|
q.length.Add(1)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *queue[T]) Dequeue() T {
|
||||||
|
headAddr := (*unsafe.Pointer)(unsafe.Pointer(&q.head))
|
||||||
|
for {
|
||||||
|
var result T
|
||||||
|
head := atomic.LoadPointer(headAddr)
|
||||||
|
n := q.head.next.Load()
|
||||||
|
if n == nil {
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
if atomic.CompareAndSwapPointer(headAddr, head, unsafe.Pointer(n)) {
|
||||||
|
q.length.Add(^uint64(0)) // Переполнение намеренное, это отнимает единицу от счетчика.
|
||||||
|
if r := n.val.Load(); r != nil {
|
||||||
|
return *r
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *queue[T]) Len() uint64 {
|
||||||
|
return q.length.Load()
|
||||||
|
}
|
19
pkg/types/queue_test.go
Normal file
19
pkg/types/queue_test.go
Normal file
@ -0,0 +1,19 @@
|
|||||||
|
package types
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestQueue(t *testing.T) {
|
||||||
|
q := NewQueue[uint8]()
|
||||||
|
assert.Equal(t, uint64(0), q.Len())
|
||||||
|
assert.Equal(t, uint8(0), q.Dequeue())
|
||||||
|
|
||||||
|
q.Enqueue(1)
|
||||||
|
q.Enqueue(2)
|
||||||
|
assert.Equal(t, uint64(2), q.Len())
|
||||||
|
assert.Equal(t, uint8(1), q.Dequeue())
|
||||||
|
assert.Equal(t, uint8(2), q.Dequeue())
|
||||||
|
assert.Equal(t, uint8(0), q.Dequeue())
|
||||||
|
}
|
63
pkg/types/ttl_map.go
Normal file
63
pkg/types/ttl_map.go
Normal file
@ -0,0 +1,63 @@
|
|||||||
|
package types
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type TTLMap[K comparable, V any] struct {
|
||||||
|
rw sync.RWMutex
|
||||||
|
m map[K]entry[V]
|
||||||
|
ttl time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
type entry[V any] struct {
|
||||||
|
value V
|
||||||
|
expireAt *time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *entry[V]) expired() bool {
|
||||||
|
if e.expireAt == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return time.Now().After(*e.expireAt)
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewTTLMap[K comparable, V any](ttl time.Duration) *TTLMap[K, V] {
|
||||||
|
return &TTLMap[K, V]{m: make(map[K]entry[V]), ttl: ttl}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TTLMap[K, V]) Set(k K, v V) {
|
||||||
|
t.SetWithTTL(k, v, t.ttl)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TTLMap[K, V]) SetWithTTL(k K, v V, ttl time.Duration) {
|
||||||
|
defer t.rw.Unlock()
|
||||||
|
t.rw.Lock()
|
||||||
|
var expireAt *time.Time
|
||||||
|
if ttl > 0 {
|
||||||
|
at := time.Now().Add(ttl)
|
||||||
|
expireAt = &at
|
||||||
|
}
|
||||||
|
t.m[k] = entry[V]{value: v, expireAt: expireAt}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *TTLMap[K, V]) Get(k K) (val V, ok bool) {
|
||||||
|
defer t.rw.RUnlock()
|
||||||
|
t.rw.RLock()
|
||||||
|
v, ok := t.m[k]
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if v.expired() {
|
||||||
|
defer func() {
|
||||||
|
t.rw.RUnlock()
|
||||||
|
t.rw.Lock()
|
||||||
|
delete(t.m, k)
|
||||||
|
t.rw.Unlock()
|
||||||
|
t.rw.RLock()
|
||||||
|
}()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
return v.value, true
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user