packages/docs-core/lib/Realtime/WebsocketConnection.ts

   1 import type { LoggerInterface } from '@proton/utils/logs'
   2 import { ConnectionCloseReason, SERVER_HEARTBEAT_INTERVAL } from '@proton/docs-proto'
   3 import { ApiResult, type WebsocketConnectionInterface, type WebsocketCallbacks } from '@proton/docs-shared'
   4 import type { WebsocketStateInterface } from './WebsocketState'
   5 import { WebsocketState } from './WebsocketState'
   6 import metrics from '@proton/metrics'
   7 import { isLocalEnvironment } from '../Util/isDevOrBlack'
   8 import { getWebSocketServerURL } from './getWebSocketServerURL'
   9 import type { MetricService } from '../Services/Metrics/MetricService'
  10
  11 /**
  12  * The heartbeat mechanism is temporarily disabled due to the fact that we cannot renew our heartbeat when receiving
  13  * a "ping" from the client as these messages are not exposed to our app—they are handled transparently by the browser.
  14  * The solution is to transition to a manual ping/pong mechanism.
  15  * See DRVDOC-535
  16  */
  17 const HeartbeatEnabled: false = false
  18
  19 /**
  20  * We will automatically close the connection if the document's visibility state goes to hidden and this amount of time elapses.
  21  */
  22 export const TIME_TO_WAIT_BEFORE_CLOSING_CONNECTION_AFTER_DOCUMENT_HIDES = 60_000 * 60
  23
  24 export const DebugConnection = {
  25   enabled: isLocalEnvironment() && false,
  26   url: 'ws://localhost:4000/websockets',
  27 }
  28
  29 export class WebsocketConnection implements WebsocketConnectionInterface {
  30   socket?: WebSocket
  31   readonly state: WebsocketStateInterface = new WebsocketState()
  32   private pingTimeout: ReturnType<typeof setTimeout> | undefined = undefined
  33   reconnectTimeout: ReturnType<typeof setTimeout> | undefined = undefined
  34   private destroyed = false
  35
  36   private didReceiveReadyMessageFromRTS = false
  37   closeConnectionDueToGoingAwayTimer: ReturnType<typeof setTimeout> | undefined = undefined
  38
  39   constructor(
  40     readonly callbacks: WebsocketCallbacks,
  41     readonly metricService: MetricService,
  42     private logger: LoggerInterface,
  43     private appVersion: string,
  44   ) {
  45     window.addEventListener('offline', this.handleWindowWentOfflineEvent)
  46     window.addEventListener('online', this.handleWindowCameOnlineEvent)
  47
  48     document.addEventListener('visibilitychange', this.handleVisibilityChangeEvent)
  49   }
  50
  51   handleVisibilityChangeEvent = (): void => {
  52     this.logger.info('Document visibility changed:', document.visibilityState)
  53
  54     if (document.visibilityState === 'visible') {
  55       if (this.closeConnectionDueToGoingAwayTimer) {
  56         clearTimeout(this.closeConnectionDueToGoingAwayTimer)
  57         this.closeConnectionDueToGoingAwayTimer = undefined
  58       }
  59
  60       if (!this.socket) {
  61         this.logger.info('Document became visible, reconnecting')
  62         this.queueReconnection({ skipDelay: true })
  63       }
  64     } else if (document.visibilityState === 'hidden') {
  65       this.closeConnectionDueToGoingAwayTimer = setTimeout(() => {
  66         this.logger.info('Closing connection due to user being away for too long')
  67         this.disconnect(ConnectionCloseReason.CODES.NORMAL_CLOSURE)
  68       }, TIME_TO_WAIT_BEFORE_CLOSING_CONNECTION_AFTER_DOCUMENT_HIDES)
  69
  70       if (this.reconnectTimeout) {
  71         clearTimeout(this.reconnectTimeout)
  72         this.reconnectTimeout = undefined
  73       }
  74     }
  75   }
  76
  77   handleWindowWentOfflineEvent = (): void => {
  78     this.disconnect(ConnectionCloseReason.CODES.NORMAL_CLOSURE)
  79   }
  80
  81   handleWindowCameOnlineEvent = (): void => {
  82     this.queueReconnection({ skipDelay: true })
  83   }
  84
  85   /**
  86    * In some cases, a client may lose their connection to the websocket without even realizing it.
  87    * The heartbeat explicitely closes the connection if we do not receive any message from the server,
  88    * including a "ping" message.
  89    * https://github.com/websockets/ws?tab=readme-ov-file#how-to-detect-and-close-broken-connections
  90    * */
  91   private heartbeat(): void {
  92     if (!HeartbeatEnabled) {
  93       return
  94     }
  95
  96     clearTimeout(this.pingTimeout)
  97
  98     this.pingTimeout = setTimeout(() => {
  99       this.logger.info('Closing connection due to heartbeat timeout')
 100       this.socket?.close(ConnectionCloseReason.CODES.NORMAL_CLOSURE)
 101     }, SERVER_HEARTBEAT_INTERVAL + 2_500)
 102   }
 103
 104   destroy(): void {
 105     this.destroyed = true
 106     this.state.destroy()
 107
 108     clearTimeout(this.pingTimeout)
 109     clearTimeout(this.reconnectTimeout)
 110
 111     window.removeEventListener('offline', this.handleWindowWentOfflineEvent)
 112     window.removeEventListener('online', this.handleWindowCameOnlineEvent)
 113     document.removeEventListener('visibilitychange', this.handleVisibilityChangeEvent)
 114
 115     this.disconnect(ConnectionCloseReason.CODES.NORMAL_CLOSURE)
 116   }
 117
 118   disconnect(code: number): void {
 119     this.socket?.close(code)
 120
 121     if (this.reconnectTimeout) {
 122       clearTimeout(this.reconnectTimeout)
 123       this.reconnectTimeout = undefined
 124     }
 125   }
 126
 127   buildConnectionUrl(params: { serverUrl: string; token: string }): string {
 128     const url = `${DebugConnection.enabled ? DebugConnection.url : params.serverUrl}/?token=${params.token}`
 129
 130     return url
 131   }
 132
 133   async getTokenOrFailConnection(): Promise<ApiResult<{ token: string }>> {
 134     const urlAndTokenResult = await this.callbacks.getUrlAndToken()
 135
 136     if (urlAndTokenResult.isFailed()) {
 137       this.logger.error('Failed to get realtime URL and token:', urlAndTokenResult.getError())
 138       this.state.didFailToFetchToken()
 139
 140       this.callbacks.onFailToGetToken(urlAndTokenResult.getError().code)
 141
 142       this.queueReconnection()
 143
 144       return ApiResult.fail(urlAndTokenResult.getError())
 145     }
 146
 147     return ApiResult.ok(urlAndTokenResult.getValue())
 148   }
 149
 150   async connect(abortSignal?: () => boolean): Promise<void> {
 151     if (this.destroyed) {
 152       throw new Error('Attempted to connect to a destroyed WebsocketConnection')
 153     }
 154
 155     if (document.visibilityState !== 'visible') {
 156       this.logger.warn('Attempting to connect socket while document is not visible')
 157       return
 158     }
 159
 160     if (this.state.isConnected || this.socket) {
 161       this.logger.warn('Attempted to connect while already connected')
 162       return
 163     }
 164
 165     clearTimeout(this.reconnectTimeout)
 166
 167     this.logger.info('Fetching url and token for websocket connection')
 168
 169     const urlAndTokenResult = await this.getTokenOrFailConnection()
 170     if (urlAndTokenResult.isFailed()) {
 171       return
 172     }
 173
 174     const { token } = urlAndTokenResult.getValue()
 175     const url = getWebSocketServerURL()
 176     const connectionUrl = this.buildConnectionUrl({
 177       serverUrl: url,
 178       token,
 179     })
 180
 181     if (abortSignal && abortSignal()) {
 182       this.logger.info('Aborting connection attempt due to abort signal')
 183       return
 184     }
 185
 186     this.logger.info('Opening websocket connection')
 187
 188     this.socket = new WebSocket(connectionUrl, [this.appVersion])
 189     this.socket.binaryType = 'arraybuffer'
 190
 191     this.callbacks.onConnecting()
 192
 193     this.socket.onopen = () => {
 194       this.logger.info(
 195         `Websocket connection opened; readyState: ${this.socket?.readyState} bufferAmount: ${this.socket?.bufferedAmount}`,
 196       )
 197
 198       this.heartbeat()
 199
 200       this.state.didOpen()
 201
 202       this.callbacks.onOpen()
 203     }
 204
 205     this.socket.onmessage = async (event) => {
 206       this.heartbeat()
 207       this.callbacks.onMessage(new Uint8Array(event.data))
 208     }
 209
 210     this.socket.onerror = (event) => {
 211       /** socket errors are completely opaque and convey no info. So we do not log an error here as to not pollute Sentry */
 212       this.logger.info('Websocket error:', event)
 213
 214       this.handleSocketClose(ConnectionCloseReason.CODES.INTERNAL_ERROR, 'Websocket error')
 215     }
 216
 217     this.socket.onclose = (event) => {
 218       this.logger.info('Websocket closed:', event.code, event.reason)
 219
 220       this.handleSocketClose(event.code, event.reason)
 221     }
 222   }
 223
 224   handleSocketClose(code: number, message: string): void {
 225     this.socket = undefined
 226     this.state.didClose()
 227
 228     const reason = ConnectionCloseReason.create({
 229       code,
 230       message,
 231     })
 232
 233     this.metricService.reportRealtimeDisconnect(reason)
 234
 235     if (this.state.isConnected) {
 236       this.callbacks.onClose(reason)
 237     } else {
 238       this.callbacks.onFailToConnect(reason)
 239     }
 240
 241     this.logDisconnectMetric(reason)
 242
 243     if (reason.props.code !== ConnectionCloseReason.CODES.UNAUTHORIZED && !this.destroyed) {
 244       this.queueReconnection()
 245     }
 246   }
 247
 248   queueReconnection(options: { skipDelay: boolean } = { skipDelay: false }): void {
 249     if (document.visibilityState !== 'visible') {
 250       this.logger.info('Not queueing reconnection because document is not visible')
 251       return
 252     }
 253
 254     const reconnectDelay = options.skipDelay ? 0 : this.state.getBackoff()
 255     this.logger.info(`Reconnecting in ${reconnectDelay}ms`)
 256     clearTimeout(this.reconnectTimeout)
 257
 258     this.reconnectTimeout = setTimeout(() => {
 259       if (document.visibilityState === 'visible') {
 260         void this.connect()
 261       }
 262     }, reconnectDelay)
 263   }
 264
 265   markAsReadyToAcceptMessages() {
 266     this.didReceiveReadyMessageFromRTS = true
 267   }
 268
 269   private logDisconnectMetric(reason: ConnectionCloseReason): void {
 270     if (
 271       [
 272         ConnectionCloseReason.CODES.TLS_HANDSHAKE,
 273         ConnectionCloseReason.CODES.TIMEOUT,
 274         ConnectionCloseReason.CODES.PROTOCOL_ERROR,
 275       ].includes(reason.props.code)
 276     ) {
 277       metrics.docs_failed_websocket_connections_total.increment({
 278         retry: 'false',
 279         type: 'network_error',
 280       })
 281     } else if (
 282       [
 283         ConnectionCloseReason.CODES.INTERNAL_ERROR,
 284         ConnectionCloseReason.CODES.UNAUTHORIZED,
 285         ConnectionCloseReason.CODES.BAD_GATEWAY,
 286       ].includes(reason.props.code)
 287     ) {
 288       metrics.docs_failed_websocket_connections_total.increment({
 289         retry: 'false',
 290         type: 'server_error',
 291       })
 292     } else {
 293       metrics.docs_failed_websocket_connections_total.increment({
 294         retry: 'false',
 295         type: 'unknown',
 296       })
 297     }
 298   }
 299
 300   public canBroadcastMessages(): boolean {
 301     if (
 302       !this.didReceiveReadyMessageFromRTS ||
 303       !this.socket ||
 304       this.socket.readyState !== WebSocket.OPEN ||
 305       !this.state.isConnected
 306     ) {
 307       return false
 308     }
 309
 310     return true
 311   }
 312
 313   public isConnected(): boolean {
 314     return this.state.isConnected && this.socket?.readyState === WebSocket.OPEN
 315   }
 316
 317   async broadcastMessage(data: Uint8Array): Promise<void> {
 318     if (!this.didReceiveReadyMessageFromRTS) {
 319       this.logger.error('Cannot send message, RTS is not ready to accept messages')
 320       return
 321     }
 322
 323     if (!this.canBroadcastMessages()) {
 324       this.logger.error('Cannot send message, socket is not open')
 325       return
 326     }
 327
 328     this.socket?.send(data)
 329   }
 330 }