gomax/v1/core/pollers/pollers.go

package pollers

import (
	"context"
	"crypto/tls"
	"errors"
	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"
	"github.com/rs/zerolog/log"
	"net"
	"net/http"
	"sync"
	"time"
)

const (
	defaultInterval       = 200 * time.Millisecond
	defaultWhStartTimeout = 10 * time.Second
)

type WebhookOption interface {
	apply(opts *webhookOptions)
}

type LongPollerOption interface {
	apply(opts *longPollerOptions)
}

type Webhook struct {
	srv *http.Server
	ch  chan model.Update
	wg  *sync.WaitGroup
}

type LongPoller struct {
	api          *maxbot.Api
	pollInterval time.Duration
	logger       zerolog.Logger
	wg           sync.WaitGroup
}

type TestPoller struct {
	ch chan model.Update
	mu sync.Mutex
}

type webhookOptions struct {
	addr              string
	tlsConfig         *tls.Config
	readTimeout       time.Duration
	readHeaderTimeout time.Duration
	writeTimeout      time.Duration
	idleTimeout       time.Duration
	maxHeaderBytes    int
	baseContext       func(l net.Listener) context.Context
	connContext       func(ctx context.Context, conn net.Conn) context.Context
}

type longPollerOptions struct {
	pollInterval time.Duration
	logger       zerolog.Logger
}

type longPollerOptionWrapper struct {
	fn func(opts *longPollerOptions)
}

type webhookOptionWrapper struct {
	fn func(opts *webhookOptions)
}

func NewWebhookPoller(api *maxbot.Api, secret string, options ...WebhookOption) *Webhook {
	ch := make(chan model.Update, 1)

	handler := api.GetHandler(func(ctx context.Context, update model.Update) {
		select {
		case ch <- update:
		case <-ctx.Done():
		}
	}, secret)

	srv := &http.Server{
		Handler: handler,
	}

	wh := &Webhook{
		srv: srv,
		ch:  ch,
	}

	opts := newWebhookOptions()
	for _, opt := range options {
		opt.apply(opts)
	}

	wh.srv.Addr = opts.addr
	wh.srv.TLSConfig = opts.tlsConfig
	wh.srv.ReadTimeout = opts.readTimeout
	wh.srv.ReadHeaderTimeout = opts.readHeaderTimeout
	wh.srv.WriteTimeout = opts.writeTimeout
	wh.srv.IdleTimeout = opts.idleTimeout
	wh.srv.MaxHeaderBytes = opts.maxHeaderBytes
	wh.srv.BaseContext = opts.baseContext
	wh.srv.ConnContext = opts.connContext

	return wh
}

func (w *Webhook) Start(ctx context.Context) (<-chan model.Update, error) {
	l, err := net.Listen("tcp", w.srv.Addr)
	if err != nil {
		return nil, err
	}

	go func() {
		var err error

		if w.srv.TLSConfig == nil {
			err = w.srv.Serve(l)
		} else {
			err = w.srv.ServeTLS(l, "", "")
		}

		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Ctx(ctx).Error().Err(err).Msg("failed to start webhook server")
		}
	}()

	return w.ch, nil
}

func (w *Webhook) Stop(ctx context.Context) error {
	defer close(w.ch)

	if err := w.srv.Shutdown(ctx); err != nil {
		return err
	}
	return nil
}

func NewLongPoller(api *maxbot.Api, options ...LongPollerOption) *LongPoller {
	lp := &LongPoller{
		api: api,
	}

	opts := newLongPollerOptions()
	for _, opt := range options {
		opt.apply(opts)
	}

	lp.pollInterval = opts.pollInterval
	lp.logger = opts.logger

	return lp
}

func (lp *LongPoller) Start(ctx context.Context) (<-chan model.Update, error) {
	ch := make(chan model.Update)

	lp.wg.Go(func() {
		defer close(ch)
		lp.poll(ctx, ch)
	})

	return ch, nil
}

func (lp *LongPoller) Stop(ctx context.Context) error {
	done := make(chan struct{})

	go func() {
		lp.wg.Wait()
		done <- struct{}{}
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-done:
		return nil
	}
}

func (tp *TestPoller) Inject(ctx context.Context, updates []model.Update) {
	ch := make(chan model.Update, 1)

	go func() {
		defer close(ch)

		for _, upd := range updates {
			select {
			case <-ctx.Done():
			case ch <- upd:
			}
		}
	}()

	tp.mu.Lock()
	tp.ch = ch
	defer tp.mu.Unlock()
}

func (tp *TestPoller) Start(_ context.Context) (<-chan model.Update, error) {
	return tp.ch, nil
}

func (tp *TestPoller) Stop(_ context.Context) error {
	return nil
}

func WithPollInterval(interval time.Duration) LongPollerOption {
	return newLpOption(func(opts *longPollerOptions) {
		opts.pollInterval = interval
	})
}

func WithZerolog(logger zerolog.Logger) LongPollerOption {
	return newLpOption(func(opts *longPollerOptions) {
		opts.logger = logger
	})
}

func WithAddress(addr string) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.addr = addr
	})
}

func WithTLSConfig(cfg *tls.Config) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.tlsConfig = cfg
	})
}

func WithReadTimeout(timeout time.Duration) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.readTimeout = timeout
	})
}

func WithReadHeaderTimeout(timeout time.Duration) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.readHeaderTimeout = timeout
	})
}

func WithWriteTimeout(timeout time.Duration) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.writeTimeout = timeout
	})
}

func WithIdleTimeout(timeout time.Duration) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.idleTimeout = timeout
	})
}

func WithMaxHeaderBytes(bytes int) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.maxHeaderBytes = bytes
	})
}

func WithBaseContext(fn func(l net.Listener) context.Context) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.baseContext = fn
	})
}

func WithConnContext(fn func(ctx context.Context, c net.Conn) context.Context) WebhookOption {
	return newWhOption(func(opts *webhookOptions) {
		opts.connContext = fn
	})
}

func (lp *LongPoller) poll(ctx context.Context, ch chan model.Update) {
	var updates []model.Update
	var marker int64
	var err error

	// todo: backoff on errors

	for {
		select {
		case <-ctx.Done():
			return
		case <-time.Tick(lp.pollInterval):
			start := time.Now()
			updates, marker, err = lp.api.Subscriptions.GetUpdates(ctx, marker)
			latency := time.Since(start)
			if err != nil && !errors.Is(err, context.Canceled) {
				lp.logger.Error().Err(err).Send()
			}
			lp.logger.Debug().
				Int("received", len(updates)).
				Dur("latency", latency).
				Msg("received updates")

			for _, update := range updates {
				select {
				case ch <- update:
				}
			}
		}
	}
}

func newLongPollerOptions() *longPollerOptions {
	return &longPollerOptions{
		pollInterval: defaultInterval,
		logger:       zerolog.Nop(),
	}
}

func newWebhookOptions() *webhookOptions {
	return &webhookOptions{
		addr: ":8080",
	}
}

func newLpOption(fn func(opts *longPollerOptions)) LongPollerOption {
	return longPollerOptionWrapper{fn: fn}
}

func newWhOption(fn func(options *webhookOptions)) WebhookOption {
	return webhookOptionWrapper{fn: fn}
}

func (low longPollerOptionWrapper) apply(opts *longPollerOptions) {
	low.fn(opts)
}

func (wow webhookOptionWrapper) apply(opts *webhookOptions) {
	wow.fn(opts)
}