Compare commits

...

2 Commits

Author SHA1 Message Date
b13cd05b45 wip: experimental branch with finite-state machine processor (doesn't work yet) 2024-05-14 15:29:35 +03:00
2cd46585a8 add hot-reload & debugger
All checks were successful
continuous-integration/drone/push Build is passing
2024-05-14 15:26:28 +03:00
32 changed files with 1293 additions and 52 deletions

10
.air.toml Normal file
View File

@ -0,0 +1,10 @@
root = "."
[build]
exclude_regex = ["_test\\.go"]
bin = "/usr/bin/make debug"
cmd = "/usr/bin/make build"
kill_delay = 0
log = "build-errors.log"
send_interrupt = true
stop_on_error = true

View File

@ -19,3 +19,5 @@ steps:
trigger:
event:
- push
branch:
- master

2
.gitignore vendored
View File

@ -22,6 +22,8 @@ vendor/
# Go workspace file
go.work
.idea/
.config/
tmp/
# Env files
.env

View File

@ -2,20 +2,37 @@ SHELL = /bin/bash -o pipefail
export PATH := $(shell go env GOPATH)/bin:$(PATH)
ROOT_DIR=$(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
BIN=$(ROOT_DIR)/build/vegapokerbot
BIN_NAME=vegapokerbot
BIN=$(ROOT_DIR)/build/$(BIN_NAME)
GO_VERSION=$(shell go version | sed -e 's/go version //')
BIN_DIR=$(ROOT_DIR)/build
build: deps fmt
@echo "> building with ${GO_VERSION}"
ifeq ($(DEBUG), true)
@echo "> building debug with ${GO_VERSION}"
@CGO_ENABLED=0 go build -buildvcs=false -gcflags "all=-N -l" -tags=release -o $(BIN) .
else
@echo "> building release with ${GO_VERSION}"
@CGO_ENABLED=0 go build -buildvcs=false -tags=release -o $(BIN) .
endif
@echo $(BIN)
docker:
@docker buildx build --platform linux/amd64 --tag $CI_APP_IMAGE --file Dockerfile .
run:
ifeq ($(DEBUG), true)
@killall -s 9 dlv > /dev/null 2>&1 || true
@killall -s 9 ${BIN_NAME} > /dev/null 2>&1 || true
@air
else
@${BIN} run
endif
debug:
@echo "> starting under delve"
@dlv --listen=:40000 --headless=true --api-version=2 --accept-multiclient --log exec ${BIN} run
fmt:
@echo "> fmt"
@gofmt -l -s -w `go list -buildvcs=false -f '{{.Dir}}' ${ROOT_DIR}/... | grep -v /vendor/`

View File

@ -1,29 +1,23 @@
services:
db:
image: postgres:latest
restart: always
environment:
POSTGRES_USER: app
POSTGRES_PASSWORD: app
POSTGRES_DATABASE: app
networks:
- default
app:
image: "gitea.neur0tx.site/neur0toxine/vegapokerbot:latest"
networks:
- default
links:
- db
env_file:
- .env
restart: always
labels:
traefik.enable: "true"
traefik.http.routers.vegapokerbot.rule: Host(`vegapokerbot.neur0tx.site`) && PathPrefix(`/webhook8c31c4b2d65a87b4f3a6f10f8eb166fba9a2e5dc7696bc5291a7e69641dc5c21`)
traefik.http.routers.vegapokerbot.rule: Host(`vegapokerbot.example.com`) && PathPrefix(`/webhook`)
traefik.http.routers.vegapokerbot.entrypoints: websecure
traefik.http.services.vegapokerbot.loadbalancer.server.port: "3333"
traefik.http.routers.vegapokerbot.tls: true
traefik.http.routers.vegapokerbot.tls.certresolver: letsencrypt
networks:
default:
driver: bridge

View File

@ -22,8 +22,10 @@ services:
- db
environment:
GOCACHE: /go
DEBUG: true # Set to false to disable hot-reload & debugger.
ports:
- 3333:3333
- 40000:40000
command: make build run
networks:

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,13 +4,16 @@ 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"
"gitea.neur0tx.site/Neur0toxine/vegapokerbot/pkg/types"
"github.com/mymmrac/telego"
"go.uber.org/zap"
"gorm.io/gorm"
"net/url"
"time"
)
type App struct {
@ -20,6 +23,7 @@ type App struct {
Config *config.Config
Repositories *db.Repositories
Handlers map[handler.Type]handler.Handler
updates types.Queue[*telego.Update]
db *gorm.DB
}
@ -113,6 +117,8 @@ func (a *App) initTelegram() error {
}
func (a *App) initHandlers() {
a.updates = types.NewQueue[*telego.Update]()
fsmwizard.PopulateStates(a)
a.Handlers = map[handler.Type]handler.Handler{
handler.Noop: handler.NewNoopHandler(a.Logger, a.Config.Debug),
handler.Message: handler.NewMessageHandler(a),
@ -138,6 +144,17 @@ func (a *App) handler(update telego.Update) handler.Handler {
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) {
defer func() {
if r := recover(); r != nil {
@ -155,9 +172,11 @@ func (a *App) longPoll() error {
if err != nil {
return err
}
go a.processUpdates()
defer a.Telegram.StopLongPolling()
for update := range updates {
go a.processUpdate(update)
update := update
a.updates.Enqueue(&update)
}
return nil
}
@ -182,8 +201,10 @@ func (a *App) listenWebhook() error {
defer func() {
_ = a.Telegram.StopWebhook()
}()
go a.processUpdates()
for update := range updates {
go a.processUpdate(update)
update := update
a.updates.Enqueue(&update)
}
return nil
}

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,25 @@ 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)
return nil
return fsmwizard.Get(wh.MyChatMember.From.ID).Handle(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
if w == nil {
return &fsmwizard.Wizard{Data: wh}
}
w.Data = wh
return w
})
// 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) MoveForHandle(mc fsm.MachineControls[Wizard], stateID fsm.StateID, pl *Wizard) {
s.LogError(mc.MoveForHandle(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.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
}

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,26 +19,30 @@ 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)
}
return fsmwizard.Get(wh.Message.From.ID).Handle(func(w *fsmwizard.Wizard) *fsmwizard.Wizard {
if w == nil {
return &fsmwizard.Wizard{Data: wh}
}
w.Data = wh
return w
})
if util.MatchCommand("poll", wh.Message) {
return group.NewPoll(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)
}
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

@ -8,3 +8,10 @@ import (
func MatchCommand(command string, msg *telego.Message) bool {
return th.CommandEqual(command)(telego.Update{Message: msg})
}
func HasCommand(msg *telego.Message) bool {
if msg == nil {
return false
}
return th.CommandRegexp.MatchString(msg.Text)
}

View File

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

View File

@ -32,7 +32,7 @@ func init() {
}
for _, tag := range tags {
localizers[tag.String()] = &localizer{loc: i18n.NewLocalizer(bundle, tag.String())}
localizers[tag.String()] = &localizer{loc: i18n.NewLocalizer(bundle, tag.String()), tag: tag}
}
}

View File

@ -1,14 +1,19 @@
package locale
import "github.com/nicksnyder/go-i18n/v2/i18n"
import (
"github.com/nicksnyder/go-i18n/v2/i18n"
"golang.org/x/text/language"
)
type Localizer interface {
Message(string) string
Template(string, interface{}) string
Tag() language.Tag
}
type localizer struct {
loc *i18n.Localizer
tag language.Tag
}
func (l *localizer) Message(str string) string {
@ -21,3 +26,7 @@ func (l *localizer) Template(str string, tpl interface{}) string {
TemplateData: tpl,
})
}
func (l *localizer) Tag() language.Tag {
return l.tag
}

147
pkg/fsm/machine.go Normal file
View 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
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
// 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
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
}