From b5a056e8232c6b52ea2dbb2957cd268847d88557 Mon Sep 17 00:00:00 2001 From: XYSK-lilong007 <267018309+XYSK-lilong007@users.noreply.github.com> Date: Fri, 13 Mar 2026 08:18:41 +0800 Subject: [PATCH] fix(telegram): reconnect long polling after disconnect --- pkg/channels/telegram/telegram.go | 148 +++++++++++++++++++++---- pkg/channels/telegram/telegram_test.go | 80 +++++++++++++ 2 files changed, 209 insertions(+), 19 deletions(-) diff --git a/pkg/channels/telegram/telegram.go b/pkg/channels/telegram/telegram.go index 4a8d34a9f..fc28e91c2 100644 --- a/pkg/channels/telegram/telegram.go +++ b/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) { @@ -106,24 +116,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{ - Timeout: 30, - }) + updates, err := c.startLongPolling(c.ctx) 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 - - bh.HandleMessage(func(ctx *th.Context, message telego.Message) error { - return c.handleMessage(ctx, &message) - }, th.AnyMessage()) + c.setBotHandler(bh) c.SetRunning(true) logger.InfoCF("telegram", "Telegram bot connected", map[string]any{ @@ -132,13 +136,7 @@ func (c *TelegramChannel) Start(ctx context.Context) error { c.startCommandRegistration(c.ctx, commands.BuiltinDefinitions()) - go func() { - if err = bh.Start(); err != nil { - logger.ErrorCF("telegram", "Bot handler failed", map[string]any{ - "error": err.Error(), - }) - } - }() + go c.runPollingLoop(c.ctx, bh) return nil } @@ -148,8 +146,8 @@ func (c *TelegramChannel) Stop(ctx context.Context) error { c.SetRunning(false) // Stop the bot handler - if c.bh != nil { - _ = c.bh.StopWithContext(ctx) + if bh := c.currentBotHandler(); bh != nil { + _ = bh.StopWithContext(ctx) } // Cancel our context (stops long polling) @@ -163,6 +161,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 diff --git a/pkg/channels/telegram/telegram_test.go b/pkg/channels/telegram/telegram_test.go index c2186d0a3..93d8b1348 100644 --- a/pkg/channels/telegram/telegram_test.go +++ b/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") + } +}