From 224b01e7a4db154a21aebe8fac5130faeaf1738c Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Tue, 15 Apr 2025 12:09:11 +0300
Subject: [PATCH] client: Wait for updates manager to finish on disconnect

---
 pkg/connector/client.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/pkg/connector/client.go b/pkg/connector/client.go
index 37ba9788..7c0f1a8b 100644
--- a/pkg/connector/client.go
+++ b/pkg/connector/client.go
@@ -64,6 +64,7 @@ type TelegramClient struct {
 	userLogin      *bridgev2.UserLogin
 	client         *telegram.Client
 	updatesManager *updates.Manager
+	updatesCloseC  chan struct{}
 	clientCtx      context.Context
 	clientCancel   context.CancelFunc
 	clientCloseC   <-chan struct{}
@@ -504,12 +505,22 @@ func (t *TelegramClient) Connect(ctx context.Context) {
 		t.sendBadCredentialsOrUnknownError(err)
 		return
 	}
+	t.updatesCloseC = make(chan struct{})
 	go func() {
-		err = t.updatesManager.Run(t.clientCtx, t.client.API(), t.telegramUserID, updates.AuthOptions{})
-		if err != nil && !errors.Is(err, context.Canceled) {
-			zerolog.Ctx(t.clientCtx).Err(err).Msg("failed to run updates manager")
-			t.Disconnect()
-			t.Connect(t.main.Bridge.Log.WithContext(context.Background()))
+		defer close(t.updatesCloseC)
+		for {
+			err = t.updatesManager.Run(t.clientCtx, t.client.API(), t.telegramUserID, updates.AuthOptions{})
+			if err == nil || errors.Is(err, context.Canceled) {
+				return
+			}
+
+			zerolog.Ctx(t.clientCtx).Err(err).Msg("failed to run updates manager, retrying")
+
+			select {
+			case <-t.clientCtx.Done():
+				return
+			case <-time.After(5 * time.Second):
+			}
 		}
 	}()
 
@@ -542,6 +553,9 @@ func (t *TelegramClient) Disconnect() {
 	if t.clientCloseC != nil {
 		<-t.clientCloseC
 	}
+	if t.updatesCloseC != nil {
+		<-t.updatesCloseC
+	}
 }
 
 func (t *TelegramClient) getInputUser(ctx context.Context, id int64) (*tg.InputUser, error) {