130 lines
2.4 KiB
Go
130 lines
2.4 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"github.com/abelith/gomax/v1/isolation"
|
|
"github.com/abelith/gomax/v1/state"
|
|
"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)
|
|
}
|