feat: added convenient state groups api
This commit is contained in:
abelith committed
1 parent
7bb08f60f3
commit
7a4e34528f
5 files changed
+63
-11
No files matched your search
@@ -0,0 +1,129 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/abelith/gomax/v1/fsm/state"
|
||||
"github.com/abelith/gomax/v1/isolation"
|
||||
"github.com/google/uuid"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type FSMStorage struct {
|
||||
m sync.Map
|
||||
}
|
||||
|
||||
type EventIsolation struct {
|
||||
mu sync.Mutex
|
||||
locks map[isolation.Scope]*lockEntry
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
type fsmKey struct {
|
||||
userID int64
|
||||
chatID int64
|
||||
}
|
||||
|
||||
type eventIsolationKey struct {
|
||||
userID int
|
||||
chatID int
|
||||
}
|
||||
|
||||
type lockEntry struct {
|
||||
owner uuid.UUID
|
||||
timer *time.Timer
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func NewFSMStorage() *FSMStorage {
|
||||
return &FSMStorage{}
|
||||
}
|
||||
|
||||
func (f *FSMStorage) Get(_ context.Context, scope isolation.Scope) (state.Snapshot, error) {
|
||||
st, ok := f.m.Load(fsmKey{
|
||||
userID: scope.UserID,
|
||||
chatID: scope.ChatID,
|
||||
})
|
||||
if !ok {
|
||||
return state.Snapshot{}, state.ErrStateNotFound
|
||||
}
|
||||
|
||||
return st.(state.Snapshot), nil
|
||||
}
|
||||
|
||||
func (f *FSMStorage) Put(_ context.Context, snapshot state.Snapshot) error {
|
||||
f.m.Store(fsmKey{
|
||||
userID: snapshot.Scope.UserID,
|
||||
chatID: snapshot.Scope.ChatID,
|
||||
}, snapshot)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *FSMStorage) Delete(_ context.Context, scope isolation.Scope) error {
|
||||
f.m.Delete(fsmKey{
|
||||
userID: scope.UserID,
|
||||
chatID: scope.ChatID,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func NewEventIsolation(ttl time.Duration) *EventIsolation {
|
||||
return &EventIsolation{
|
||||
locks: make(map[isolation.Scope]*lockEntry),
|
||||
ttl: ttl,
|
||||
}
|
||||
}
|
||||
|
||||
func (ei *EventIsolation) Lock(ctx context.Context, key isolation.Key) error {
|
||||
scope := key.Scope()
|
||||
|
||||
for {
|
||||
ei.mu.Lock()
|
||||
entry, busy := ei.locks[scope]
|
||||
if !busy {
|
||||
ei.acquireLocked(scope, key.UpdateID)
|
||||
ei.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
done := entry.done
|
||||
ei.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ei *EventIsolation) Unlock(_ context.Context, key isolation.Key) error {
|
||||
scope := key.Scope()
|
||||
|
||||
ei.mu.Lock()
|
||||
defer ei.mu.Unlock()
|
||||
|
||||
entry, ok := ei.locks[scope]
|
||||
if !ok || entry.owner != key.UpdateID {
|
||||
return isolation.ErrNotOwner
|
||||
}
|
||||
ei.releaseLocked(scope, entry)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ei *EventIsolation) acquireLocked(scope isolation.Scope, owner uuid.UUID) {
|
||||
entry := &lockEntry{owner: owner, done: make(chan struct{})}
|
||||
entry.timer = time.AfterFunc(ei.ttl, func() {
|
||||
ei.mu.Lock()
|
||||
defer ei.mu.Unlock()
|
||||
if cur, ok := ei.locks[scope]; ok && cur == entry {
|
||||
ei.releaseLocked(scope, entry)
|
||||
}
|
||||
})
|
||||
ei.locks[scope] = entry
|
||||
}
|
||||
|
||||
func (ei *EventIsolation) releaseLocked(scope isolation.Scope, entry *lockEntry) {
|
||||
entry.timer.Stop()
|
||||
delete(ei.locks, scope)
|
||||
close(entry.done)
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
package state
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"github.com/abelith/gomax/v1/isolation"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrStateNotFound = fmt.Errorf("state not found")
|
||||
ErrKeyNotFound = fmt.Errorf("state: key not found")
|
||||
)
|
||||
|
||||
type State struct {
|
||||
group string
|
||||
name string
|
||||
}
|
||||
|
||||
type Group struct {
|
||||
name string
|
||||
states map[string]*State
|
||||
}
|
||||
|
||||
type Codec interface {
|
||||
Marshal(v any) ([]byte, error)
|
||||
Unmarshal(data []byte, v any) error
|
||||
}
|
||||
|
||||
type JSONCodec struct{}
|
||||
|
||||
func (JSONCodec) Marshal(v any) ([]byte, error) {
|
||||
return json.Marshal(v)
|
||||
}
|
||||
|
||||
func (JSONCodec) Unmarshal(d []byte, v any) error {
|
||||
return json.Unmarshal(d, v)
|
||||
}
|
||||
|
||||
type Snapshot struct {
|
||||
Scope isolation.Scope `json:"id"`
|
||||
State string `json:"state"`
|
||||
Data map[string][]byte `json:"data"`
|
||||
}
|
||||
|
||||
type Context struct {
|
||||
scope isolation.Scope
|
||||
state State
|
||||
data map[string][]byte
|
||||
codec Codec
|
||||
clear bool
|
||||
}
|
||||
|
||||
type Option func(*Context)
|
||||
|
||||
func NewStateGroup(name string) *Group {
|
||||
return &Group{
|
||||
name: name,
|
||||
states: make(map[string]*State),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Group) State(name string) State {
|
||||
state := g.states[name]
|
||||
|
||||
if state != nil {
|
||||
return *state
|
||||
}
|
||||
|
||||
if state = g.states[name]; state != nil {
|
||||
return *state
|
||||
}
|
||||
|
||||
state = &State{
|
||||
group: g.name,
|
||||
name: name,
|
||||
}
|
||||
|
||||
g.states[name] = state
|
||||
return *state
|
||||
}
|
||||
|
||||
func (s *State) String() string {
|
||||
return fmt.Sprintf("%s:%s", s.group, s.name)
|
||||
}
|
||||
|
||||
func fromString(s string) State {
|
||||
pieces := strings.Split(s, ":")
|
||||
if len(pieces) < 2 {
|
||||
return State{}
|
||||
}
|
||||
|
||||
return State{
|
||||
group: pieces[0],
|
||||
name: pieces[1],
|
||||
}
|
||||
}
|
||||
|
||||
func WithCodec(c Codec) Option { return func(sc *Context) { sc.codec = c } }
|
||||
|
||||
func New(scope isolation.Scope, opts ...Option) *Context {
|
||||
sc := &Context{
|
||||
scope: scope,
|
||||
data: make(map[string][]byte),
|
||||
codec: JSONCodec{},
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(sc)
|
||||
}
|
||||
return sc
|
||||
}
|
||||
|
||||
func Restore(s Snapshot, opts ...Option) *Context {
|
||||
sc := New(s.Scope, opts...)
|
||||
sc.state = fromString(s.State)
|
||||
for k, v := range s.Data {
|
||||
sc.data[k] = v
|
||||
}
|
||||
return sc
|
||||
}
|
||||
|
||||
func (sc *Context) Snapshot() Snapshot {
|
||||
data := make(map[string][]byte, len(sc.data))
|
||||
for k, v := range sc.data {
|
||||
data[k] = v
|
||||
}
|
||||
return Snapshot{Scope: sc.scope, State: sc.state.String(), Data: data}
|
||||
}
|
||||
|
||||
func (sc *Context) Set(state State) {
|
||||
sc.state = state
|
||||
}
|
||||
|
||||
func (sc *Context) State() State { return sc.state }
|
||||
|
||||
func (sc *Context) Clear() {
|
||||
sc.state = State{}
|
||||
sc.clear = true
|
||||
sc.data = make(map[string][]byte)
|
||||
}
|
||||
|
||||
func (sc *Context) IsEmpty() bool {
|
||||
return sc.state == State{} && sc.clear
|
||||
}
|
||||
|
||||
func (sc *Context) Store(key string, val any) error {
|
||||
b, err := sc.codec.Marshal(val)
|
||||
if err != nil {
|
||||
return fmt.Errorf("sc.Store: failed to encode %q: %w", key, err)
|
||||
}
|
||||
sc.data[key] = b
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sc *Context) Delete(key string) {
|
||||
delete(sc.data, key)
|
||||
}
|
||||
|
||||
func (sc *Context) Load(key string) Value {
|
||||
b, ok := sc.data[key]
|
||||
return Value{key: key, raw: b, ok: ok, codec: sc.codec}
|
||||
}
|
||||
|
||||
type Value struct {
|
||||
key string
|
||||
raw []byte
|
||||
ok bool
|
||||
codec Codec
|
||||
}
|
||||
|
||||
func (v Value) Exists() bool { return v.ok }
|
||||
|
||||
func (v Value) Raw() []byte { return v.raw }
|
||||
|
||||
func (v Value) Decode(dst any) error {
|
||||
if !v.ok {
|
||||
return fmt.Errorf("%w: %q", ErrKeyNotFound, v.key)
|
||||
}
|
||||
if err := v.codec.Unmarshal(v.raw, dst); err != nil {
|
||||
return fmt.Errorf("state: decode %q as %T: %w", v.key, dst, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (v Value) Str() (string, error) {
|
||||
var s string
|
||||
if err := v.Decode(&s); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
func (v Value) Int() (int, error) {
|
||||
var i int
|
||||
if err := v.Decode(&i); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return i, nil
|
||||
}
|
||||
|
||||
func (v Value) Bool() (bool, error) {
|
||||
var b bool
|
||||
if err := v.Decode(&b); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
func (v Value) Float() (float64, error) {
|
||||
var f float64
|
||||
if err := v.Decode(&f); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
func (v Value) Bytes() ([]byte, error) {
|
||||
var b []byte
|
||||
if err := v.Decode(&b); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func (v Value) StrOr(def string) string {
|
||||
if s, err := v.Str(); err == nil {
|
||||
return s
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
func (v Value) IntOr(def int) int {
|
||||
if i, err := v.Int(); err == nil {
|
||||
return i
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
func (v Value) BoolOr(def bool) bool {
|
||||
if b, err := v.Bool(); err == nil {
|
||||
return b
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
func (v Value) FloatOr(def float64) float64 {
|
||||
if f, err := v.Float(); err == nil {
|
||||
return f
|
||||
}
|
||||
return def
|
||||
}
|
||||
Reference in new issue
Block a user