diff --git a/v1/core/dispatcher.go b/v1/core/dispatcher.go index 4ce6772..c569ebb 100644 --- a/v1/core/dispatcher.go +++ b/v1/core/dispatcher.go @@ -39,6 +39,7 @@ type UpdateHandler interface { type Poller interface { Start(ctx context.Context) (<-chan model.Update, error) + Stop(ctx context.Context) error } type DispatcherOption interface { @@ -109,6 +110,10 @@ func (d *Dispatcher) Start(ctx context.Context) error { return nil } +func (d *Dispatcher) Stop(ctx context.Context) error { + return d.poller.Stop(ctx) +} + func (d *Dispatcher) handleUpdates(ch <-chan model.Update) { for upd := range ch { d.wg.Go(func() { diff --git a/v1/core/pollers.go b/v1/core/pollers.go index fed31fd..1fdd5c1 100644 --- a/v1/core/pollers.go +++ b/v1/core/pollers.go @@ -34,6 +34,7 @@ type Webhook struct { type LongPoller struct { api *maxbot.Api pollInterval time.Duration + wg sync.WaitGroup } type TestPoller struct { @@ -148,45 +149,32 @@ func NewLongPoller(api *maxbot.Api, options ...LongPollerOption) *LongPoller { } 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() { + lp.wg.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 ch <- update: - } - } - } - } - }() + 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) @@ -264,6 +252,38 @@ func WithConnContext(fn func(ctx context.Context, c net.Conn) context.Context) W }) } +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) { + 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 ch <- update: + } + } + } + } +} + func newLongPollerOptions() *longPollerOptions { return &longPollerOptions{ pollInterval: defaultInterval,