From 26d068eeb2e568da96cda064ffa67ac32398a059 Mon Sep 17 00:00:00 2001
From: alina sireneva <alina@tei.su>
Date: Tue, 24 Sep 2024 22:26:17 +0300
Subject: [PATCH] fix: improved ping handling

---
 packages/core/src/network/mtproto-session.ts  |  2 +-
 packages/core/src/network/network-manager.ts  |  8 ++++++++
 .../core/src/network/session-connection.ts    | 20 +++++++++++++------
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/packages/core/src/network/mtproto-session.ts b/packages/core/src/network/mtproto-session.ts
index d9440085..a48eb146 100644
--- a/packages/core/src/network/mtproto-session.ts
+++ b/packages/core/src/network/mtproto-session.ts
@@ -167,7 +167,6 @@ export class MtprotoSession {
 
         timers.clearTimeout(this.current429Timeout)
         this.resetState(withAuthKey)
-        this.resetLastPing(true)
     }
 
     resetAuthKey(): void {
@@ -231,6 +230,7 @@ export class MtprotoSession {
         this.getStateSchedule.clear()
         this.chains.clear()
         this.chainsPendingFails.clear()
+        this.resetLastPing(true)
     }
 
     enqueueRpc(rpc: PendingRpc, force?: boolean): boolean {
diff --git a/packages/core/src/network/network-manager.ts b/packages/core/src/network/network-manager.ts
index 363da5a1..ce839fb9 100644
--- a/packages/core/src/network/network-manager.ts
+++ b/packages/core/src/network/network-manager.ts
@@ -108,6 +108,13 @@ export interface NetworkManagerExtraParams {
      * > If you need to handle incoming updates, use a {@link Dispatcher} instead.
      */
     middlewares?: Middleware<RpcCallMiddlewareContext, unknown>[]
+
+    /**
+     * Ping interval in milliseconds.
+     *
+     * @default  60000 (1 minute)
+     */
+    pingInterval?: number
 }
 
 /** Options that can be customized when making an RPC call */
@@ -255,6 +262,7 @@ export class DcConnectionManager {
             enableErrorReporting: this.manager.params.enableErrorReporting,
             salts: this._salts,
             platform: this.manager.params.platform,
+            pingInterval: this.manager.params.pingInterval ?? 60_000,
         })
 
         const mainParams = baseConnectionParams()
diff --git a/packages/core/src/network/session-connection.ts b/packages/core/src/network/session-connection.ts
index c634625f..00f015f1 100644
--- a/packages/core/src/network/session-connection.ts
+++ b/packages/core/src/network/session-connection.ts
@@ -27,6 +27,7 @@ import { TransportError } from './transports/abstract.js'
 export interface SessionConnectionParams extends PersistentConnectionParams {
     initConnection: tl.RawInitConnectionRequest
     inactivityTimeout?: number
+    pingInterval: number
     enableErrorReporting: boolean
     layer: number
     disableUpdates?: boolean
@@ -43,7 +44,6 @@ export interface SessionConnectionParams extends PersistentConnectionParams {
 }
 
 const TEMP_AUTH_KEY_EXPIRY = 86400 // 24 hours
-const PING_INTERVAL = 60000 // 1 minute
 const GET_STATE_INTERVAL = 1500 // 1.5 seconds
 
 // destroy_auth_key#d1435160 = DestroyAuthKeyRes;
@@ -84,6 +84,10 @@ export class SessionConnection extends PersistentConnection {
     private _crypto: ICryptoProvider
     private _salts: ServerSaltManager
 
+    // todo: we should probably do adaptive ping interval based on rtt like tdlib:
+    // https://github.com/tdlib/td/blob/91aa6c9e4d0774eabf4f8d7f3aa51239032059a6/td/mtproto/SessionConnection.h
+    private _pingInterval: number
+
     constructor(
         params: SessionConnectionParams,
         readonly _session: MtprotoSession,
@@ -91,6 +95,8 @@ export class SessionConnection extends PersistentConnection {
         super(params, _session.log.create('conn'))
         this._flushTimer.onTimeout(this._flush.bind(this))
 
+        this._pingInterval = params.pingInterval
+
         this._readerMap = params.readerMap
         this._writerMap = params.writerMap
         this._crypto = params.crypto
@@ -1564,7 +1570,7 @@ export class SessionConnection extends PersistentConnection {
             // between multiple connections using the same session
             this._flushTimer.emitWhenIdle()
         } else {
-            const nextPingTime = this._session.lastPingTime + PING_INTERVAL
+            const nextPingTime = this._session.lastPingTime + this._pingInterval
             const nextGetScheduleTime = this._session.getStateSchedule.raw[0]?.getState || Infinity
 
             this._flushTimer.emitBefore(Math.min(nextPingTime, nextGetScheduleTime))
@@ -1629,17 +1635,19 @@ export class SessionConnection extends PersistentConnection {
             messageCount += 1
         }
 
-        if (now - this._session.lastPingTime > PING_INTERVAL) {
+        if (now - this._session.lastPingTime > this._pingInterval) {
             if (!this._session.lastPingMsgId.isZero()) {
-                this.log.warn("didn't receive pong for previous ping (msg_id = %l)", this._session.lastPingMsgId)
-                this._session.pendingMessages.delete(this._session.lastPingMsgId)
+                this.log.warn("didn't receive pong for previous ping (msg_id = %l). are we offline?", this._session.lastPingMsgId)
+                this._resetSession()
+                return
             }
 
             pingId = randomLong()
             const obj: mtp.RawMt_ping_delay_disconnect = {
                 _: 'mt_ping_delay_disconnect',
                 pingId,
-                disconnectDelay: 75,
+                // ÷ 1000 * 1.25
+                disconnectDelay: Math.round(this._pingInterval * 0.00125),
             }
 
             this._session.lastPingTime = Date.now()