package core import ( "context" "errors" mcontext "github.com/abelith/gomax/v1/context" "github.com/abelith/gomax/v1/core/pollers" "github.com/abelith/gomax/v1/fsm/memory" "github.com/abelith/gomax/v1/fsm/state" "github.com/abelith/gomax/v1/isolation" "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" "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, scope isolation.Scope) (state.Snapshot, error) Put(ctx context.Context, state state.Snapshot) error Delete(ctx context.Context, scope isolation.Scope) error } type UpdateHandler interface { HandleUpdate(c *mcontext.Context) error Use(mw ...Middleware) } type Poller interface { Start(ctx context.Context) (<-chan model.Update, error) Stop(ctx context.Context) error } type DispatcherOption interface { apply(d *dispatcherOptions) } type Dispatcher struct { api *maxbot.Api poller Poller fsmStorage FSMStorage eventIsolation EventIsolation onStart func(ctx context.Context) error onStop func(ctx context.Context) error rt *Router logger zerolog.Logger wg sync.WaitGroup } type dispatcherOptions struct { poller Poller fsmStorage FSMStorage eventIsolation EventIsolation logger zerolog.Logger } 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 d.logger = opts.logger 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 } if d.onStart != nil { if err := d.onStart(ctx); err != nil { return err } } d.handleUpdates(ch) return nil } func (d *Dispatcher) Stop(ctx context.Context) error { if err := d.poller.Stop(ctx); err != nil { return err } done := make(chan struct{}) go func() { d.wg.Wait() done <- struct{}{} }() select { case <-ctx.Done(): return ctx.Err() case <-done: } if d.onStop != nil { return d.onStop(ctx) } return nil } func (d *Dispatcher) OnStart(f func(ctx context.Context) error) { d.onStart = f } func (d *Dispatcher) OnStop(f func(ctx context.Context) error) { d.onStop = f } 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 := d.logger.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 { d.logger.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, isolation.Scope{UserID: upd.UserID, ChatID: upd.ChatID}) if err != nil && !errors.Is(err, state.ErrStateNotFound) { d.logger.Error().Err(err).Send() return } if errors.Is(err, state.ErrStateNotFound) { st = state.New(isolation.Scope{UserID: upd.UserID, ChatID: upd.ChatID}) } else { st = state.Restore(snapshot) } c := mcontext.New(ctx, upd, st, d.api) if err := d.rt.HandleUpdate(c); err != nil { d.logger.Error().Err(err).Send() return } st = c.State() if st.IsEmpty() { if err := d.fsmStorage.Delete(ctx, isolation.Scope{UserID: upd.UserID, ChatID: upd.ChatID}); err != nil { d.logger.Error().Err(err).Send() } return } if err := d.fsmStorage.Put(ctx, st.Snapshot()); err != nil { d.logger.Error().Err(err).Send() } } 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 WithZerolog(logger zerolog.Logger) DispatcherOption { return newOption(func(options *dispatcherOptions) { options.logger = logger }) } func newDispatcherOptions(api *maxbot.Api) *dispatcherOptions { return &dispatcherOptions{ poller: pollers.NewLongPoller(api), fsmStorage: memory.NewFSMStorage(), eventIsolation: memory.NewEventIsolation(defaultTTL), logger: zerolog.Nop(), } } func newOption(fn func(options *dispatcherOptions)) DispatcherOption { return dispatcherOptionWrapper{fn: fn} } func (dow dispatcherOptionWrapper) apply(opts *dispatcherOptions) { dow.fn(opts) }