gomax/v1/core/dispatcher.go

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)
}