gomax/v1/core/dispatcher.go

package core

import (
	"context"
	"errors"
	mcontext "github.com/abelith/gomax/v1/context"
	"github.com/abelith/gomax/v1/state"
	"github.com/abelith/gomax/v1/state/memory"
	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 = 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)
	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) 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)
	<-ctx.Done()
	d.wg.Wait()
	return nil
}

func (d *Dispatcher) handleUpdates(ch <-chan model.Update) {
	for upd := range ch {
		d.wg.Go(func() {
			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(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)
}