updater.StopBot should close polling context

ext/botmapping.go

@@ -1,6 +1,7 @@
 package ext
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"io"
@@ -19,6 +20,8 @@ type botData struct {
 
 	// updateChan represents the incoming updates channel.
 	updateChan chan json.RawMessage
+	// pollingContextCloser allows one to stop polling instantly without waiting for the return
+	pollingContextCloser context.CancelFunc
 	// updateWriterControl is used to count the number of current writers on the update channel.
 	// This is required to ensure that we can safely close the channel, and thus stop processing incoming updates.
 	// While this remains non-zero, it is unsafe to close the update channel.
@@ -47,12 +50,22 @@ type botMapping struct {
 	errorLog *log.Logger
 }
 
-var ErrBotAlreadyExists = errors.New("bot already exists in bot mapping")
-var ErrBotUrlPathAlreadyExists = errors.New("url path already exists in bot mapping")
+var (
+	ErrBotAlreadyExists        = errors.New("bot already exists in bot mapping")
+	ErrBotUrlPathAlreadyExists = errors.New("url path already exists in bot mapping")
+)
+
+func (m *botMapping) addWebhookBot(b *gotgbot.Bot, urlPath string, webhookSecret string) (*botData, error) {
+	return m.addBot(b, urlPath, webhookSecret, nil)
+}
+
+func (m *botMapping) addPollingBot(b *gotgbot.Bot, ctxClose context.CancelFunc) (*botData, error) {
+	return m.addBot(b, "", "", ctxClose)
+}
 
 // addBot Adds a new bot to the botMapping structure.
 // Pass an empty urlPath/webhookSecret if using polling instead of webhooks.
-func (m *botMapping) addBot(b *gotgbot.Bot, urlPath string, webhookSecret string) (*botData, error) {
+func (m *botMapping) addBot(b *gotgbot.Bot, urlPath string, webhookSecret string, ctxClose context.CancelFunc) (*botData, error) {
 	// Clean up the URLPath such that it remains consistent.
 	urlPath = strings.TrimPrefix(urlPath, "/")
 
@@ -75,12 +88,13 @@ func (m *botMapping) addBot(b *gotgbot.Bot, urlPath string, webhookSecret string
 	}
 
 	bData := botData{
-		bot:                 b,
-		updateChan:          make(chan json.RawMessage),
-		stopUpdates:         make(chan struct{}),
-		updateWriterControl: &sync.WaitGroup{},
-		urlPath:             urlPath,
-		webhookSecret:       webhookSecret,
+		bot:                  b,
+		updateChan:           make(chan json.RawMessage),
+		pollingContextCloser: ctxClose,
+		stopUpdates:          make(chan struct{}),
+		updateWriterControl:  &sync.WaitGroup{},
+		urlPath:              urlPath,
+		webhookSecret:        webhookSecret,
 	}
 
 	m.mapping[bData.bot.Token] = bData
@@ -208,6 +222,11 @@ func (b *botData) stop() {
 		close(b.stopUpdates)
 	}
 
+	// If we have a context to close, close it. This will stop polling immediately.
+	if b.pollingContextCloser != nil {
+		b.pollingContextCloser()
+	}
+
 	// Wait for all writers to finish writing to the updateChannel
 	b.updateWriterControl.Wait()
 

ext/botmapping_test.go

@@ -18,7 +18,7 @@ func Test_botMapping(t *testing.T) {
 	t.Run("addBot", func(t *testing.T) {
 		// check that bots can be added fine
 		var err error
-		origBdata, err = bm.addBot(b, "", "")
+		origBdata, err = bm.addPollingBot(b, nil)
 		if err != nil {
 			t.Errorf("expected to be able to add a new bot fine: %s", err.Error())
 			t.FailNow()
@@ -31,7 +31,7 @@ func Test_botMapping(t *testing.T) {
 
 	t.Run("doubleAdd", func(t *testing.T) {
 		// Adding the same bot twice should fail
-		_, err := bm.addBot(b, "", "")
+		_, err := bm.addPollingBot(b, nil)
 		if err == nil {
 			t.Errorf("adding the same bot twice should throw an error")
 			t.FailNow()
@@ -84,7 +84,10 @@ func Test_botData_isUpdateChannelStopped(t *testing.T) {
 		BotClient: &gotgbot.BaseBotClient{},
 	}
 
-	bData, err := bm.addBot(b, "", "")
+	ctxCancelled := false
+	bData, err := bm.addPollingBot(b, func() {
+		ctxCancelled = true
+	})
 	if err != nil {
 		t.Errorf("bot with token %s should not have failed to be added", b.Token)
 		return
@@ -99,4 +102,7 @@ func Test_botData_isUpdateChannelStopped(t *testing.T) {
 		t.Errorf("bot with token %s should be stopped", b.Token)
 		return
 	}
+	if !ctxCancelled {
+		t.Errorf("bot with token %s should have a cancelled context ", b.Token)
+	}
 }

ext/updater.go

@@ -157,18 +157,20 @@ func (u *Updater) StartPolling(b *gotgbot.Bot, opts *PollingOpts) error {
 		}
 	}
 
-	bData, err := u.botMapping.addBot(b, "", "")
+	ctx, closeFn := context.WithCancel(context.Background())
+
+	bData, err := u.botMapping.addPollingBot(b, closeFn)
 	if err != nil {
 		return fmt.Errorf("failed to add bot with long polling: %w", err)
 	}
 
 	go u.Dispatcher.Start(b, bData.updateChan)
-	go u.pollingLoop(bData, reqOpts, v)
+	go u.pollingLoop(ctx, bData, reqOpts, v)
 
 	return nil
 }
 
-func (u *Updater) pollingLoop(bData *botData, opts *gotgbot.RequestOpts, v map[string]string) {
+func (u *Updater) pollingLoop(ctx context.Context, bData *botData, opts *gotgbot.RequestOpts, v map[string]string) {
 	bData.updateWriterControl.Add(1)
 	defer bData.updateWriterControl.Done()
 
@@ -180,8 +182,13 @@ func (u *Updater) pollingLoop(bData *botData, opts *gotgbot.RequestOpts, v map[s
 
 		// Manually craft the getUpdate calls to improve memory management, reduce json parsing overheads, and
 		// unnecessary reallocation of url.Values in the polling loop.
-		r, err := bData.bot.Request("getUpdates", v, nil, opts)
+		r, err := bData.bot.RequestWithContext(ctx, "getUpdates", v, nil, opts)
 		if err != nil {
+			if errors.Is(err, context.Canceled) {
+				// context cancelled; means the bot was stopped gracefully through Updater.StopBot
+				return
+			}
+
 			if u.UnhandledErrFunc != nil {
 				u.UnhandledErrFunc(err)
 			} else {
@@ -317,7 +324,7 @@ func (u *Updater) AddWebhook(b *gotgbot.Bot, urlPath string, opts *AddWebhookOpt
 		secretToken = opts.SecretToken
 	}
 
-	bData, err := u.botMapping.addBot(b, urlPath, secretToken)
+	bData, err := u.botMapping.addWebhookBot(b, urlPath, secretToken)
 	if err != nil {
 		return fmt.Errorf("failed to add webhook for bot: %w", err)
 	}

ext/updater_test.go

@@ -152,6 +152,54 @@ func concurrentTest(t *testing.T) {
 	time.Sleep(delay * 2)
 }
 
+func TestUpdater_StopBot(t *testing.T) {
+	server := basicTestServer(t, map[string]*testEndpoint{
+		"getUpdates": {
+			delay: time.Second * 3, // server close will take 3s
+			reply: `{"ok": true, "result": []}`,
+		},
+	})
+	defer server.Close()
+
+	reqOpts := &gotgbot.RequestOpts{
+		APIURL:  server.URL,
+		Timeout: time.Second * 4,
+	}
+
+	b := &gotgbot.Bot{
+		User:      gotgbot.User{},
+		Token:     "SOME_TOKEN",
+		BotClient: &gotgbot.BaseBotClient{},
+	}
+
+	d := ext.NewDispatcher(&ext.DispatcherOpts{MaxRoutines: 1})
+	u := ext.NewUpdater(d, nil)
+
+	err := u.StartPolling(b, &ext.PollingOpts{
+		GetUpdatesOpts: &gotgbot.GetUpdatesOpts{
+			RequestOpts: reqOpts,
+		},
+	})
+	if err != nil {
+		t.Errorf("failed to start polling: %v", err)
+		return
+	}
+
+	// sleep to ensure polling is fully underway
+	time.Sleep(time.Millisecond * 500)
+
+	// Should still have lots of time until timeout, so... kill the bots.
+	start := time.Now()
+	u.StopAllBots()
+
+	// If it took longer than 1ms then we know something went wrong; ctx should've stopped immediately.
+	since := time.Since(start)
+	t.Logf("stopping took %dms", since.Milliseconds())
+	if since > time.Millisecond {
+		t.Errorf("stopping all bots took %dms; shouldve taken less than 1ms", since.Milliseconds())
+	}
+}
+
 func TestUpdaterDisallowsEmptyWebhooks(t *testing.T) {
 	b := &gotgbot.Bot{
 		Token:     "SOME_TOKEN",