archify-vscode-ext/src/sync/connection.ts

import WebSocket from 'ws';
import * as Y from 'yjs';
import * as awarenessProtocol from 'y-protocols/awareness';
import * as syncProtocol from 'y-protocols/sync';
import * as decoding from 'lib0/decoding';
import * as encoding from 'lib0/encoding';

const MESSAGE_SYNC = 0;
const MESSAGE_AWARENESS = 1;
const MESSAGE_QUERY_AWARENESS = 3;

export type ConnectionState =
  | 'connecting'
  | 'online'
  | 'offline'
  /** No session for the server, or the server refused it. */
  | 'signed-out'
  /** The server no longer lets this account in (removed from the space). */
  | 'forbidden'
  /** The server speaks another protocol version. */
  | 'outdated';

export interface ConnectionOptions {
  url: string;
  /** A fresh access token; `force` after the server refused the last one. */
  token: (force: boolean) => Promise<string | undefined>;
  doc: Y.Doc;
  awareness: awarenessProtocol.Awareness;
  onState: (state: ConnectionState) => void;
  /** After the first full exchange of each connection. */
  onSynced: () => void;
  log: (message: string) => void;
}

/**
 * One space's WebSocket to the server: the y-websocket protocol (sync step 1
 * and 2, updates, awareness), reconnecting with backoff. Edits made while
 * offline go out on reconnect through the state-vector exchange.
 */
export class LiveConnection {
  state: ConnectionState = 'connecting';
  synced = false;
  private socket?: WebSocket;
  private attempt = 0;
  private timer?: NodeJS.Timeout;
  private closed = false;
  private forceToken = false;

  constructor(private readonly options: ConnectionOptions) {
    options.doc.on('update', this.onDocUpdate);
    options.awareness.on('update', this.onAwarenessUpdate);
    void this.connect();
  }

  /** Connects at once, e.g. after signing in. */
  reconnectNow(): void {
    if (this.closed || this.state === 'online' || this.state === 'connecting') return;
    clearTimeout(this.timer);
    this.attempt = 0;
    void this.connect();
  }

  private setState(state: ConnectionState): void {
    if (this.state === state) return;
    this.state = state;
    this.options.onState(state);
  }

  private async connect(): Promise<void> {
    if (this.closed) return;
    this.setState('connecting');
    let token: string | undefined;
    try {
      token = await this.options.token(this.forceToken);
    } catch (error) {
      this.options.log(`sync: no token (${String(error)})`);
    }
    this.forceToken = false;
    if (this.closed) return;
    if (!token) {
      this.setState('signed-out');
      return;
    }
    const socket = new WebSocket(this.options.url, { headers: { authorization: `Bearer ${token}` }, handshakeTimeout: 10_000 });
    socket.binaryType = 'arraybuffer';
    this.socket = socket;
    let refused = false;

    socket.on('open', () => {
      const step1 = encoding.createEncoder();
      encoding.writeVarUint(step1, MESSAGE_SYNC);
      syncProtocol.writeSyncStep1(step1, this.options.doc);
      this.send(encoding.toUint8Array(step1));
      if (this.options.awareness.getLocalState() !== null) this.sendAwareness([this.options.doc.clientID]);
    });
    socket.on('message', (data: ArrayBuffer) => this.receive(new Uint8Array(data)));
    socket.on('unexpected-response', (_request, response) => {
      refused = true;
      const status = response.statusCode ?? 0;
      this.options.log(`sync: the server refused the connection (${status})`);
      if (status === 401) {
        // Try once more with a refreshed token; without a session, stop.
        this.forceToken = true;
        this.setState('signed-out');
        this.schedule(true);
      } else if (status === 403 || status === 404) {
        this.setState('forbidden');
      } else if (status === 426) {
        this.setState('outdated');
      } else {
        this.setState('offline');
        this.schedule();
      }
      socket.terminate();
    });
    socket.on('error', (error) => this.options.log(`sync: ${error.message}`));
    socket.on('close', () => {
      if (this.socket !== socket) return;
      this.socket = undefined;
      this.synced = false;
      const others = [...this.options.awareness.getStates().keys()].filter((id) => id !== this.options.doc.clientID);
      awarenessProtocol.removeAwarenessStates(this.options.awareness, others, this);
      if (this.closed || refused) return;
      this.setState('offline');
      this.schedule();
    });
  }

  /** Backoff: 0.5 s doubling to 30 s, with jitter. */
  private schedule(once = false): void {
    if (this.closed) return;
    clearTimeout(this.timer);
    if (once && this.attempt > 0) return;
    const delay = Math.min(30_000, 500 * 2 ** this.attempt) * (0.75 + Math.random() / 2);
    this.attempt++;
    this.timer = setTimeout(() => void this.connect(), delay);
  }

  private receive(bytes: Uint8Array): void {
    const decoder = decoding.createDecoder(bytes);
    const type = decoding.readVarUint(decoder);
    if (type === MESSAGE_SYNC) {
      const reply = encoding.createEncoder();
      encoding.writeVarUint(reply, MESSAGE_SYNC);
      const sub = syncProtocol.readSyncMessage(decoder, reply, this.options.doc, this);
      if (encoding.length(reply) > 1) this.send(encoding.toUint8Array(reply));
      if (sub === syncProtocol.messageYjsSyncStep2 && !this.synced) {
        this.synced = true;
        this.attempt = 0;
        this.setState('online');
        this.options.onSynced();
      }
    } else if (type === MESSAGE_AWARENESS) {
      awarenessProtocol.applyAwarenessUpdate(this.options.awareness, decoding.readVarUint8Array(decoder), this);
    } else if (type === MESSAGE_QUERY_AWARENESS) {
      this.sendAwareness([...this.options.awareness.getStates().keys()]);
    }
  }

  private send(bytes: Uint8Array): void {
    if (this.socket?.readyState === WebSocket.OPEN) this.socket.send(bytes);
  }

  private sendAwareness(clients: number[]): void {
    const message = encoding.createEncoder();
    encoding.writeVarUint(message, MESSAGE_AWARENESS);
    encoding.writeVarUint8Array(message, awarenessProtocol.encodeAwarenessUpdate(this.options.awareness, clients));
    this.send(encoding.toUint8Array(message));
  }

  private readonly onDocUpdate = (update: Uint8Array, origin: unknown) => {
    if (origin === this) return;
    const message = encoding.createEncoder();
    encoding.writeVarUint(message, MESSAGE_SYNC);
    syncProtocol.writeUpdate(message, update);
    this.send(encoding.toUint8Array(message));
  };

  private readonly onAwarenessUpdate = (
    { added, updated, removed }: { added: number[]; updated: number[]; removed: number[] },
    origin: unknown,
  ) => {
    if (origin === this) return;
    this.sendAwareness([...added, ...updated, ...removed]);
  };

  dispose(): void {
    this.closed = true;
    clearTimeout(this.timer);
    // Tell the others this editor left before closing.
    awarenessProtocol.removeAwarenessStates(this.options.awareness, [this.options.doc.clientID], 'dispose');
    this.options.doc.off('update', this.onDocUpdate);
    this.options.awareness.off('update', this.onAwarenessUpdate);
    this.socket?.close();
    this.socket = undefined;
  }
}