111 lines
2.1 KiB
Go
111 lines
2.1 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"github.com/abelith/gomax/v1/state"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type FSMStorage struct {
|
|
m sync.Map
|
|
}
|
|
|
|
type EventIsolation struct {
|
|
mu sync.Mutex
|
|
locks map[eventIsolationKey]*lockEntry
|
|
ttl time.Duration
|
|
}
|
|
|
|
type eventIsolationKey struct {
|
|
userID int
|
|
chatID int
|
|
}
|
|
|
|
type lockEntry struct {
|
|
timer *time.Timer
|
|
done chan struct{}
|
|
}
|
|
|
|
func NewFSMStorage() *FSMStorage {
|
|
return &FSMStorage{}
|
|
}
|
|
|
|
func (f *FSMStorage) Get(_ context.Context, userID int) (state.Snapshot, error) {
|
|
st, ok := f.m.Load(userID)
|
|
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(snapshot.ID, snapshot)
|
|
return nil
|
|
}
|
|
|
|
func (f *FSMStorage) Delete(_ context.Context, userID int) error {
|
|
f.m.Delete(userID)
|
|
return nil
|
|
}
|
|
|
|
func NewEventIsolation(ttl time.Duration) *EventIsolation {
|
|
return &EventIsolation{
|
|
locks: make(map[eventIsolationKey]*lockEntry),
|
|
ttl: ttl,
|
|
}
|
|
}
|
|
|
|
func (ei *EventIsolation) Lock(ctx context.Context, userID, chatID int) error {
|
|
key := eventIsolationKey{userID: userID, chatID: chatID}
|
|
|
|
for {
|
|
ei.mu.Lock()
|
|
entry, busy := ei.locks[key]
|
|
if !busy {
|
|
ei.acquireLocked(key)
|
|
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, userID, chatID int) error {
|
|
key := eventIsolationKey{userID: userID, chatID: chatID}
|
|
|
|
ei.mu.Lock()
|
|
defer ei.mu.Unlock()
|
|
|
|
if entry, ok := ei.locks[key]; ok {
|
|
ei.releaseLocked(key, entry)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ei *EventIsolation) acquireLocked(key eventIsolationKey) {
|
|
entry := &lockEntry{done: make(chan struct{})}
|
|
entry.timer = time.AfterFunc(ei.ttl, func() {
|
|
ei.mu.Lock()
|
|
defer ei.mu.Unlock()
|
|
if cur, ok := ei.locks[key]; ok && cur == entry {
|
|
ei.releaseLocked(key, entry)
|
|
}
|
|
})
|
|
ei.locks[key] = entry
|
|
}
|
|
|
|
func (ei *EventIsolation) releaseLocked(key eventIsolationKey, entry *lockEntry) {
|
|
entry.timer.Stop()
|
|
delete(ei.locks, key)
|
|
close(entry.done)
|
|
}
|