fix: changed dispatcher lifecycle interfaces

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() {

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,