package core import ( "context" "errors" mcontext "github.com/abelith/gomax/v1/context" "github.com/abelith/gomax/v1/isolation" "github.com/abelith/gomax/v1/state" "github.com/abelith/gomax/v1/state/memory" "github.com/google/uuid" maxbot "github.com/max-messenger/max-bot-api-client-go/v2" "github.com/max-messenger/max-bot-api-client-go/v2/model" "github.com/rs/zerolog/log" "sync" "time" ) const ( defaultTTL = 30 * time.Second defaultLockTimeout = 10 * time.Second defaultHandlerTimeout = 10 * time.Second ) type EventIsolation interface { Lock(ctx context.Context, key isolation.Key) error Unlock(ctx context.Context, key isolation.Key) error } type FSMStorage interface { Get(ctx context.Context, userID int) (state.Snapshot, error) Put(ctx context.Context, state state.Snapshot) error Delete(ctx context.Context, userID int) error } type UpdateHandler interface { HandleUpdate(c *mcontext.Context) error Use(mw ...Middleware) } type Poller interface { Start(ctx context.Context) (<-chan model.Update, error) } type DispatcherOption interface { apply(d *dispatcherOptions) } type Dispatcher struct { api *maxbot.Api poller Poller fsmStorage FSMStorage eventIsolation EventIsolation rt *Router wg sync.WaitGroup } type dispatcherOptions struct { poller Poller fsmStorage FSMStorage eventIsolation EventIsolation } type dispatcherOptionWrapper struct { fn func(options *dispatcherOptions) } func NewDispatcher(api *maxbot.Api, options ...DispatcherOption) *Dispatcher { d := &Dispatcher{ api: api, rt: &Router{}, } opts := newDispatcherOptions(api) for _, opt := range options { opt.apply(opts) } d.poller = opts.poller d.fsmStorage = opts.fsmStorage d.eventIsolation = opts.eventIsolation return d } func (d *Dispatcher) Handle(handler UpdateHandler, filters ...func(c *mcontext.Context) bool) { d.rt.Handle(handler, filters...) } func (d *Dispatcher) HandleFunc(handler func(c *mcontext.Context) error, filters ...func(c *mcontext.Context) bool) { d.rt.HandleFunc(handler, filters...) } func (d *Dispatcher) HandleUpdate(c *mcontext.Context) error { return d.rt.HandleUpdate(c) } func (d *Dispatcher) Use(mw ...Middleware) { d.rt.Use(mw...) } func (d *Dispatcher) Start(ctx context.Context) error { ch, err := d.poller.Start(ctx) if err != nil { return err } d.handleUpdates(ch) d.wg.Wait() return nil } func (d *Dispatcher) handleUpdates(ch <-chan model.Update) { for upd := range ch { d.wg.Go(func() { key := isolation.Key{ UpdateID: uuid.New(), UserID: upd.UserID, ChatID: upd.ChatID, } l := log.With(). Str("update_id", key.UpdateID.String()). Int64("user_id", key.UserID). Int64("chat_id", key.ChatID). Logger() lockCtx, cancel := context.WithTimeout(context.Background(), defaultLockTimeout) defer cancel() if err := d.eventIsolation.Lock(lockCtx, key); err != nil { log.Error(). Err(err). Int64("user_id", key.UserID). Int64("chat_id", key.ChatID). Msg("failed to lock on update") return } defer func() { unlockCtx, unlockCancel := context.WithTimeout(context.Background(), defaultLockTimeout) defer unlockCancel() err := d.eventIsolation.Unlock(unlockCtx, key) switch { case errors.Is(err, isolation.ErrNotOwner): l.Warn().Msg("lock expired before handler finished") case err != nil: l.Warn().Err(err).Msg("failed to unlock on update") } }() handleCtx, cancel := context.WithTimeout(context.Background(), defaultHandlerTimeout) defer cancel() d.handleUpdate(handleCtx, upd) }) } } func (d *Dispatcher) handleUpdate(ctx context.Context, upd model.Update) { var st *state.Context snapshot, err := d.fsmStorage.Get(ctx, int(upd.UserID)) if err != nil && !errors.Is(err, state.ErrStateNotFound) { log.Error().Err(err).Send() return } if errors.Is(err, state.ErrStateNotFound) { st = state.New(int(upd.UserID)) } else { st = state.Restore(snapshot) } c := mcontext.New(ctx, upd, st, d.api) if err := d.rt.HandleUpdate(c); err != nil { log.Error().Err(err).Send() return } st = c.State() if st.State() == state.ClearState { if err := d.fsmStorage.Delete(ctx, int(upd.UserID)); err != nil { log.Error().Err(err).Send() return } } if st.State() != "" { if err := d.fsmStorage.Put(ctx, st.Snapshot()); err != nil { log.Error().Err(err).Send() return } } } func WithPoller(poller Poller) DispatcherOption { return newOption(func(opts *dispatcherOptions) { opts.poller = poller }) } func WithFSMStorage(storage FSMStorage) DispatcherOption { return newOption(func(opts *dispatcherOptions) { opts.fsmStorage = storage }) } func WithEventIsolation(eventIsolation EventIsolation) DispatcherOption { return newOption(func(opts *dispatcherOptions) { opts.eventIsolation = eventIsolation }) } func newDispatcherOptions(api *maxbot.Api) *dispatcherOptions { return &dispatcherOptions{ poller: NewLongPoller(api), fsmStorage: memory.NewFSMStorage(), eventIsolation: memory.NewEventIsolation(defaultTTL), } } func newOption(fn func(options *dispatcherOptions)) DispatcherOption { return dispatcherOptionWrapper{fn: fn} } func (dow dispatcherOptionWrapper) apply(opts *dispatcherOptions) { dow.fn(opts) }