Compare commits

..

2 Commits
master ... fsm

29 changed files with 1264 additions and 67 deletions

View File

@ -15,19 +15,6 @@ steps:
from_secret: CI_APP_IMAGE
tags:
- latest
- name: deploy
image: alpine:latest
environment:
DEPLOYER_AUTH:
from_secret: DEPLOYER_AUTH
DEPLOYER_URL:
from_secret: DEPLOYER_URL
commands:
- apk add --no-cache curl
- curl -H "$DEPLOYER_AUTH" $DEPLOYER_URL
when:
branch:
- master
trigger:
event:

3
go.mod
View File

@ -12,6 +12,7 @@ require (
github.com/maypok86/otter v1.2.1
github.com/mymmrac/telego v0.29.2
github.com/nicksnyder/go-i18n/v2 v2.4.0
github.com/stretchr/testify v1.9.0
go.uber.org/zap v1.27.0
golang.org/x/text v0.14.0
gopkg.in/yaml.v3 v3.0.1
@ -26,6 +27,7 @@ require (
github.com/bytedance/sonic v1.11.3 // indirect
github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // 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/fasthttp/router v1.5.0 // 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/cpuid/v2 v2.2.6 // 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/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/valyala/bytebufferpool v1.0.0 // indirect

View File

@ -4,6 +4,7 @@ import (
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/config"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/db"
"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/locale"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/logger"
@ -113,6 +114,7 @@ func (a *App) initTelegram() error {
}
func (a *App) initHandlers() {
fsmwizard.PopulateStates(a)
a.Handlers = map[handler.Type]handler.Handler{
handler.Noop: handler.NewNoopHandler(a.Logger, a.Config.Debug),
handler.Message: handler.NewMessageHandler(a),

View File

@ -24,7 +24,7 @@ func (c *Chat) ByID(id uint64) (*model.Chat, error) {
func (c *Chat) ByTelegramID(id int64) (*model.Chat, error) {
var chat model.Chat
if err := c.db.Model(&model.Chat{}).Where("telegram_id = ?", id).First(&chat).Error; err != nil {
if err := c.db.Model(model.Chat{TelegramID: id}).First(&chat).Error; err != nil {
return nil, util.HandleRecordNotFound(err)
}
return &chat, nil
@ -32,7 +32,7 @@ func (c *Chat) ByTelegramID(id int64) (*model.Chat, error) {
func (c *Chat) ByTelegramIDWithIntegrations(id int64) (*model.Chat, error) {
var chat model.Chat
if err := c.db.Model(&model.Chat{}).Where("telegram_id = ?", id).Preload("Integrations").First(&chat).Error; err != nil {
if err := c.db.Model(model.Chat{TelegramID: id}).Preload("Integrations").First(&chat).Error; err != nil {
return nil, util.HandleRecordNotFound(err)
}
return &chat, nil
@ -47,9 +47,6 @@ func (c *Chat) ByIDWithIntegrations(id uint64) (*model.Chat, error) {
}
func (c *Chat) Save(chat *model.Chat) error {
if chat.ID == 0 {
return c.db.Create(chat).Error
}
return c.db.Model(chat).Save(chat).Error
}

View File

@ -16,16 +16,13 @@ func NewIntegration(db *gorm.DB) *Integration {
func (r *Integration) LoadForChatAndType(chatID uint64, typ model.IntegrationType) (*model.Integration, error) {
var integration model.Integration
if err := r.db.Model(&model.Integration{}).Where(`chat_id = ? and "type" = ?`, chatID, typ).First(&integration).Error; err != nil {
if err := r.db.Model(model.Integration{ChatID: chatID, Type: typ}).First(&integration).Error; err != nil {
return nil, util.HandleRecordNotFound(err)
}
return &integration, nil
}
func (r *Integration) Save(integration *model.Integration) error {
if integration.ID == 0 {
return r.db.Create(integration).Error
}
return r.db.Model(integration).Save(integration).Error
}

View File

@ -24,7 +24,7 @@ func (u *User) ByID(id uint64) (*model.User, error) {
func (u *User) ByTelegramID(id int64) (*model.User, error) {
var user model.User
if err := u.db.Model(&model.User{}).Where("telegram_id = ?", id).First(&user).Error; err != nil {
if err := u.db.Model(model.User{TelegramID: id}).First(&user).Error; err != nil {
return nil, util.HandleRecordNotFound(err)
}
return &user, nil
@ -32,7 +32,7 @@ func (u *User) ByTelegramID(id int64) (*model.User, error) {
func (u *User) ByTelegramIDs(ids []int64) ([]model.User, error) {
var users []model.User
if err := u.db.Model(&model.User{}).Where("telegram_id in (?)", ids).Find(&users).Error; err != nil {
if err := u.db.Model(model.User{}).Where("telegram_id in (?)", ids).Find(&users).Error; err != nil {
return nil, util.HandleRecordNotFound(err)
}
return users, nil
@ -47,8 +47,5 @@ func (u *User) ByIDWithChats(id uint64) (*model.User, error) {
}
func (u *User) Save(user *model.User) error {
if user.ID == 0 {
return u.db.Create(user).Error
}
return u.db.Model(user).Save(user).Error
}

View File

@ -2,6 +2,7 @@ package handler
import (
"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/util"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/locale"
@ -21,18 +22,26 @@ func NewChatMemberUpdatedHandler(app iface.App) *ChatMemberUpdatedHandler {
}
func (h *ChatMemberUpdatedHandler) Handle(wh telego.Update) error {
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 {
return h.handleAddToChat(wh.MyChatMember.Chat)
}
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)
fsmwizard.Get(wh.MyChatMember.From.ID).Enqueue(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
if w == nil {
return &fsmwizard.Wizard{Data: wh}
}
w.Data = wh
return w
})
return nil
// 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 {
// return h.handleAddToChat(wh.MyChatMember.Chat)
// }
// 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 {

View 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
}

View 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) MoveForHandling(mc fsm.MachineControls[Wizard], stateID fsm.StateID, pl *Wizard) {
s.LogError(mc.MoveForHandling(stateID, pl))
}

View 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
}

View 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

View 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()
}

View 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
}

View 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)
}
}

View 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
}

View 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
}

View 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
}

View 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
}

View File

@ -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.MoveForHandling(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.MoveForHandling(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
}

View File

@ -41,21 +41,18 @@ func (h *Poll) Handle(wh telego.Update) error {
return err
}
loc := h.Localizer(user.Language)
_ = loc
if len(wh.Message.Entities) == 0 ||
(len(wh.Message.Entities) > 0 && wh.Message.Entities[0].Type != telego.EntityTypeBotCommand) ||
(len(wh.Message.Entities) > 0 && wh.Message.Entities[0].Offset != 0) {
return nil
}
taskInfo := strings.TrimSpace(wh.Message.Text[wh.Message.Entities[0].Length:])
if taskInfo == "" && wh.Message.ReplyToMessage != nil {
taskInfo = wh.Message.ReplyToMessage.Text
}
var (
taskID int
canRedmine bool
)
taskInfo := strings.TrimSpace(wh.Message.Text[wh.Message.Entities[0].Length:])
if taskInfo != "" {
for _, integrationData := range chat.Integrations {
id, info := integration.New(integrationData, h.App.Log()).GetTaskInfo(taskInfo)

View File

@ -1,11 +1,8 @@
package handler
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/store"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/util"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/internal/handler/wizard"
"github.com/mymmrac/telego"
)
@ -22,22 +19,31 @@ func (h *MessageHandler) Handle(wh telego.Update) error {
return nil
}
if util.MatchCommand("start", wh.Message) {
return wizard.NewRegister(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
}
fsmwizard.Get(wh.Message.From.ID).Enqueue(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
if w == nil {
return &fsmwizard.Wizard{Data: wh}
}
w.Data = wh
return w
})
return nil
if util.MatchCommand("poll", wh.Message) {
return group.NewPoll(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
}
setup, found := store.RedmineSetups.Get(wh.Message.Chat.ID)
if found {
return wizard.NewRedmineSetup(h.App, wh.Message.From.ID, wh.Message.Chat.ID, setup).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)
// if util.MatchCommand("start", wh.Message) {
// return wizard.NewRegister(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
// }
//
// if util.MatchCommand("poll", wh.Message) {
// return group.NewPoll(h.App, wh.Message.From.ID, wh.Message.Chat.ID).Handle(wh)
// }
//
// setup, found := store.RedmineSetups.Get(wh.Message.Chat.ID)
// if found {
// return wizard.NewRedmineSetup(h.App, wh.Message.From.ID, wh.Message.Chat.ID, setup).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)
}

View File

@ -44,7 +44,7 @@ func (p Payload) Vote() (val Vote) {
}
func (p Payload) KeyboardChoice() (val KBChooserData) {
if p.Action != PayloadActionChooseKeyboard {
if p.Action != PayloadActionVote {
return
}
_ = json.Unmarshal(p.Data, &val)
@ -65,7 +65,7 @@ type Member struct {
}
type KBChooserData struct {
Type int64 `json:"k"`
Type uint8 `json:"k"`
}
type Vote struct {
@ -83,7 +83,7 @@ func NewKeyboardChooserPayload(userID, chatID int64, kbType uint8) *Payload {
User: userID,
Chat: chatID,
Data: marshal(KBChooserData{
Type: int64(kbType),
Type: kbType,
}),
}
}

View File

@ -1,2 +1,3 @@
// Package wizard contains setup wizard implementation.
// Deprecated: will be replaced with fsmwizard in the future.
package wizard

186
pkg/fsm/machine.go Normal file
View File

@ -0,0 +1,186 @@
package fsm
import (
"context"
"errors"
"fmt"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/types"
"runtime"
"time"
)
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]
// Enqueue the state input. Handle func will accept the current payload and modify it based on user input.
Enqueue(MachineStateProvider[T])
// Reset the machine to its initial state.
Reset()
}
// Machine is a finite-state machine implementation.
type Machine[T any] struct {
ctx context.Context
cancel context.CancelFunc
payload *T
transitions types.Queue[MachineStateProvider[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)
}
ctx, cancel := context.WithCancel(context.Background())
m := &Machine[T]{
ctx: ctx,
cancel: cancel,
state: initialState,
transitions: types.NewQueue[MachineStateProvider[T]](),
stateRouter: router,
initialState: initialState,
states: stateMap,
errHandler: errHandler,
}
go m.handleQueue()
runtime.SetFinalizer(m, func(m *Machine[T]) {
m.cancel()
})
return m
}
// 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]) MoveForHandling(id StateID, payload *T) error {
m.handleNow = true
return m.Move(id, payload)
}
func (m *Machine[T]) Enqueue(provider MachineStateProvider[T]) {
m.transitions.Enqueue(provider)
}
func (m *Machine[T]) handleQueue() {
defer func() {
if r := recover(); r != nil {
go m.handleQueue()
}
}()
for {
select {
case <-m.ctx.Done():
return
case <-time.After(time.Nanosecond * 10):
}
transition := m.transitions.Dequeue()
if transition == nil {
continue
}
m.handle(transition)
}
}
// Handle the input.
func (m *Machine[T]) handle(provider MachineStateProvider[T]) {
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
}
for {
st.Handle(m.payload, m)
if m.handleNow { // MoveForHandling was called, trying to handle again.
m.handleNow = false
continue
}
return
}
}
// 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
View 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
// MoveForHandling is the same as Move but it also triggers Handle immediately for the next state.
MoveForHandling(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
View 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
View 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
View 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
View 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
}