feat: implemented update isolation

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) {
+func (d *Dispatcher) handleUpdate(ctx context.Context, upd model.Update) {
-	ctx := context.Background()
 	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(),
+		fsmStorage:     memory.NewFSMStorage(),
-		eventIsolation: memory.NewMemoryEventIsolation(),
+		eventIsolation: memory.NewEventIsolation(defaultTTL),
 	}
 }
 

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 {
+func NewEventIsolation(ttl time.Duration) *EventIsolation {
-	return &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)
 }