Merge PR #1455

pkg/channels/telegram/telegram.go

@@ -9,6 +9,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mymmrac/telego"
@@ -41,7 +42,8 @@ var (
 type TelegramChannel struct {
 	*channels.BaseChannel
 	bot     *telego.Bot
-	bh      *th.BotHandler
+	bh      telegramBotHandler
+	bhMu    sync.Mutex
 	config  *config.Config
 	chatIDs map[string]int64
 	ctx     context.Context
@@ -49,6 +51,14 @@ type TelegramChannel struct {
 
 	registerFunc     func(context.Context, []commands.Definition) error
 	commandRegCancel context.CancelFunc
+	startPollingFunc func(context.Context) (<-chan telego.Update, error)
+	newHandlerFunc   func(<-chan telego.Update) (telegramBotHandler, error)
+	sleepFunc        func(context.Context, time.Duration) bool
+}
+
+type telegramBotHandler interface {
+	Start() error
+	StopWithContext(context.Context) error
 }
 
 func NewTelegramChannel(cfg *config.Config, bus *bus.MessageBus) (*TelegramChannel, error) {
@@ -107,24 +117,18 @@ func (c *TelegramChannel) Start(ctx context.Context) error {
 
 	c.ctx, c.cancel = context.WithCancel(ctx)
 
-	updates, err := c.bot.UpdatesViaLongPolling(c.ctx, &telego.GetUpdatesParams{
+	updates, err := c.startLongPolling(c.ctx)
-		Timeout: 30,
-	})
 	if err != nil {
 		c.cancel()
 		return fmt.Errorf("failed to start long polling: %w", err)
 	}
 
-	bh, err := th.NewBotHandler(c.bot, updates)
+	bh, err := c.newBotHandler(updates)
 	if err != nil {
 		c.cancel()
 		return fmt.Errorf("failed to create bot handler: %w", err)
 	}
-	c.bh = bh
+	c.setBotHandler(bh)
-
-	bh.HandleMessage(func(ctx *th.Context, message telego.Message) error {
-		return c.handleMessage(ctx, &message)
-	}, th.AnyMessage())
 
 	c.SetRunning(true)
 	logger.InfoCF("telegram", "Telegram bot connected", map[string]any{
@@ -133,13 +137,7 @@ func (c *TelegramChannel) Start(ctx context.Context) error {
 
 	c.startCommandRegistration(c.ctx, commands.BuiltinDefinitions())
 
-	go func() {
+	go c.runPollingLoop(c.ctx, bh)
-		if err = bh.Start(); err != nil {
-			logger.ErrorCF("telegram", "Bot handler failed", map[string]any{
-				"error": err.Error(),
-			})
-		}
-	}()
 
 	return nil
 }
@@ -149,8 +147,8 @@ func (c *TelegramChannel) Stop(ctx context.Context) error {
 	c.SetRunning(false)
 
 	// Stop the bot handler
-	if c.bh != nil {
+	if bh := c.currentBotHandler(); bh != nil {
-		_ = c.bh.StopWithContext(ctx)
+		_ = bh.StopWithContext(ctx)
 	}
 
 	// Cancel our context (stops long polling)
@@ -164,6 +162,118 @@ func (c *TelegramChannel) Stop(ctx context.Context) error {
 	return nil
 }
 
+func (c *TelegramChannel) startLongPolling(ctx context.Context) (<-chan telego.Update, error) {
+	if c.startPollingFunc != nil {
+		return c.startPollingFunc(ctx)
+	}
+	return c.bot.UpdatesViaLongPolling(ctx, &telego.GetUpdatesParams{
+		Timeout: 30,
+	})
+}
+
+func (c *TelegramChannel) newBotHandler(updates <-chan telego.Update) (telegramBotHandler, error) {
+	if c.newHandlerFunc != nil {
+		return c.newHandlerFunc(updates)
+	}
+
+	bh, err := th.NewBotHandler(c.bot, updates)
+	if err != nil {
+		return nil, err
+	}
+	bh.HandleMessage(func(ctx *th.Context, message telego.Message) error {
+		return c.handleMessage(ctx, &message)
+	}, th.AnyMessage())
+	return bh, nil
+}
+
+func (c *TelegramChannel) setBotHandler(bh telegramBotHandler) {
+	c.bhMu.Lock()
+	defer c.bhMu.Unlock()
+	c.bh = bh
+}
+
+func (c *TelegramChannel) currentBotHandler() telegramBotHandler {
+	c.bhMu.Lock()
+	defer c.bhMu.Unlock()
+	return c.bh
+}
+
+func (c *TelegramChannel) sleep(ctx context.Context, delay time.Duration) bool {
+	if c.sleepFunc != nil {
+		return c.sleepFunc(ctx, delay)
+	}
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return false
+	case <-timer.C:
+		return true
+	}
+}
+
+func (c *TelegramChannel) runPollingLoop(ctx context.Context, bh telegramBotHandler) {
+	backoff := time.Second
+	for {
+		if err := bh.Start(); err != nil && ctx.Err() == nil {
+			logger.ErrorCF("telegram", "Bot handler failed", map[string]any{
+				"error": err.Error(),
+			})
+		}
+		if ctx.Err() != nil {
+			return
+		}
+
+		logger.WarnC("telegram", "Updates channel closed, restarting long polling")
+		if !c.sleep(ctx, backoff) {
+			return
+		}
+
+		for {
+			updates, err := c.startLongPolling(ctx)
+			if err != nil {
+				logger.ErrorCF("telegram", "Failed to restart long polling", map[string]any{
+					"error":       err.Error(),
+					"retry_after": backoff.String(),
+				})
+				if backoff < 30*time.Second {
+					backoff *= 2
+					if backoff > 30*time.Second {
+						backoff = 30 * time.Second
+					}
+				}
+				if !c.sleep(ctx, backoff) {
+					return
+				}
+				continue
+			}
+
+			nextHandler, err := c.newBotHandler(updates)
+			if err != nil {
+				logger.ErrorCF("telegram", "Failed to recreate bot handler", map[string]any{
+					"error":       err.Error(),
+					"retry_after": backoff.String(),
+				})
+				if backoff < 30*time.Second {
+					backoff *= 2
+					if backoff > 30*time.Second {
+						backoff = 30 * time.Second
+					}
+				}
+				if !c.sleep(ctx, backoff) {
+					return
+				}
+				continue
+			}
+
+			c.setBotHandler(nextHandler)
+			bh = nextHandler
+			backoff = time.Second
+			break
+		}
+	}
+}
+
 func (c *TelegramChannel) Send(ctx context.Context, msg bus.OutboundMessage) error {
 	if !c.IsRunning() {
 		return channels.ErrNotRunning

pkg/channels/telegram/telegram_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"io"
 	"strings"
 	"testing"
 	"time"
@@ -460,3 +461,82 @@ func TestHandleMessage_ReplyThread_NonForum_NoIsolation(t *testing.T) {
 	assert.Empty(t, inbound.Metadata["parent_peer_kind"])
 	assert.Empty(t, inbound.Metadata["parent_peer_id"])
 }
+
+type fakeTelegramBotHandler struct {
+	start func() error
+	stop  func(context.Context) error
+}
+
+func (f *fakeTelegramBotHandler) Start() error {
+	if f.start != nil {
+		return f.start()
+	}
+	return nil
+}
+
+func (f *fakeTelegramBotHandler) StopWithContext(ctx context.Context) error {
+	if f.stop != nil {
+		return f.stop(ctx)
+	}
+	return nil
+}
+
+func TestRunPollingLoop_RestartsAfterHandlerStops(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	started := make(chan string, 2)
+	first := &fakeTelegramBotHandler{
+		start: func() error {
+			started <- "first"
+			return nil
+		},
+	}
+	second := &fakeTelegramBotHandler{
+		start: func() error {
+			started <- "second"
+			<-ctx.Done()
+			return io.EOF
+		},
+	}
+
+	ch := &TelegramChannel{
+		BaseChannel: channels.NewBaseChannel("telegram", nil, nil, nil),
+		sleepFunc: func(ctx context.Context, _ time.Duration) bool {
+			return ctx.Err() == nil
+		},
+	}
+
+	startPollingCalls := 0
+	ch.startPollingFunc = func(context.Context) (<-chan telego.Update, error) {
+		startPollingCalls++
+		return make(chan telego.Update), nil
+	}
+	ch.newHandlerFunc = func(<-chan telego.Update) (telegramBotHandler, error) {
+		return second, nil
+	}
+
+	done := make(chan struct{})
+	go func() {
+		ch.runPollingLoop(ctx, first)
+		close(done)
+	}()
+
+	if got := <-started; got != "first" {
+		t.Fatalf("first start = %q, want %q", got, "first")
+	}
+	if got := <-started; got != "second" {
+		t.Fatalf("second start = %q, want %q", got, "second")
+	}
+	if startPollingCalls != 1 {
+		t.Fatalf("startPollingCalls = %d, want %d", startPollingCalls, 1)
+	}
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("runPollingLoop did not stop after context cancellation")
+	}
+}