gomax/v1/core/pollers.go

package core

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

type LongPoller struct {
	api          *maxbot.Api
	pollInterval time.Duration
}

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
}

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(_ context.Context) (<-chan model.Update, error) {
	errCh := make(chan error, 1)
	defer close(errCh)

	go func() {
		var err error
		defer func() {
			errCh <- err
			if err != nil && !errors.Is(err, http.ErrServerClosed) {
				log.Fatal().Err(err).Msg("webhook server error")
			}
		}()

		if w.srv.TLSConfig == nil {
			err = w.srv.ListenAndServe()
			return
		}

		err = w.srv.ListenAndServeTLS("", "")
	}()

	select {
	case err := <-errCh:
		return nil, err
	case <-time.After(defaultWhStartTimeout):
		return w.ch, nil
	}
}

func (w *Webhook) Stop(ctx context.Context) error {
	return w.srv.Shutdown(ctx)
}

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

	return lp
}

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

	ch := make(chan model.Update)

	go func() {
		defer close(ch)

		// 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) {
					log.Error().Err(err).Send()
				}
				log.Ctx(ctx).Debug().
					Int("received", len(updates)).
					Dur("latency", latency).
					Msg("received updates")

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

	return ch, 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 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 newLongPollerOptions() *longPollerOptions {
	return &longPollerOptions{
		pollInterval: defaultInterval,
	}
}

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

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

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