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 backoff struct {
	base time.Duration
	max  time.Duration

	window []bool
	pos    int
	count  int
	errors int

	sinceAdjust int
	interval    time.Duration
}

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 newBackoff(
	base, max time.Duration,
	windowSize int,
) *backoff {
	if base <= 0 || max < base || windowSize <= 0 {
		panic("backoff: invalid configuration")
	}

	return &backoff{
		base:     base,
		max:      max,
		window:   make([]bool, windowSize),
		interval: base,
	}
}

func (b *backoff) observe(err error) time.Duration {
	failed := err != nil

	if b.count == len(b.window) && b.window[b.pos] {
		b.errors--
	} else {
		b.count++
	}

	b.window[b.pos] = failed
	if failed {
		b.errors++
	}

	b.pos = (b.pos + 1) % len(b.window)
	b.sinceAdjust++

	if b.count < len(b.window) ||
		b.sinceAdjust < len(b.window) {
		return b.interval
	}

	b.sinceAdjust = 0
	errorRate := float64(b.errors) / float64(b.count)

	switch {
	case errorRate >= 0.3:
		b.interval = min(b.interval*2, b.max)
	case errorRate <= 0.1:
		b.interval = max(b.interval/2, b.base)
	}

	return b.interval
}

func (lp *LongPoller) poll(ctx context.Context, ch chan model.Update) {
	b := newBackoff(
		lp.pollInterval,
		30*time.Second,
		20,
	)

	var marker int64

	for {
		if err := wait(ctx, b.interval); err != nil {
			return
		}

		start := time.Now()
		updates, nextMarker, err := lp.api.Subscriptions.GetUpdates(ctx, marker)
		latency := time.Since(start)

		if err != nil {
			if errors.Is(err, context.Canceled) ||
				ctx.Err() != nil {
				return
			}

			delay := b.observe(err)

			lp.logger.Error().
				Err(err).
				Dur("latency", latency).
				Dur("next_interval", delay).
				Msg("failed to receive updates")

			continue
		}

		b.observe(nil)
		marker = nextMarker

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

func wait(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()

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