diff --git a/core/dispatcher.go b/core/dispatcher.go index 40ac155..756a7cf 100644 --- a/core/dispatcher.go +++ b/core/dispatcher.go @@ -10,9 +10,17 @@ import ( "github.com/max-messenger/max-bot-api-client-go/v2/model" "github.com/rs/zerolog/log" "sync" + "time" ) -type EventIsolation interface{} +const ( + defaultTTL = 5 * time.Second +) + +type EventIsolation interface { + Lock(ctx context.Context, userID, chatID int) error + Unlock(ctx context.Context, userID, chatID int) error +} type FSMStorage interface { Get(ctx context.Context, userID int) (state.Snapshot, error) @@ -91,15 +99,28 @@ func (d *Dispatcher) Start(ctx context.Context) error { } func (d *Dispatcher) handleUpdates(ch <-chan model.Update) { - for update := range ch { + for upd := range ch { d.wg.Go(func() { - d.handleUpdate(update) + ctx := context.Background() + + _ = d.eventIsolation.Lock(ctx, int(upd.UserID), int(upd.ChatID)) // todo: handle error somehow (retry or smth) + defer func() { + err := d.eventIsolation.Unlock(ctx, int(upd.UserID), int(upd.ChatID)) + if err != nil { + log.Warn(). + Err(err). + Int64("user_id", upd.UserID). + Int64("chat_id", upd.ChatID). + Msg("failed to unlock on update") + } + }() + + d.handleUpdate(ctx, upd) }) } } -func (d *Dispatcher) handleUpdate(upd model.Update) { - ctx := context.Background() +func (d *Dispatcher) handleUpdate(ctx context.Context, upd model.Update) { var st *state.Context snapshot, err := d.fsmStorage.Get(ctx, int(upd.UserID)) @@ -156,8 +177,8 @@ func WithEventIsolation(eventIsolation EventIsolation) DispatcherOption { func newDispatcherOptions(api *maxbot.Api) *dispatcherOptions { return &dispatcherOptions{ poller: NewLongPoller(api), - fsmStorage: memory.NewMemoryFSMStorage(), - eventIsolation: memory.NewMemoryEventIsolation(), + fsmStorage: memory.NewFSMStorage(), + eventIsolation: memory.NewEventIsolation(defaultTTL), } } diff --git a/state/memory/memory.go b/state/memory/memory.go index 15e077b..caeccf6 100644 --- a/state/memory/memory.go +++ b/state/memory/memory.go @@ -4,15 +4,30 @@ import ( "context" "github.com/abelith/gomax/state" "sync" + "time" ) type FSMStorage struct { m sync.Map } -type EventIsolation struct{} +type EventIsolation struct { + mu sync.Mutex + locks map[eventIsolationKey]*lockEntry + ttl time.Duration +} -func NewMemoryFSMStorage() *FSMStorage { +type eventIsolationKey struct { + userID int + chatID int +} + +type lockEntry struct { + timer *time.Timer + done chan struct{} +} + +func NewFSMStorage() *FSMStorage { return &FSMStorage{} } @@ -35,6 +50,61 @@ func (f *FSMStorage) Delete(_ context.Context, userID int) error { return nil } -func NewMemoryEventIsolation() *EventIsolation { - return &EventIsolation{} +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) }