Merge pull request #1450 from gotd/fix/reset-backoff-on-connection-ready

fix(telegram): reset backoff on connection ready
gotd · Oct 8, 2024 · a2c34af · a2c34af
2 parents 024d4de + 951db34
commit a2c34af
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 13 deletions.
diff --git a/examples/go.mod b/examples/go.mod
@@ -33,7 +33,7 @@ require (
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/gotd/ige v0.2.2 // indirect
 	github.com/gotd/neo v0.1.5 // indirect
-	github.com/klauspost/compress v1.17.9 // indirect
+	github.com/klauspost/compress v1.17.10 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect

diff --git a/examples/go.sum b/examples/go.sum
@@ -51,6 +51,7 @@ github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+o
 github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
 github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
+github.com/klauspost/compress v1.17.10/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=

diff --git a/telegram/connect.go b/telegram/connect.go
@@ -61,23 +61,37 @@ func (c *Client) reconnectUntilClosed(ctx context.Context) error {
 	// Note that we currently have no timeout on connection, so this is
 	// potentially eternal.
 	b := tdsync.SyncBackoff(backoff.WithContext(c.connBackoff(), ctx))
-
-	return backoff.RetryNotify(func() error {
-		if err := c.runUntilRestart(ctx); err != nil {
-			if c.isPermanentError(err) {
-				return backoff.Permanent(err)
+	g := tdsync.NewCancellableGroup(ctx)
+	g.Go(func(ctx context.Context) error {
+		for {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-c.ready.Ready():
+				// Reset backoff on successful connection.
+				b.Reset()
 			}
-			return err
 		}
+	})
+	g.Go(func(ctx context.Context) error {
+		return backoff.RetryNotify(func() error {
+			if err := c.runUntilRestart(ctx); err != nil {
+				if c.isPermanentError(err) {
+					return backoff.Permanent(err)
+				}
+				return err
+			}
 
-		return nil
-	}, b, func(err error, timeout time.Duration) {
-		c.log.Info("Restarting connection", zap.Error(err), zap.Duration("backoff", timeout))
+			return nil
+		}, b, func(err error, timeout time.Duration) {
+			c.log.Info("Restarting connection", zap.Error(err), zap.Duration("backoff", timeout))
 
-		c.connMux.Lock()
-		c.conn = c.createPrimaryConn(nil)
-		c.connMux.Unlock()
+			c.connMux.Lock()
+			c.conn = c.createPrimaryConn(nil)
+			c.connMux.Unlock()
+		})
 	})
+	return g.Wait()
 }
 
 func (c *Client) onReady() {