mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2025-11-29 02:42:51 +00:00
Fix: Prevent race condition and goroutine leak in TgBot
Addresses a critical race condition on the global `botCancel` variable, which could occur if `Tgbot.OnReceive()` was called concurrently (e.g., during rapid panel restarts or unexpected behavior). Changes in tgbot.go: - Added `tgBotMutex sync.Mutex` to ensure thread safety. - Protected `botCancel` creation and assignment in `OnReceive()` using the mutex, and added a check to prevent overwriting an active context, which avoids goroutine leaks. - Protected the cancellation and cleanup logic in `StopBot()` with the mutex.
This commit is contained in:
parent
9ac8fdf3d3
commit
bb78b9495a
1 changed files with 22 additions and 6 deletions
|
|
@ -42,6 +42,8 @@ var (
|
||||||
|
|
||||||
// botCancel stores the function to cancel the context, stopping Long Polling gracefully.
|
// botCancel stores the function to cancel the context, stopping Long Polling gracefully.
|
||||||
botCancel context.CancelFunc
|
botCancel context.CancelFunc
|
||||||
|
// tgBotMutex protects concurrent access to botCancel variable
|
||||||
|
tgBotMutex sync.Mutex
|
||||||
|
|
||||||
botHandler *th.BotHandler
|
botHandler *th.BotHandler
|
||||||
adminIds []int64
|
adminIds []int64
|
||||||
|
|
@ -328,6 +330,9 @@ func (t *Tgbot) Stop() {
|
||||||
// StopBot safely stops the Telegram bot's Long Polling operation by cancelling its context.
|
// StopBot safely stops the Telegram bot's Long Polling operation by cancelling its context.
|
||||||
// This is the global function called from main.go's signal handler and t.Stop().
|
// This is the global function called from main.go's signal handler and t.Stop().
|
||||||
func StopBot() {
|
func StopBot() {
|
||||||
|
tgBotMutex.Lock()
|
||||||
|
defer tgBotMutex.Unlock()
|
||||||
|
|
||||||
if botCancel != nil {
|
if botCancel != nil {
|
||||||
logger.Info("Sending cancellation signal to Telegram bot...")
|
logger.Info("Sending cancellation signal to Telegram bot...")
|
||||||
|
|
||||||
|
|
@ -376,16 +381,27 @@ func (t *Tgbot) OnReceive() {
|
||||||
// updates, _ := bot.UpdatesViaLongPolling(context.Background(), ¶ms)
|
// updates, _ := bot.UpdatesViaLongPolling(context.Background(), ¶ms)
|
||||||
|
|
||||||
// --- GRACEFUL SHUTDOWN FIX: Context creation ---
|
// --- GRACEFUL SHUTDOWN FIX: Context creation ---
|
||||||
// Create a context with cancellation and store the cancel function.
|
tgBotMutex.Lock()
|
||||||
var ctx context.Context
|
|
||||||
ctx, botCancel = context.WithCancel(context.Background())
|
|
||||||
// --------------------------------------------------
|
|
||||||
|
|
||||||
// Get updates channel using the created context. This channel will close when ctx is cancelled.
|
// Create a context with cancellation and store the cancel function.
|
||||||
updates, _ := bot.UpdatesViaLongPolling(ctx, ¶ms) // <<<
|
var ctx context.Context
|
||||||
|
|
||||||
|
// Check if botCancel is already set (to prevent race condition overwrite and goroutine leak)
|
||||||
|
if botCancel == nil {
|
||||||
|
ctx, botCancel = context.WithCancel(context.Background())
|
||||||
|
} else {
|
||||||
|
// If botCancel is already set, use a non-cancellable context for this redundant call.
|
||||||
|
// This prevents overwriting the active botCancel and causing a goroutine leak from the previous call.
|
||||||
|
logger.Warning("TgBot OnReceive called concurrently. Using background context for redundant call.")
|
||||||
|
ctx = context.Background() // <<< ИЗМЕНЕНИЕ
|
||||||
|
}
|
||||||
|
|
||||||
|
tgBotMutex.Unlock()
|
||||||
|
|
||||||
|
// Get updates channel using the context.
|
||||||
|
updates, _ := bot.UpdatesViaLongPolling(ctx, ¶ms)
|
||||||
|
|
||||||
botHandler, _ = th.NewBotHandler(bot, updates)
|
botHandler, _ = th.NewBotHandler(bot, updates)
|
||||||
|
|
||||||
botHandler.HandleMessage(func(ctx *th.Context, message telego.Message) error {
|
botHandler.HandleMessage(func(ctx *th.Context, message telego.Message) error {
|
||||||
delete(userStates, message.Chat.ID)
|
delete(userStates, message.Chat.ID)
|
||||||
t.SendMsgToTgbot(message.Chat.ID, t.I18nBot("tgbot.keyboardClosed"), tu.ReplyKeyboardRemove())
|
t.SendMsgToTgbot(message.Chat.ID, t.I18nBot("tgbot.keyboardClosed"), tu.ReplyKeyboardRemove())
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue